Coverage for /builds/debichem-team/python-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, default: ``numpy.random`` 

49 Random number generator. This must have the ``random`` method 

50 with the same signature as ``numpy.random.random``. 

51 

52 logfile: file object or str (optional) 

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

54 Use '-' for stdout. 

55 

56 trajectory: Trajectory object or str (optional) 

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

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

59 trajectory. 

60 

61 communicator: MPI communicator (optional) 

62 Communicator used to distribute random numbers to all tasks. 

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

64 

65 append_trajectory: bool (optional) 

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

67 overwritten each time the dynamics is restarted from scratch. 

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

69 file instead. 

70 

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

72 that acts on velocity components of randomly chosen particles. 

73 The algorithm randomly decorrelates velocities, so dynamical properties 

74 like diffusion or viscosity cannot be properly measured. 

75 

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

77 """ 

78 self.temp = units.kB * temperature_K 

79 self.andersen_prob = andersen_prob 

80 self.fix_com = fixcm 

81 self.rng = rng 

82 if communicator is None: 

83 communicator = DummyMPI() 

84 self.communicator = communicator 

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

86 logfile, loginterval, 

87 append_trajectory=append_trajectory) 

88 

89 def set_temperature(self, temperature_K): 

90 self.temp = units.kB * temperature_K 

91 

92 def set_andersen_prob(self, andersen_prob): 

93 self.andersen_prob = andersen_prob 

94 

95 def set_timestep(self, timestep): 

96 self.dt = timestep 

97 

98 def boltzmann_random(self, width, size): 

99 x = self.rng.random(size=size) 

100 y = self.rng.random(size=size) 

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

102 return z 

103 

104 def get_maxwell_boltzmann_velocities(self): 

105 natoms = len(self.atoms) 

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

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

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

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

110 

111 def step(self, forces=None): 

112 atoms = self.atoms 

113 

114 if forces is None: 

115 forces = atoms.get_forces(md=True) 

116 

117 self.v = atoms.get_velocities() 

118 

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

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

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

122 # - self.andersen_chance # andersen_chance <= andersen_prob 

123 # a dummy communicator will be used for serial runs 

124 

125 if self.fix_com: 

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

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

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

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

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

131 self.v += self.random_com_velocity 

132 

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

134 

135 # apply Andersen thermostat 

136 self.random_velocity = self.get_maxwell_boltzmann_velocities() 

137 self.andersen_chance = self.rng.random(size=self.v.shape) 

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

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

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

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

142 

143 x = atoms.get_positions() 

144 if self.fix_com: 

145 old_com = atoms.get_center_of_mass() 

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

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

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

149 if self.fix_com: 

150 atoms.set_center_of_mass(old_com) 

151 

152 # recalc velocities after RATTLE constraints are applied 

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

154 forces = atoms.get_forces(md=True) 

155 

156 # Update the velocities 

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

158 

159 if self.fix_com: 

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

161 

162 # Second part of RATTLE taken care of here 

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

164 

165 return forces