Coverage for ase / md / andersen.py: 93.55%

1"""Andersen dynamics class.""" 

2 

3import warnings 

4 

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

6 

7from ase import Atoms, units 

8from ase.md.md import MolecularDynamics 

9 

10 

11class Andersen(MolecularDynamics): 

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

13 

14 Andersen, J. Chem. Phys. 72, 2384–2393 (1980). 

15 https://doi.org/10.1063/1.439486 

16 """ 

17 

18 def __init__( 

19 self, 

20 atoms: Atoms, 

21 timestep: float, 

22 temperature_K: float, 

23 andersen_prob: float, 

24 fixcm: bool = True, 

25 rng=random, 

26 **kwargs, 

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 **kwargs : dict, optional 

53 Additional arguments passed to 

54 :class:`~ase.md.md.MolecularDynamics` 

55 base class. 

56 

57 Notes 

58 ----- 

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

60 that acts on velocity components of randomly chosen particles. 

61 The algorithm randomly decorrelates velocities, so dynamical properties 

62 like diffusion or viscosity cannot be properly measured. 

63 

64 """ 

65 if 'communicator' in kwargs: 

66 msg = ( 

67 '`communicator` has been deprecated since ASE 3.25.0 ' 

68 'and will be removed in ASE 3.26.0. Use `comm` instead.' 

69 ) 

70 warnings.warn(msg, FutureWarning) 

71 kwargs['comm'] = kwargs.pop('communicator') 

72 self.temp = units.kB * temperature_K 

73 self.andersen_prob = andersen_prob 

74 self.fix_com = fixcm 

75 self.rng = rng 

76 MolecularDynamics.__init__(self, atoms, timestep, **kwargs) 

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(size=size) 

89 y = self.rng.random(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( 

118 self.v.shape 

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

120 self.comm.broadcast(self.random_com_velocity, 0) 

121 self.v += self.random_com_velocity 

122 

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

124 

125 # apply Andersen thermostat 

126 self.random_velocity = self.get_maxwell_boltzmann_velocities() 

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

128 self.comm.broadcast(self.random_velocity, 0) 

129 self.comm.broadcast(self.andersen_chance, 0) 

130 self.v[self.andersen_chance <= self.andersen_prob] = ( 

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

132 ) 

133 

134 x = atoms.get_positions() 

135 if self.fix_com: 

136 old_com = atoms.get_center_of_mass() 

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

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

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

140 if self.fix_com: 

141 atoms.set_center_of_mass(old_com) 

142 

143 # recalc velocities after RATTLE constraints are applied 

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

145 forces = atoms.get_forces(md=True) 

146 

147 # Update the velocities 

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

149 

150 if self.fix_com: 

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

152 

153 # Second part of RATTLE taken care of here 

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

155 

156 return forces