1"""Berendsen NPT dynamics class."""
2
3import numpy as np
4import warnings
5
6from ase.md.nvtberendsen import NVTBerendsen
7import ase.units as units
8
9
10class NPTBerendsen(NVTBerendsen):
11    def __init__(self, atoms, timestep, temperature=None,
12                 *, temperature_K=None, pressure=None, pressure_au=None,
13                 taut=0.5e3 * units.fs, taup=1e3 * units.fs,
14                 compressibility=None, compressibility_au=None, fixcm=True,
15                 trajectory=None,
16                 logfile=None, loginterval=1, append_trajectory=False):
17        """Berendsen (constant N, P, T) molecular dynamics.
18
19        This dynamics scale the velocities and volumes to maintain a constant
20        pressure and temperature.  The shape of the simulation cell is not
21        altered, if that is desired use Inhomogenous_NPTBerendsen.
22
23        Parameters:
24
25        atoms: Atoms object
26            The list of atoms.
27
28        timestep: float
29            The time step in ASE time units.
30
31        temperature: float
32            The desired temperature, in Kelvin.
33
34        temperature_K: float
35            Alias for ``temperature``.
36
37        pressure: float (deprecated)
38            The desired pressure, in bar (1 bar = 1e5 Pa).  Deprecated,
39            use ``pressure_au`` instead.
40
41        pressure: float
42            The desired pressure, in atomic units (eV/Å^3).
43
44        taut: float
45            Time constant for Berendsen temperature coupling in ASE
46            time units.  Default: 0.5 ps.
47
48        taup: float
49            Time constant for Berendsen pressure coupling.  Default: 1 ps.
50
51        compressibility: float (deprecated)
52            The compressibility of the material, in bar-1.  Deprecated,
53            use ``compressibility_au`` instead.
54
55        compressibility_au: float
56            The compressibility of the material, in atomic units (Å^3/eV).
57
58        fixcm: bool (optional)
59            If True, the position and momentum of the center of mass is
60            kept unperturbed.  Default: True.
61
62        trajectory: Trajectory object or str (optional)
63            Attach trajectory object.  If *trajectory* is a string a
64            Trajectory will be constructed.  Use *None* for no
65            trajectory.
66
67        logfile: file object or str (optional)
68            If *logfile* is a string, a file with that name will be opened.
69            Use '-' for stdout.
70
71        loginterval: int (optional)
72            Only write a log line for every *loginterval* time steps.
73            Default: 1
74
75        append_trajectory: boolean (optional)
76            Defaults to False, which causes the trajectory file to be
77            overwriten each time the dynamics is restarted from scratch.
78            If True, the new structures are appended to the trajectory
79            file instead.
80
81
82        """
83
84        NVTBerendsen.__init__(self, atoms, timestep, temperature=temperature,
85                              temperature_K=temperature_K,
86                              taut=taut, fixcm=fixcm, trajectory=trajectory,
87                              logfile=logfile, loginterval=loginterval,
88                              append_trajectory=append_trajectory)
89        self.taup = taup
90        self.pressure = self._process_pressure(pressure, pressure_au)
91        if compressibility is not None and compressibility_au is not None:
92            raise TypeError(
93                "Do not give both 'compressibility' and 'compressibility_au'")
94        if compressibility is not None:
95            # Specified in bar, convert to atomic units
96            warnings.warn(FutureWarning(
97                "Specify the compressibility in atomic units."))
98            self.set_compressibility(
99                compressibility_au=compressibility / (1e5 * units.Pascal))
100        else:
101            self.set_compressibility(compressibility_au=compressibility_au)
102
103    def set_taup(self, taup):
104        self.taup = taup
105
106    def get_taup(self):
107        return self.taup
108
109    def set_pressure(self, pressure=None, *, pressure_au=None,
110                     pressure_bar=None):
111        self.pressure = self._process_pressure(pressure, pressure_bar,
112                                               pressure_au)
113
114    def get_pressure(self):
115        return self.pressure
116
117    def set_compressibility(self, *, compressibility_au):
118        self.compressibility = compressibility_au
119
120    def get_compressibility(self):
121        return self.compressibility
122
123    def set_timestep(self, timestep):
124        self.dt = timestep
125
126    def get_timestep(self):
127        return self.dt
128
129    def scale_positions_and_cell(self):
130        """ Do the Berendsen pressure coupling,
131        scale the atom position and the simulation cell."""
132
133        taupscl = self.dt / self.taup
134        stress = self.atoms.get_stress(voigt=False, include_ideal_gas=True)
135        old_pressure = -stress.trace() / 3
136        scl_pressure = (1.0 - taupscl * self.compressibility / 3.0 *
137                        (self.pressure - old_pressure))
138
139        #print("old_pressure", old_pressure, self.pressure)
140        #print("volume scaling by:", scl_pressure)
141
142        cell = self.atoms.get_cell()
143        cell = scl_pressure * cell
144        self.atoms.set_cell(cell, scale_atoms=True)
145
146    def step(self, forces=None):
147        """ move one timestep forward using Berenden NPT molecular dynamics."""
148
149        NVTBerendsen.scale_velocities(self)
150        self.scale_positions_and_cell()
151
152        # one step velocity verlet
153        atoms = self.atoms
154
155        if forces is None:
156            forces = atoms.get_forces(md=True)
157
158        p = self.atoms.get_momenta()
159        p += 0.5 * self.dt * forces
160
161        if self.fix_com:
162            # calculate the center of mass
163            # momentum and subtract it
164            psum = p.sum(axis=0) / float(len(p))
165            p = p - psum
166
167        self.atoms.set_positions(
168            self.atoms.get_positions() +
169            self.dt * p / self.atoms.get_masses()[:, np.newaxis])
170
171        # We need to store the momenta on the atoms before calculating
172        # the forces, as in a parallel Asap calculation atoms may
173        # migrate during force calculations, and the momenta need to
174        # migrate along with the atoms.  For the same reason, we
175        # cannot use self.masses in the line above.
176
177        self.atoms.set_momenta(p)
178        forces = self.atoms.get_forces(md=True)
179        atoms.set_momenta(self.atoms.get_momenta() + 0.5 * self.dt * forces)
180
181        return forces
182
183    def _process_pressure(self, pressure, pressure_au):
184        """Handle that pressure can be specified in multiple units.
185
186        For at least a transition period, Berendsen NPT dynamics in ASE can
187        have the pressure specified in either bar or atomic units (eV/Å^3).
188
189        Two parameters:
190
191        pressure: None or float
192            The original pressure specification in bar.
193            A warning is issued if this is not None.
194
195        pressure_au: None or float
196            Pressure in ev/Å^3.
197
198        Exactly one of the two pressure parameters must be different from
199        None, otherwise an error is issued.
200
201        Return value: Pressure in eV/Å^3.
202        """
203        if (pressure is not None) + (pressure_au is not None) != 1:
204            raise TypeError("Exactly one of the parameters 'pressure',"
205                            + " and 'pressure_au' must"
206                            + " be given")
207
208        if pressure is not None:
209            w = ("The 'pressure' parameter is deprecated, please"
210                 + " specify the pressure in atomic units (eV/Å^3)"
211                 + " using the 'pressure_au' parameter.")
212            warnings.warn(FutureWarning(w))
213            return pressure * (1e5 * units.Pascal)
214        else:
215            return pressure_au
216
217
218class Inhomogeneous_NPTBerendsen(NPTBerendsen):
219    """Berendsen (constant N, P, T) molecular dynamics.
220
221    This dynamics scale the velocities and volumes to maintain a constant
222    pressure and temperature.  The size of the unit cell is allowed to change
223    independently in the three directions, but the angles remain constant.
224
225    Usage: NPTBerendsen(atoms, timestep, temperature, taut, pressure, taup)
226
227    atoms
228        The list of atoms.
229
230    timestep
231        The time step.
232
233    temperature
234        The desired temperature, in Kelvin.
235
236    taut
237        Time constant for Berendsen temperature coupling.
238
239    fixcm
240        If True, the position and momentum of the center of mass is
241        kept unperturbed.  Default: True.
242
243    pressure
244        The desired pressure, in bar (1 bar = 1e5 Pa).
245
246    taup
247        Time constant for Berendsen pressure coupling.
248
249    compressibility
250        The compressibility of the material, water 4.57E-5 bar-1, in bar-1
251
252    mask
253        Specifies which axes participate in the barostat.  Default (1, 1, 1)
254        means that all axes participate, set any of them to zero to disable
255        the barostat in that direction.
256    """
257
258    def __init__(self, atoms, timestep, temperature=None,
259                 *, temperature_K=None,
260                 taut=0.5e3 * units.fs, pressure=None,
261                 pressure_au=None, taup=1e3 * units.fs,
262                 compressibility=None, compressibility_au=None,
263                 mask=(1, 1, 1), fixcm=True, trajectory=None,
264                 logfile=None, loginterval=1):
265
266        NPTBerendsen.__init__(self, atoms, timestep, temperature=temperature,
267                              temperature_K=temperature_K,
268                              taut=taut, taup=taup, pressure=pressure,
269                              pressure_au=pressure_au,
270                              compressibility=compressibility,
271                              compressibility_au=compressibility_au,
272                              fixcm=fixcm, trajectory=trajectory,
273                              logfile=logfile, loginterval=loginterval)
274        self.mask = mask
275
276    def scale_positions_and_cell(self):
277        """ Do the Berendsen pressure coupling,
278        scale the atom position and the simulation cell."""
279
280        taupscl = self.dt * self.compressibility / self.taup / 3.0
281        stress = - self.atoms.get_stress(include_ideal_gas=True)
282        if stress.shape == (6,):
283            stress = stress[:3]
284        elif stress.shape == (3, 3):
285            stress = [stress[i][i] for i in range(3)]
286        else:
287            raise ValueError('Cannot use a stress tensor of shape ' +
288                             str(stress.shape))
289        pbc = self.atoms.get_pbc()
290        scl_pressurex = 1.0 - taupscl * (self.pressure - stress[0]) \
291            * pbc[0] * self.mask[0]
292        scl_pressurey = 1.0 - taupscl * (self.pressure - stress[1]) \
293            * pbc[1] * self.mask[1]
294        scl_pressurez = 1.0 - taupscl * (self.pressure - stress[2]) \
295            * pbc[2] * self.mask[2]
296        cell = self.atoms.get_cell()
297        cell = np.array([scl_pressurex * cell[0],
298                         scl_pressurey * cell[1],
299                         scl_pressurez * cell[2]])
300        self.atoms.set_cell(cell, scale_atoms=True)
301