Coverage for /builds/ase/ase/ase/md/nvtberendsen.py: 80.36%

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

2 

3import warnings 

4from typing import Optional 

5 

6import numpy as np 

7 

8from ase import Atoms 

9from ase.md.md import MolecularDynamics 

10 

11 

12class NVTBerendsen(MolecularDynamics): 

13 def __init__( 

14 self, 

15 atoms: Atoms, 

16 timestep: float, 

17 temperature: Optional[float] = None, 

18 taut: Optional[float] = None, 

19 fixcm: bool = True, 

20 *, 

21 temperature_K: Optional[float] = None, 

22 **kwargs, 

23 ): 

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

25 

26 Parameters 

27 ---------- 

28 atoms: Atoms object 

29 The list of atoms. 

30 

31 timestep: float 

32 The time step in ASE time units. 

33 

34 temperature: float 

35 The desired temperature, in Kelvin. 

36 

37 temperature_K: float 

38 Alias for *temperature* 

39 

40 taut: float 

41 Time constant for Berendsen temperature coupling in ASE 

42 time units. 

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

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

50 base class. 

51 

52 """ 

53 if 'communicator' in kwargs: 

54 msg = ( 

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

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

57 ) 

58 warnings.warn(msg, FutureWarning) 

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

60 

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

62 

63 if taut is None: 

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

65 self.taut = taut 

66 self.temperature = self._process_temperature( 

67 temperature, temperature_K, 'K' 

68 ) 

69 

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

71 

72 def set_taut(self, taut): 

73 self.taut = taut 

74 

75 def get_taut(self): 

76 return self.taut 

77 

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

79 self.temperature = self._process_temperature( 

80 temperature, temperature_K, 'K' 

81 ) 

82 

83 def get_temperature(self): 

84 return self.temperature 

85 

86 def set_timestep(self, timestep): 

87 self.dt = timestep 

88 

89 def get_timestep(self): 

90 return self.dt 

91 

92 def scale_velocities(self): 

93 """Do the NVT Berendsen velocity scaling""" 

94 tautscl = self.dt / self.taut 

95 old_temperature = self.atoms.get_temperature() 

96 

97 scl_temperature = np.sqrt( 

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

99 ) 

100 # Limit the velocity scaling to reasonable values 

101 if scl_temperature > 1.1: 

102 scl_temperature = 1.1 

103 if scl_temperature < 0.9: 

104 scl_temperature = 0.9 

105 

106 p = self.atoms.get_momenta() 

107 p = scl_temperature * p 

108 self.atoms.set_momenta(p) 

109 return 

110 

111 def step(self, forces=None): 

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

113 self.scale_velocities() 

114 

115 # one step velocity verlet 

116 atoms = self.atoms 

117 

118 if forces is None: 

119 forces = atoms.get_forces(md=True) 

120 

121 p = self.atoms.get_momenta() 

122 p += 0.5 * self.dt * forces 

123 

124 if self.fix_com: 

125 # calculate the center of mass 

126 # momentum and subtract it 

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

128 p = p - psum 

129 

130 self.atoms.set_positions( 

131 self.atoms.get_positions() 

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

133 ) 

134 

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

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

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

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

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

140 

141 self.atoms.set_momenta(p) 

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

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

144 

145 return forces