Coverage for /builds/kinetik161/ase/ase/md/andersen.py: 96.77%

1"""Andersen dynamics class.""" 

2from typing import IO, Optional, Union 

3 

4from numpy import cos, log, ones, pi, random, repeat 

5 

6from ase import Atoms, units 

7from ase.md.md import MolecularDynamics 

8from ase.parallel import DummyMPI, world 

9 

10 

11class Andersen(MolecularDynamics): 

12 """Andersen (constant N, V, T) molecular dynamics.""" 

13 

14 def __init__( 

15 self, 

16 atoms: Atoms, 

17 timestep: float, 

18 temperature_K: float, 

19 andersen_prob: float, 

20 fixcm: bool = True, 

21 trajectory: Optional[str] = None, 

22 logfile: Optional[Union[IO, str]] = None, 

23 loginterval: int = 1, 

24 communicator=world, 

25 rng=random, 

26 append_trajectory: bool = False, 

27 ): 

28 """" 

29 Parameters: 

30 

31 atoms: Atoms object 

32 The list of atoms. 

33 

34 timestep: float 

35 The time step in ASE time units. 

36 

37 temperature_K: float 

38 The desired temperature, in Kelvin. 

39 

40 andersen_prob: float 

41 A random collision probability, typically 1e-4 to 1e-1. 

42 With this probability atoms get assigned random velocity components. 

43 

44 fixcm: bool (optional) 

45 If True, the position and momentum of the center of mass is 

46 kept unperturbed. Default: True. 

47 

48 rng: RNG object (optional) 

49 Random number generator, by default numpy.random. Must have a 

50 random_sample method matching the signature of 

51 numpy.random.random_sample. 

52 

53 logfile: file object or str (optional) 

54 If *logfile* is a string, a file with that name will be opened. 

55 Use '-' for stdout. 

56 

57 trajectory: Trajectory object or str (optional) 

58 Attach trajectory object. If *trajectory* is a string a 

59 Trajectory will be constructed. Use *None* (the default) for no 

60 trajectory. 

61 

62 communicator: MPI communicator (optional) 

63 Communicator used to distribute random numbers to all tasks. 

64 Default: ase.parallel.world. Set to None to disable communication. 

65 

66 append_trajectory: bool (optional) 

67 Defaults to False, which causes the trajectory file to be 

68 overwritten each time the dynamics is restarted from scratch. 

69 If True, the new structures are appended to the trajectory 

70 file instead. 

71 

72 The temperature is imposed by stochastic collisions with a heat bath 

73 that acts on velocity components of randomly chosen particles. 

74 The algorithm randomly decorrelates velocities, so dynamical properties 

75 like diffusion or viscosity cannot be properly measured. 

76 

77 H. C. Andersen, J. Chem. Phys. 72 (4), 2384–2393 (1980) 

78 """ 

79 self.temp = units.kB * temperature_K 

80 self.andersen_prob = andersen_prob 

81 self.fix_com = fixcm 

82 self.rng = rng 

83 if communicator is None: 

84 communicator = DummyMPI() 

85 self.communicator = communicator 

86 MolecularDynamics.__init__(self, atoms, timestep, trajectory, 

87 logfile, loginterval, 

88 append_trajectory=append_trajectory) 

89 

90 def set_temperature(self, temperature_K): 

91 self.temp = units.kB * temperature_K 

92 

93 def set_andersen_prob(self, andersen_prob): 

94 self.andersen_prob = andersen_prob 

95 

96 def set_timestep(self, timestep): 

97 self.dt = timestep 

98 

99 def boltzmann_random(self, width, size): 

100 x = self.rng.random_sample(size=size) 

101 y = self.rng.random_sample(size=size) 

102 z = width * cos(2 * pi * x) * (-2 * log(1 - y))**0.5 

103 return z 

104 

105 def get_maxwell_boltzmann_velocities(self): 

106 natoms = len(self.atoms) 

107 masses = repeat(self.masses, 3).reshape(natoms, 3) 

108 width = (self.temp / masses)**0.5 

109 velos = self.boltzmann_random(width, size=(natoms, 3)) 

110 return velos # [[x, y, z],] components for each atom 

111 

112 def step(self, forces=None): 

113 atoms = self.atoms 

114 

115 if forces is None: 

116 forces = atoms.get_forces(md=True) 

117 

118 self.v = atoms.get_velocities() 

119 

120 # Random atom-wise variables are stored as attributes and broadcasted: 

121 # - self.random_com_velocity # added to all atoms if self.fix_com 

122 # - self.random_velocity # added to some atoms if the per-atom 

123 # - self.andersen_chance # andersen_chance <= andersen_prob 

124 # a dummy communicator will be used for serial runs 

125 

126 if self.fix_com: 

127 # add random velocity to center of mass to prepare Andersen 

128 width = (self.temp / sum(self.masses))**0.5 

129 self.random_com_velocity = (ones(self.v.shape) 

130 * self.boltzmann_random(width, (3))) 

131 self.communicator.broadcast(self.random_com_velocity, 0) 

132 self.v += self.random_com_velocity 

133 

134 self.v += 0.5 * forces / self.masses * self.dt 

135 

136 # apply Andersen thermostat 

137 self.random_velocity = self.get_maxwell_boltzmann_velocities() 

138 self.andersen_chance = self.rng.random_sample(size=self.v.shape) 

139 self.communicator.broadcast(self.random_velocity, 0) 

140 self.communicator.broadcast(self.andersen_chance, 0) 

141 self.v[self.andersen_chance <= self.andersen_prob] \ 

142 = self.random_velocity[self.andersen_chance <= self.andersen_prob] 

143 

144 x = atoms.get_positions() 

145 if self.fix_com: 

146 old_com = atoms.get_center_of_mass() 

147 self.v -= self._get_com_velocity(self.v) 

148 # Step: x^n -> x^(n+1) - this applies constraints if any 

149 atoms.set_positions(x + self.v * self.dt) 

150 if self.fix_com: 

151 atoms.set_center_of_mass(old_com) 

152 

153 # recalc velocities after RATTLE constraints are applied 

154 self.v = (atoms.get_positions() - x) / self.dt 

155 forces = atoms.get_forces(md=True) 

156 

157 # Update the velocities 

158 self.v += 0.5 * forces / self.masses * self.dt 

159 

160 if self.fix_com: 

161 self.v -= self._get_com_velocity(self.v) 

162 

163 # Second part of RATTLE taken care of here 

164 atoms.set_momenta(self.v * self.masses) 

165 

166 return forces