Coverage for /builds/ase/ase/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 :class:~ase.md.md.MolecularDynamics 

50 base class. 

51 

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

53 that acts on velocity components of randomly chosen particles. 

54 The algorithm randomly decorrelates velocities, so dynamical properties 

55 like diffusion or viscosity cannot be properly measured. 

56 

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

58 """ 

59 if 'communicator' in kwargs: 

60 msg = ( 

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

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

63 ) 

64 warnings.warn(msg, FutureWarning) 

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

66 self.temp = units.kB * temperature_K 

67 self.andersen_prob = andersen_prob 

68 self.fix_com = fixcm 

69 self.rng = rng 

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

71 

72 def set_temperature(self, temperature_K): 

73 self.temp = units.kB * temperature_K 

74 

75 def set_andersen_prob(self, andersen_prob): 

76 self.andersen_prob = andersen_prob 

77 

78 def set_timestep(self, timestep): 

79 self.dt = timestep 

80 

81 def boltzmann_random(self, width, size): 

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

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

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

85 return z 

86 

87 def get_maxwell_boltzmann_velocities(self): 

88 natoms = len(self.atoms) 

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

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

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

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

93 

94 def step(self, forces=None): 

95 atoms = self.atoms 

96 

97 if forces is None: 

98 forces = atoms.get_forces(md=True) 

99 

100 self.v = atoms.get_velocities() 

101 

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

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

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

105 # - self.andersen_chance # andersen_chance <= andersen_prob 

106 # a dummy communicator will be used for serial runs 

107 

108 if self.fix_com: 

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

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

111 self.random_com_velocity = ones( 

112 self.v.shape 

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

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

115 self.v += self.random_com_velocity 

116 

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

118 

119 # apply Andersen thermostat 

120 self.random_velocity = self.get_maxwell_boltzmann_velocities() 

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

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

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

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

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

126 ) 

127 

128 x = atoms.get_positions() 

129 if self.fix_com: 

130 old_com = atoms.get_center_of_mass() 

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

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

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

134 if self.fix_com: 

135 atoms.set_center_of_mass(old_com) 

136 

137 # recalc velocities after RATTLE constraints are applied 

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

139 forces = atoms.get_forces(md=True) 

140 

141 # Update the velocities 

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

143 

144 if self.fix_com: 

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

146 

147 # Second part of RATTLE taken care of here 

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

149 

150 return forces