1"""Langevin dynamics class."""
2
3import numpy as np
4
5from ase.md.md import MolecularDynamics
6from ase.parallel import world, DummyMPI
7from ase import units
8
9
10class Langevin(MolecularDynamics):
11    """Langevin (constant N, V, T) molecular dynamics."""
12
13    # Helps Asap doing the right thing.  Increment when changing stuff:
14    _lgv_version = 4
15
16    def __init__(self, atoms, timestep, temperature=None, friction=None,
17                 fixcm=True, *, temperature_K=None, trajectory=None,
18                 logfile=None, loginterval=1, communicator=world,
19                 rng=None, append_trajectory=False):
20        """
21        Parameters:
22
23        atoms: Atoms object
24            The list of atoms.
25
26        timestep: float
27            The time step in ASE time units.
28
29        temperature: float (deprecated)
30            The desired temperature, in electron volt.
31
32        temperature_K: float
33            The desired temperature, in Kelvin.
34
35        friction: float
36            A friction coefficient, typically 1e-4 to 1e-2.
37
38        fixcm: bool (optional)
39            If True, the position and momentum of the center of mass is
40            kept unperturbed.  Default: True.
41
42        rng: RNG object (optional)
43            Random number generator, by default numpy.random.  Must have a
44            standard_normal method matching the signature of
45            numpy.random.standard_normal.
46
47        logfile: file object or str (optional)
48            If *logfile* is a string, a file with that name will be opened.
49            Use '-' for stdout.
50
51        trajectory: Trajectory object or str (optional)
52            Attach trajectory object.  If *trajectory* is a string a
53            Trajectory will be constructed.  Use *None* (the default) for no
54            trajectory.
55
56        communicator: MPI communicator (optional)
57            Communicator used to distribute random numbers to all tasks.
58            Default: ase.parallel.world. Set to None to disable communication.
59
60        append_trajectory: bool (optional)
61            Defaults to False, which causes the trajectory file to be
62            overwritten each time the dynamics is restarted from scratch.
63            If True, the new structures are appended to the trajectory
64            file instead.
65
66        The temperature and friction are normally scalars, but in principle one
67        quantity per atom could be specified by giving an array.
68
69        RATTLE constraints can be used with these propagators, see:
70        E. V.-Eijnden, and G. Ciccotti, Chem. Phys. Lett. 429, 310 (2006)
71
72        The propagator is Equation 23 (Eq. 39 if RATTLE constraints are used)
73        of the above reference.  That reference also contains another
74        propagator in Eq. 21/34; but that propagator is not quasi-symplectic
75        and gives a systematic offset in the temperature at large time steps.
76        """
77        if friction is None:
78            raise TypeError("Missing 'friction' argument.")
79        self.fr = friction
80        self.temp = units.kB * self._process_temperature(temperature,
81                                                         temperature_K, 'eV')
82        self.fix_com = fixcm
83        if communicator is None:
84            communicator = DummyMPI()
85        self.communicator = communicator
86        if rng is None:
87            self.rng = np.random
88        else:
89            self.rng = rng
90        MolecularDynamics.__init__(self, atoms, timestep, trajectory,
91                                   logfile, loginterval,
92                                   append_trajectory=append_trajectory)
93        self.updatevars()
94
95    def todict(self):
96        d = MolecularDynamics.todict(self)
97        d.update({'temperature_K': self.temp / units.kB,
98                  'friction': self.fr,
99                  'fixcm': self.fix_com})
100        return d
101
102    def set_temperature(self, temperature=None, temperature_K=None):
103        self.temp = units.kB * self._process_temperature(temperature,
104                                                         temperature_K, 'eV')
105        self.updatevars()
106
107    def set_friction(self, friction):
108        self.fr = friction
109        self.updatevars()
110
111    def set_timestep(self, timestep):
112        self.dt = timestep
113        self.updatevars()
114
115    def updatevars(self):
116        dt = self.dt
117        T = self.temp
118        fr = self.fr
119        masses = self.masses
120        sigma = np.sqrt(2 * T * fr / masses)
121
122        self.c1 = dt / 2. - dt * dt * fr / 8.
123        self.c2 = dt * fr / 2 - dt * dt * fr * fr / 8.
124        self.c3 = np.sqrt(dt) * sigma / 2. - dt**1.5 * fr * sigma / 8.
125        self.c5 = dt**1.5 * sigma / (2 * np.sqrt(3))
126        self.c4 = fr / 2. * self.c5
127
128    def step(self, forces=None):
129        atoms = self.atoms
130        natoms = len(atoms)
131
132        if forces is None:
133            forces = atoms.get_forces(md=True)
134
135        # This velocity as well as xi, eta and a few other variables are stored
136        # as attributes, so Asap can do its magic when atoms migrate between
137        # processors.
138        self.v = atoms.get_velocities()
139
140        self.xi = self.rng.standard_normal(size=(natoms, 3))
141        self.eta = self.rng.standard_normal(size=(natoms, 3))
142
143        # When holonomic constraints for rigid linear triatomic molecules are
144        # present, ask the constraints to redistribute xi and eta within each
145        # triple defined in the constraints. This is needed to achieve the
146        # correct target temperature.
147        for constraint in self.atoms.constraints:
148            if hasattr(constraint, 'redistribute_forces_md'):
149                constraint.redistribute_forces_md(atoms, self.xi, rand=True)
150                constraint.redistribute_forces_md(atoms, self.eta, rand=True)
151
152        self.communicator.broadcast(self.xi, 0)
153        self.communicator.broadcast(self.eta, 0)
154
155        # First halfstep in the velocity.
156        self.v += (self.c1 * forces / self.masses - self.c2 * self.v +
157                   self.c3 * self.xi - self.c4 * self.eta)
158
159        # Full step in positions
160        x = atoms.get_positions()
161        if self.fix_com:
162            old_com = atoms.get_center_of_mass()
163        # Step: x^n -> x^(n+1) - this applies constraints if any.
164        atoms.set_positions(x + self.dt * self.v + self.c5 * self.eta)
165        if self.fix_com:
166            atoms.set_center_of_mass(old_com)
167
168        # recalc velocities after RATTLE constraints are applied
169        self.v = (self.atoms.get_positions() - x -
170                  self.c5 * self.eta) / self.dt
171        forces = atoms.get_forces(md=True)
172
173        # Update the velocities
174        self.v += (self.c1 * forces / self.masses - self.c2 * self.v +
175                   self.c3 * self.xi - self.c4 * self.eta)
176
177        if self.fix_com:  # subtract center of mass vel
178            self.v -= self._get_com_velocity(self.v)
179
180        # Second part of RATTLE taken care of here
181        atoms.set_momenta(self.v * self.masses)
182
183        return forces
184