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 def __init__( 

15 self, 

16 atoms: Atoms, 

17 timestep: float, 

18 temperature_K: float, 

19 andersen_prob: float, 

20 fixcm: bool = True, 

21 rng=random, 

22 **kwargs, 

23 ): 

24 """ 

25 Parameters 

26 ---------- 

27 atoms: Atoms object 

28 The list of atoms. 

29 

30 timestep: float 

31 The time step in ASE time units. 

32 

33 temperature_K: float 

34 The desired temperature, in Kelvin. 

35 

36 andersen_prob: float 

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

38 With this probability atoms get assigned random velocity components. 

39 

40 fixcm: bool (optional) 

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

42 kept unperturbed. Default: True. 

43 

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

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

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

47 

48 **kwargs : dict, optional 

49 Additional arguments passed to 

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

51 base class. 

52 

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

54 that acts on velocity components of randomly chosen particles. 

55 The algorithm randomly decorrelates velocities, so dynamical properties 

56 like diffusion or viscosity cannot be properly measured. 

57 

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

59 """ 

60 if 'communicator' in kwargs: 

61 msg = ( 

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

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

64 ) 

65 warnings.warn(msg, FutureWarning) 

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

67 self.temp = units.kB * temperature_K 

68 self.andersen_prob = andersen_prob 

69 self.fix_com = fixcm 

70 self.rng = rng 

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

72 

73 def set_temperature(self, temperature_K): 

74 self.temp = units.kB * temperature_K 

75 

76 def set_andersen_prob(self, andersen_prob): 

77 self.andersen_prob = andersen_prob 

78 

79 def set_timestep(self, timestep): 

80 self.dt = timestep 

81 

82 def boltzmann_random(self, width, size): 

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

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

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

86 return z 

87 

88 def get_maxwell_boltzmann_velocities(self): 

89 natoms = len(self.atoms) 

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

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

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

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

94 

95 def step(self, forces=None): 

96 atoms = self.atoms 

97 

98 if forces is None: 

99 forces = atoms.get_forces(md=True) 

100 

101 self.v = atoms.get_velocities() 

102 

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

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

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

106 # - self.andersen_chance # andersen_chance <= andersen_prob 

107 # a dummy communicator will be used for serial runs 

108 

109 if self.fix_com: 

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

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

112 self.random_com_velocity = ones( 

113 self.v.shape 

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

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

116 self.v += self.random_com_velocity 

117 

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

119 

120 # apply Andersen thermostat 

121 self.random_velocity = self.get_maxwell_boltzmann_velocities() 

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

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

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

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

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

127 ) 

128 

129 x = atoms.get_positions() 

130 if self.fix_com: 

131 old_com = atoms.get_center_of_mass() 

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

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

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

135 if self.fix_com: 

136 atoms.set_center_of_mass(old_com) 

137 

138 # recalc velocities after RATTLE constraints are applied 

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

140 forces = atoms.get_forces(md=True) 

141 

142 # Update the velocities 

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

144 

145 if self.fix_com: 

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

147 

148 # Second part of RATTLE taken care of here 

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

150 

151 return forces