Coverage for ase / md / nvtberendsen.py: 84.00%

1"""Berendsen NVT dynamics class.""" 

2 

3import numpy as np 

4 

5from ase import Atoms 

6from ase.md.md import MolecularDynamics 

7 

8 

9class NVTBerendsen(MolecularDynamics): 

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

11 

12 Berendsen *et al.*, J. Chem. Phys. 81, 3684–3690 (1984). 

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

14 """ 

15 

16 def __init__( 

17 self, 

18 atoms: Atoms, 

19 timestep: float, 

20 temperature: float | None = None, 

21 taut: float | None = None, 

22 fixcm: bool = True, 

23 *, 

24 temperature_K: float | None = None, 

25 **kwargs, 

26 ): 

27 """ 

28 Parameters 

29 ---------- 

30 atoms: Atoms object 

31 The list of atoms. 

32 

33 timestep: float 

34 The time step in ASE time units. 

35 

36 temperature: float 

37 The desired temperature, in Kelvin. 

38 

39 temperature_K: float 

40 Alias for *temperature* 

41 

42 taut: float 

43 Time constant for Berendsen temperature coupling in ASE 

44 time units. 

45 

46 fixcm: bool (optional) 

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

48 kept unperturbed. Default: True. 

49 

50 **kwargs : dict, optional 

51 Additional arguments passed to :class:~ase.md.md.MolecularDynamics 

52 base class. 

53 

54 """ 

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

56 

57 if taut is None: 

58 raise TypeError("Missing 'taut' argument.") 

59 self.taut = taut 

60 self.temperature = self._process_temperature( 

61 temperature, temperature_K, 'K' 

62 ) 

63 

64 self.fix_com = fixcm # will the center of mass be held fixed? 

65 

66 def set_taut(self, taut): 

67 self.taut = taut 

68 

69 def get_taut(self): 

70 return self.taut 

71 

72 def set_temperature(self, temperature=None, *, temperature_K=None): 

73 self.temperature = self._process_temperature( 

74 temperature, temperature_K, 'K' 

75 ) 

76 

77 def get_temperature(self): 

78 return self.temperature 

79 

80 def set_timestep(self, timestep): 

81 self.dt = timestep 

82 

83 def get_timestep(self): 

84 return self.dt 

85 

86 def scale_velocities(self): 

87 """Do the NVT Berendsen velocity scaling""" 

88 tautscl = self.dt / self.taut 

89 old_temperature = self.atoms.get_temperature() 

90 

91 scl_temperature = np.sqrt( 

92 1.0 + (self.temperature / old_temperature - 1.0) * tautscl 

93 ) 

94 # Limit the velocity scaling to reasonable values 

95 if scl_temperature > 1.1: 

96 scl_temperature = 1.1 

97 if scl_temperature < 0.9: 

98 scl_temperature = 0.9 

99 

100 p = self.atoms.get_momenta() 

101 p = scl_temperature * p 

102 self.atoms.set_momenta(p) 

103 return 

104 

105 def step(self, forces=None): 

106 """Move one timestep forward using Berenden NVT molecular dynamics.""" 

107 self.scale_velocities() 

108 

109 # one step velocity verlet 

110 atoms = self.atoms 

111 

112 if forces is None: 

113 forces = atoms.get_forces(md=True) 

114 

115 p = self.atoms.get_momenta() 

116 p += 0.5 * self.dt * forces 

117 

118 if self.fix_com: 

119 # calculate the center of mass 

120 # momentum and subtract it 

121 psum = p.sum(axis=0) / float(len(p)) 

122 p = p - psum 

123 

124 self.atoms.set_positions( 

125 self.atoms.get_positions() 

126 + self.dt * p / self.atoms.get_masses()[:, np.newaxis] 

127 ) 

128 

129 # We need to store the momenta on the atoms before calculating 

130 # the forces, as in a parallel Asap calculation atoms may 

131 # migrate during force calculations, and the momenta need to 

132 # migrate along with the atoms. For the same reason, we 

133 # cannot use self.masses in the line above. 

134 

135 self.atoms.set_momenta(p) 

136 forces = self.atoms.get_forces(md=True) 

137 atoms.set_momenta(self.atoms.get_momenta() + 0.5 * self.dt * forces) 

138 

139 return forces