Coverage for /builds/debichem-team/python-ase/ase/md/nose_hoover_chain.py: 100.00%

1from __future__ import annotations 

2 

3import numpy as np 

4 

5import ase.units 

6from ase import Atoms 

7from ase.md.md import MolecularDynamics 

8 

9# Coefficients for the fourth-order Suzuki-Yoshida integration scheme 

10# Ref: H. Yoshida, Phys. Lett. A 150, 5-7, 262-268 (1990). 

11# https://doi.org/10.1016/0375-9601(90)90092-3 

12FOURTH_ORDER_COEFFS = [ 

13 1 / (2 - 2 ** (1 / 3)), 

14 -(2 ** (1 / 3)) / (2 - 2 ** (1 / 3)), 

15 1 / (2 - 2 ** (1 / 3)), 

16] 

17 

18 

19class NoseHooverChainNVT(MolecularDynamics): 

20 """Isothermal molecular dynamics with Nose-Hoover chain. 

21 

22 This implementation is based on the Nose-Hoover chain equations and 

23 the Liouville-operator derived integrator for non-Hamiltonian systems [1-3]. 

24 

25 - [1] G. J. Martyna, M. L. Klein, and M. E. Tuckerman, J. Chem. Phys. 97, 

26 2635 (1992). https://doi.org/10.1063/1.463940 

27 - [2] M. E. Tuckerman, J. Alejandre, R. López-Rendón, A. L. Jochim, 

28 and G. J. Martyna, J. Phys. A: Math. Gen. 39, 5629 (2006). 

29 https://doi.org/10.1088/0305-4470/39/19/S18 

30 - [3] M. E. Tuckerman, Statistical Mechanics: Theory and Molecular 

31 Simulation, Oxford University Press (2010). 

32 

33 While the algorithm and notation for the thermostat are largely adapted 

34 from Ref. [4], the core equations are detailed in the implementation 

35 note available at 

36 https://github.com/lan496/lan496.github.io/blob/main/notes/nose_hoover_chain/main.pdf. 

37 

38 - [4] M. E. Tuckerman, Statistical Mechanics: Theory and Molecular 

39 Simulation, 2nd ed. (Oxford University Press, 2009). 

40 """ 

41 

42 def __init__( 

43 self, 

44 atoms: Atoms, 

45 timestep: float, 

46 temperature_K: float, 

47 tdamp: float, 

48 tchain: int = 3, 

49 tloop: int = 1, 

50 **kwargs, 

51 ): 

52 """ 

53 Parameters 

54 ---------- 

55 atoms: ase.Atoms 

56 The atoms object. 

57 timestep: float 

58 The time step in ASE time units. 

59 temperature_K: float 

60 The target temperature in K. 

61 tdamp: float 

62 The characteristic time scale for the thermostat in ASE time units. 

63 Typically, it is set to 100 times of `timestep`. 

64 tchain: int 

65 The number of thermostat variables in the Nose-Hoover chain. 

66 tloop: int 

67 The number of sub-steps in thermostat integration. 

68 trajectory: str or None 

69 If `trajectory` is str, `Trajectory` will be instantiated. 

70 Set `None` for no trajectory. 

71 logfile: IO or str or None 

72 If `logfile` is str, a file with that name will be opened. 

73 Set `-` to output into stdout. 

74 loginterval: int 

75 Write a log line for every `loginterval` time steps. 

76 **kwargs : dict, optional 

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

78 base class. 

79 """ 

80 super().__init__( 

81 atoms=atoms, 

82 timestep=timestep, 

83 **kwargs, 

84 ) 

85 assert self.masses.shape == (len(self.atoms), 1) 

86 

87 self._thermostat = NoseHooverChainThermostat( 

88 masses=self.masses, 

89 temperature_K=temperature_K, 

90 tdamp=tdamp, 

91 tchain=tchain, 

92 tloop=tloop, 

93 ) 

94 

95 # The following variables are updated during self.step() 

96 self._q = self.atoms.get_positions() 

97 self._p = self.atoms.get_momenta() 

98 

99 def step(self) -> None: 

100 dt2 = self.dt / 2 

101 self._p = self._thermostat.integrate_nhc(self._p, dt2) 

102 self._integrate_p(dt2) 

103 self._integrate_q(self.dt) 

104 self._integrate_p(dt2) 

105 self._p = self._thermostat.integrate_nhc(self._p, dt2) 

106 

107 self._update_atoms() 

108 

109 def get_conserved_energy(self) -> float: 

110 """Return the conserved energy-like quantity. 

111 

112 This method is mainly used for testing. 

113 """ 

114 conserved_energy = ( 

115 self.atoms.get_potential_energy(force_consistent=True) 

116 + self.atoms.get_kinetic_energy() 

117 + self._thermostat.get_thermostat_energy() 

118 ) 

119 return float(conserved_energy) 

120 

121 def _update_atoms(self) -> None: 

122 self.atoms.set_positions(self._q) 

123 self.atoms.set_momenta(self._p) 

124 

125 def _get_forces(self) -> np.ndarray: 

126 self._update_atoms() 

127 return self.atoms.get_forces(md=True) 

128 

129 def _integrate_q(self, delta: float) -> None: 

130 """Integrate exp(i * L_1 * delta)""" 

131 self._q += delta * self._p / self.masses 

132 

133 def _integrate_p(self, delta: float) -> None: 

134 """Integrate exp(i * L_2 * delta)""" 

135 forces = self._get_forces() 

136 self._p += delta * forces 

137 

138 

139class NoseHooverChainThermostat: 

140 """Nose-Hoover chain style thermostats. 

141 

142 See `NoseHooverChainNVT` for the references. 

143 """ 

144 def __init__( 

145 self, 

146 masses: np.ndarray, 

147 temperature_K: float, 

148 tdamp: float, 

149 tchain: int = 3, 

150 tloop: int = 1, 

151 ): 

152 """See `NoseHooverChainNVT` for the parameters.""" 

153 self._num_atoms = masses.shape[0] 

154 self._masses = masses # (num_atoms, 1) 

155 self._tdamp = tdamp 

156 self._tchain = tchain 

157 self._tloop = tloop 

158 

159 self._kT = ase.units.kB * temperature_K 

160 

161 assert tchain >= 1 

162 self._Q = np.zeros(tchain) 

163 self._Q[0] = 3 * self._num_atoms * self._kT * tdamp**2 

164 self._Q[1:] = self._kT * tdamp**2 

165 

166 # The following variables are updated during self.step() 

167 self._eta = np.zeros(self._tchain) 

168 self._p_eta = np.zeros(self._tchain) 

169 

170 def get_thermostat_energy(self) -> float: 

171 """Return energy-like contribution from the thermostat variables.""" 

172 energy = ( 

173 3 * self._num_atoms * self._kT * self._eta[0] 

174 + self._kT * np.sum(self._eta[1:]) 

175 + np.sum(0.5 * self._p_eta**2 / self._Q) 

176 ) 

177 return float(energy) 

178 

179 def integrate_nhc(self, p: np.ndarray, delta: float) -> np.ndarray: 

180 """Integrate exp(i * L_NHC * delta) and update momenta `p`.""" 

181 for _ in range(self._tloop): 

182 for coeff in FOURTH_ORDER_COEFFS: 

183 p = self._integrate_nhc_loop( 

184 p, coeff * delta / len(FOURTH_ORDER_COEFFS) 

185 ) 

186 

187 return p 

188 

189 def _integrate_nhc_loop(self, p: np.ndarray, delta: float) -> np.ndarray: 

190 delta2 = delta / 2 

191 delta4 = delta / 4 

192 

193 def _integrate_p_eta_j(p: np.ndarray, j: int) -> None: 

194 if j < self._tchain - 1: 

195 self._p_eta[j] *= np.exp( 

196 -delta4 * self._p_eta[j + 1] / self._Q[j + 1] 

197 ) 

198 

199 if j == 0: 

200 g_j = np.sum(p**2 / self._masses) \ 

201 - 3 * self._num_atoms * self._kT 

202 else: 

203 g_j = self._p_eta[j - 1] ** 2 / self._Q[j - 1] - self._kT 

204 self._p_eta[j] += delta2 * g_j 

205 

206 if j < self._tchain - 1: 

207 self._p_eta[j] *= np.exp( 

208 -delta4 * self._p_eta[j + 1] / self._Q[j + 1] 

209 ) 

210 

211 def _integrate_eta() -> None: 

212 self._eta += delta * self._p_eta / self._Q 

213 

214 def _integrate_nhc_p(p: np.ndarray) -> None: 

215 p *= np.exp(-delta * self._p_eta[0] / self._Q[0]) 

216 

217 for j in range(self._tchain): 

218 _integrate_p_eta_j(p, self._tchain - j - 1) 

219 _integrate_eta() 

220 _integrate_nhc_p(p) 

221 for j in range(self._tchain): 

222 _integrate_p_eta_j(p, j) 

223 

224 return p