Coverage for /builds/kinetik161/ase/ase/md/nvtberendsen.py: 83.02%

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

2from typing import IO, Optional, Union 

3 

4import numpy as np 

5 

6from ase import Atoms 

7from ase.md.md import MolecularDynamics 

8from ase.parallel import world 

9 

10 

11class NVTBerendsen(MolecularDynamics): 

12 def __init__( 

13 self, 

14 atoms: Atoms, 

15 timestep: float, 

16 temperature: Optional[float] = None, 

17 taut: Optional[float] = None, 

18 fixcm: bool = True, 

19 *, 

20 temperature_K: Optional[float] = None, 

21 trajectory: Optional[str] = None, 

22 logfile: Optional[Union[IO, str]] = None, 

23 loginterval: int = 1, 

24 communicator=world, 

25 append_trajectory: bool = False, 

26 ): 

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

28 

29 Parameters: 

30 

31 atoms: Atoms object 

32 The list of atoms. 

33 

34 timestep: float 

35 The time step in ASE time units. 

36 

37 temperature: float 

38 The desired temperature, in Kelvin. 

39 

40 temperature_K: float 

41 Alias for *temperature* 

42 

43 taut: float 

44 Time constant for Berendsen temperature coupling in ASE 

45 time units. 

46 

47 fixcm: bool (optional) 

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

49 kept unperturbed. Default: True. 

50 

51 trajectory: Trajectory object or str (optional) 

52 Attach trajectory object. If *trajectory* is a string a 

53 Trajectory will be constructed. Use *None* for no 

54 trajectory. 

55 

56 logfile: file object or str (optional) 

57 If *logfile* is a string, a file with that name will be opened. 

58 Use '-' for stdout. 

59 

60 loginterval: int (optional) 

61 Only write a log line for every *loginterval* time steps. 

62 Default: 1 

63 

64 append_trajectory: boolean (optional) 

65 Defaults to False, which causes the trajectory file to be 

66 overwriten each time the dynamics is restarted from scratch. 

67 If True, the new structures are appended to the trajectory 

68 file instead. 

69 

70 """ 

71 

72 MolecularDynamics.__init__(self, atoms, timestep, trajectory, 

73 logfile, loginterval, 

74 append_trajectory=append_trajectory) 

75 if taut is None: 

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

77 self.taut = taut 

78 self.temperature = self._process_temperature(temperature, 

79 temperature_K, 'K') 

80 

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

82 self.communicator = communicator 

83 

84 def set_taut(self, taut): 

85 self.taut = taut 

86 

87 def get_taut(self): 

88 return self.taut 

89 

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

91 self.temperature = self._process_temperature(temperature, 

92 temperature_K, 'K') 

93 

94 def get_temperature(self): 

95 return self.temperature 

96 

97 def set_timestep(self, timestep): 

98 self.dt = timestep 

99 

100 def get_timestep(self): 

101 return self.dt 

102 

103 def scale_velocities(self): 

104 """ Do the NVT Berendsen velocity scaling """ 

105 tautscl = self.dt / self.taut 

106 old_temperature = self.atoms.get_temperature() 

107 

108 scl_temperature = np.sqrt(1.0 + 

109 (self.temperature / old_temperature - 1.0) * 

110 tautscl) 

111 # Limit the velocity scaling to reasonable values 

112 if scl_temperature > 1.1: 

113 scl_temperature = 1.1 

114 if scl_temperature < 0.9: 

115 scl_temperature = 0.9 

116 

117 p = self.atoms.get_momenta() 

118 p = scl_temperature * p 

119 self.atoms.set_momenta(p) 

120 return 

121 

122 def step(self, forces=None): 

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

124 self.scale_velocities() 

125 

126 # one step velocity verlet 

127 atoms = self.atoms 

128 

129 if forces is None: 

130 forces = atoms.get_forces(md=True) 

131 

132 p = self.atoms.get_momenta() 

133 p += 0.5 * self.dt * forces 

134 

135 if self.fix_com: 

136 # calculate the center of mass 

137 # momentum and subtract it 

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

139 p = p - psum 

140 

141 self.atoms.set_positions( 

142 self.atoms.get_positions() + 

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

144 

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

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

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

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

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

150 

151 self.atoms.set_momenta(p) 

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

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

154 

155 return forces