1"""Andersen dynamics class."""
2
3from numpy import random, cos, pi, log, ones, repeat
4
5from ase.md.md import MolecularDynamics
6from ase.parallel import world, DummyMPI
7from ase import units
8
9
10class Andersen(MolecularDynamics):
11    """Andersen (constant N, V, T) molecular dynamics."""
12
13    def __init__(self, atoms, timestep, temperature_K, andersen_prob,
14                 fixcm=True, trajectory=None, logfile=None, loginterval=1,
15                 communicator=world, rng=random, append_trajectory=False):
16        """"
17        Parameters:
18
19        atoms: Atoms object
20            The list of atoms.
21
22        timestep: float
23            The time step in ASE time units.
24
25        temperature_K: float
26            The desired temperature, in Kelvin.
27
28        andersen_prob: float
29            A random collision probability, typically 1e-4 to 1e-1.
30            With this probability atoms get assigned random velocity components.
31
32        fixcm: bool (optional)
33            If True, the position and momentum of the center of mass is
34            kept unperturbed.  Default: True.
35
36        rng: RNG object (optional)
37            Random number generator, by default numpy.random.  Must have a
38            random_sample method matching the signature of
39            numpy.random.random_sample.
40
41        logfile: file object or str (optional)
42            If *logfile* is a string, a file with that name will be opened.
43            Use '-' for stdout.
44
45        trajectory: Trajectory object or str (optional)
46            Attach trajectory object. If *trajectory* is a string a
47            Trajectory will be constructed. Use *None* (the default) for no
48            trajectory.
49
50        communicator: MPI communicator (optional)
51            Communicator used to distribute random numbers to all tasks.
52            Default: ase.parallel.world. Set to None to disable communication.
53
54        append_trajectory: bool (optional)
55            Defaults to False, which causes the trajectory file to be
56            overwritten each time the dynamics is restarted from scratch.
57            If True, the new structures are appended to the trajectory
58            file instead.
59
60        The temperature is imposed by stochastic collisions with a heat bath
61        that acts on velocity components of randomly chosen particles.
62        The algorithm randomly decorrelates velocities, so dynamical properties
63        like diffusion or viscosity cannot be properly measured.
64
65        H. C. Andersen, J. Chem. Phys. 72 (4), 2384–2393 (1980)
66        """
67        self.temp = units.kB * temperature_K
68        self.andersen_prob = andersen_prob
69        self.fix_com = fixcm
70        self.rng = rng
71        if communicator is None:
72            communicator = DummyMPI()
73        self.communicator = communicator
74        MolecularDynamics.__init__(self, atoms, timestep, trajectory,
75                                   logfile, loginterval,
76                                   append_trajectory=append_trajectory)
77
78    def set_temperature(self, temperature_K):
79        self.temp = units.kB * temperature_K
80
81    def set_andersen_prob(self, andersen_prob):
82        self.andersen_prob = andersen_prob
83
84    def set_timestep(self, timestep):
85        self.dt = timestep
86
87    def boltzmann_random(self, width, size):
88        x = self.rng.random_sample(size=size)
89        y = self.rng.random_sample(size=size)
90        z = width * cos(2 * pi * x) * (-2 * log(1 - y))**0.5
91        return z
92
93    def get_maxwell_boltzmann_velocities(self):
94        natoms = len(self.atoms)
95        masses = repeat(self.masses, 3).reshape(natoms, 3)
96        width = (self.temp / masses)**0.5
97        velos = self.boltzmann_random(width, size=(natoms, 3))
98        return velos  # [[x, y, z],] components for each atom
99
100    def step(self, forces=None):
101        atoms = self.atoms
102
103        if forces is None:
104            forces = atoms.get_forces(md=True)
105
106        self.v = atoms.get_velocities()
107
108        # Random atom-wise variables are stored as attributes and broadcasted:
109        #  - self.random_com_velocity  # added to all atoms if self.fix_com
110        #  - self.random_velocity      # added to some atoms if the per-atom
111        #  - self.andersen_chance      # andersen_chance <= andersen_prob
112        # a dummy communicator will be used for serial runs
113
114        if self.fix_com:
115            # add random velocity to center of mass to prepare Andersen
116            width = (self.temp / sum(self.masses))**0.5
117            self.random_com_velocity = (ones(self.v.shape)
118                                        * self.boltzmann_random(width, (3)))
119            self.communicator.broadcast(self.random_com_velocity, 0)
120            self.v += self.random_com_velocity
121
122        self.v += 0.5 * forces / self.masses * self.dt
123
124        # apply Andersen thermostat
125        self.random_velocity = self.get_maxwell_boltzmann_velocities()
126        self.andersen_chance = self.rng.random_sample(size=self.v.shape)
127        self.communicator.broadcast(self.random_velocity, 0)
128        self.communicator.broadcast(self.andersen_chance, 0)
129        self.v[self.andersen_chance <= self.andersen_prob] \
130            = self.random_velocity[self.andersen_chance <= self.andersen_prob]
131
132        x = atoms.get_positions()
133        if self.fix_com:
134            old_com = atoms.get_center_of_mass()
135            self.v -= self._get_com_velocity(self.v)
136        # Step: x^n -> x^(n+1) - this applies constraints if any
137        atoms.set_positions(x + self.v * self.dt)
138        if self.fix_com:
139            atoms.set_center_of_mass(old_com)
140
141        # recalc velocities after RATTLE constraints are applied
142        self.v = (atoms.get_positions() - x) / self.dt
143        forces = atoms.get_forces(md=True)
144
145        # Update the velocities
146        self.v += 0.5 * forces / self.masses * self.dt
147
148        if self.fix_com:
149            self.v -= self._get_com_velocity(self.v)
150
151        # Second part of RATTLE taken care of here
152        atoms.set_momenta(self.v * self.masses)
153
154        return forces
155