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