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