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

1"""Andersen dynamics class.""" 

2 

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

4 

5from ase import Atoms, units 

6from ase.md.md import MolecularDynamics 

7 

8 

9class Andersen(MolecularDynamics): 

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

11 

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

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

14 """ 

15 

16 def __init__( 

17 self, 

18 atoms: Atoms, 

19 timestep: float, 

20 temperature_K: float, 

21 andersen_prob: float, 

22 fixcm: bool = True, 

23 rng=random, 

24 **kwargs, 

25 ): 

26 """ 

27 Parameters 

28 ---------- 

29 atoms: Atoms object 

30 The list of atoms. 

31 

32 timestep: float 

33 The time step in ASE time units. 

34 

35 temperature_K: float 

36 The desired temperature, in Kelvin. 

37 

38 andersen_prob: float 

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

40 With this probability atoms get assigned random velocity components. 

41 

42 fixcm: bool (optional) 

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

44 kept unperturbed. Default: True. 

45 

46 rng: RNG object, default: ``numpy.random`` 

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

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

49 

50 **kwargs : dict, optional 

51 Additional arguments passed to 

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

53 base class. 

54 

55 Notes 

56 ----- 

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

58 that acts on velocity components of randomly chosen particles. 

59 The algorithm randomly decorrelates velocities, so dynamical properties 

60 like diffusion or viscosity cannot be properly measured. 

61 

62 """ 

63 self.temp = units.kB * temperature_K 

64 self.andersen_prob = andersen_prob 

65 self.fix_com = fixcm 

66 self.rng = rng 

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

68 

69 def set_temperature(self, temperature_K): 

70 self.temp = units.kB * temperature_K 

71 

72 def set_andersen_prob(self, andersen_prob): 

73 self.andersen_prob = andersen_prob 

74 

75 def set_timestep(self, timestep): 

76 self.dt = timestep 

77 

78 def boltzmann_random(self, width, size): 

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

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

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

82 return z 

83 

84 def get_maxwell_boltzmann_velocities(self): 

85 natoms = len(self.atoms) 

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

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

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

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

90 

91 def step(self, forces=None): 

92 atoms = self.atoms 

93 

94 if forces is None: 

95 forces = atoms.get_forces(md=True) 

96 

97 self.v = atoms.get_velocities() 

98 

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

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

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

102 # - self.andersen_chance # andersen_chance <= andersen_prob 

103 # a dummy communicator will be used for serial runs 

104 

105 if self.fix_com: 

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

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

108 self.random_com_velocity = ones( 

109 self.v.shape 

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

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

112 self.v += self.random_com_velocity 

113 

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

115 

116 # apply Andersen thermostat 

117 self.random_velocity = self.get_maxwell_boltzmann_velocities() 

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

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

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

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

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

123 ) 

124 

125 x = atoms.get_positions() 

126 if self.fix_com: 

127 old_com = atoms.get_center_of_mass() 

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

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

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

131 if self.fix_com: 

132 atoms.set_center_of_mass(old_com) 

133 

134 # recalc velocities after RATTLE constraints are applied 

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

136 forces = atoms.get_forces(md=True) 

137 

138 # Update the velocities 

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

140 

141 if self.fix_com: 

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

143 

144 # Second part of RATTLE taken care of here 

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

146 

147 return forces