Coverage for /builds/ericyuan00000/ase/ase/md/nose_hoover

1# fmt: off

3from __future__ import annotations

5import numpy as np

6from scipy.special import exprel

8import ase.units

9from ase import Atoms

10from ase.md.md import MolecularDynamics

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

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

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

15FOURTH_ORDER_COEFFS = [

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

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

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

19]

22class NoseHooverChainNVT(MolecularDynamics):

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

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

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

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

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

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

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

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

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

34 Simulation, Oxford University Press (2010).

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

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

38 note available at

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

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

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

43 """

45 def __init__(

46 self,

47 atoms: Atoms,

48 timestep: float,

49 temperature_K: float,

50 tdamp: float,

51 tchain: int = 3,

52 tloop: int = 1,

53 **kwargs,

54 ):

55 """

56 Parameters

57 ----------

58 atoms: ase.Atoms

59 The atoms object.

60 timestep: float

61 The time step in ASE time units.

62 temperature_K: float

63 The target temperature in K.

64 tdamp: float

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

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

67 tchain: int

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

69 tloop: int

70 The number of sub-steps in thermostat integration.

71 **kwargs : dict, optional

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

73 base class.

74 """

75 super().__init__(

76 atoms=atoms,

77 timestep=timestep,

78 **kwargs,

79 )

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

82 num_atoms = self.atoms.get_global_number_of_atoms()

83 self._thermostat = NoseHooverChainThermostat(

84 num_atoms_global=num_atoms,

85 masses=self.masses,

86 temperature_K=temperature_K,

87 tdamp=tdamp,

88 tchain=tchain,

89 tloop=tloop,

90 )

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

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

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

96 def step(self) -> None:

97 dt2 = self.dt / 2

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

99 self._integrate_p(dt2)

100 self._integrate_q(self.dt)

101 self._integrate_p(dt2)

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

103

104 self._update_atoms()

105

106 def get_conserved_energy(self) -> float:

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

108

109 This method is mainly used for testing.

110 """

111 conserved_energy = (

112 self.atoms.get_potential_energy(force_consistent=True)

113 + self.atoms.get_kinetic_energy()

114 + self._thermostat.get_thermostat_energy()

115 )

116 return float(conserved_energy)

117

118 def _update_atoms(self) -> None:

119 self.atoms.set_positions(self._q)

120 self.atoms.set_momenta(self._p)

121

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

123 self._update_atoms()

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

125

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

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

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

129

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

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

132 forces = self._get_forces()

133 self._p += delta * forces

134

135

136class NoseHooverChainThermostat:

137 """Nose-Hoover chain style thermostats.

138

139 See `NoseHooverChainNVT` for the references.

140 """

141 def __init__(

142 self,

143 num_atoms_global: int,

144 masses: np.ndarray,

145 temperature_K: float,

146 tdamp: float,

147 tchain: int = 3,

148 tloop: int = 1,

149 ):

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

151 self._num_atoms_global = num_atoms_global

152 self._masses = masses # (len(atoms), 1)

153 self._tdamp = tdamp

154 self._tchain = tchain

155 self._tloop = tloop

156

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

158

159 assert tchain >= 1

160 self._Q = np.zeros(tchain)

161 self._Q[0] = 3 * self._num_atoms_global * self._kT * tdamp**2

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

163

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

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

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

167

168 def get_thermostat_energy(self) -> float:

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

170 energy = (

171 3 * self._num_atoms_global * self._kT * self._eta[0]

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

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

174 )

175 return float(energy)

176

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

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

179 for _ in range(self._tloop):

180 for coeff in FOURTH_ORDER_COEFFS:

181 p = self._integrate_nhc_loop(

182 p, coeff * delta / len(FOURTH_ORDER_COEFFS)

183 )

184

185 return p

186

187 def _integrate_p_eta_j(self, p: np.ndarray, j: int,

188 delta2: float, delta4: float) -> None:

189 if j < self._tchain - 1:

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

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

192 )

193

194 if j == 0:

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

196 - 3 * self._num_atoms_global * self._kT

197 else:

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

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

200

201 if j < self._tchain - 1:

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

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

204 )

205

206 def _integrate_eta(self, delta: float) -> None:

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

208

209 def _integrate_nhc_p(self, p: np.ndarray, delta: float) -> None:

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

211

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

213 delta2 = delta / 2

214 delta4 = delta / 4

215

216 for j in range(self._tchain):

217 self._integrate_p_eta_j(p, self._tchain - j - 1, delta2, delta4)

218 self._integrate_eta(delta)

219 self._integrate_nhc_p(p, delta)

220 for j in range(self._tchain):

221 self._integrate_p_eta_j(p, j, delta2, delta4)

222

223 return p

224

225

226class IsotropicMTKNPT(MolecularDynamics):

227 """Isothermal-isobaric molecular dynamics with isotropic volume fluctuations

228 by Martyna-Tobias-Klein (MTK) method [1].

229

230 See also `NoseHooverChainNVT` for the references.

231

232 - [1] G. J. Martyna, D. J. Tobias, and M. L. Klein, J. Chem. Phys. 101,

233 4177-4189 (1994). https://doi.org/10.1063/1.467468

234 """

235 def __init__(

236 self,

237 atoms: Atoms,

238 timestep: float,

239 temperature_K: float,

240 pressure_au: float,

241 tdamp: float,

242 pdamp: float,

243 tchain: int = 3,

244 pchain: int = 3,

245 tloop: int = 1,

246 ploop: int = 1,

247 **kwargs,

248 ):

249 """

250 Parameters

251 ----------

252 atoms: ase.Atoms

253 The atoms object.

254 timestep: float

255 The time step in ASE time units.

256 temperature_K: float

257 The target temperature in K.

258 pressure_au: float

259 The external pressure in eV/Ang^3.

260 tdamp: float

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

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

263 pdamp: float

264 The characteristic time scale for the barostat in ASE time units.

265 Typically, it is set to 1000 times of `timestep`.

266 tchain: int

267 The number of thermostat variables in the Nose-Hoover thermostat.

268 pchain: int

269 The number of barostat variables in the MTK barostat.

270 tloop: int

271 The number of sub-steps in thermostat integration.

272 ploop: int

273 The number of sub-steps in barostat integration.

274 **kwargs : dict, optional

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

276 base class.

277 """

278 super().__init__(

279 atoms=atoms,

280 timestep=timestep,

281 **kwargs,

282 )

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

284

285 if len(atoms.constraints) > 0:

286 raise NotImplementedError(

287 "Current implementation does not support constraints"

288 )

289

290 self._num_atoms_global = self.atoms.get_global_number_of_atoms()

291 self._thermostat = NoseHooverChainThermostat(

292 num_atoms_global=self._num_atoms_global,

293 masses=self.masses,

294 temperature_K=temperature_K,

295 tdamp=tdamp,

296 tchain=tchain,

297 tloop=tloop,

298 )

299 self._barostat = IsotropicMTKBarostat(

300 num_atoms_global=self._num_atoms_global,

301 masses=self.masses,

302 temperature_K=temperature_K,

303 pdamp=pdamp,

304 pchain=pchain,

305 ploop=ploop,

306 )

307

308 self._temperature_K = temperature_K

309 self._pressure_au = pressure_au

310

311 self._kT = ase.units.kB * self._temperature_K

312 self._volume0 = self.atoms.get_volume()

313 self._cell0 = np.array(self.atoms.get_cell())

314

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

316 self._q = self.atoms.get_positions() # positions

317 self._p = self.atoms.get_momenta() # momenta

318 self._eps = 0.0 # volume

319 self._p_eps = 0.0 # volume momenta

320

321 def step(self) -> None:

322 dt2 = self.dt / 2

323

324 self._p_eps = self._barostat.integrate_nhc_baro(self._p_eps, dt2)

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

326 self._integrate_p_cell(dt2)

327 self._integrate_p(dt2)

328 self._integrate_q(self.dt)

329 self._integrate_q_cell(self.dt)

330 self._integrate_p(dt2)

331 self._integrate_p_cell(dt2)

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

333 self._p_eps = self._barostat.integrate_nhc_baro(self._p_eps, dt2)

334

335 self._update_atoms()

336

337 def get_conserved_energy(self) -> float:

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

339

340 This method is mainly used for testing.

341 """

342 conserved_energy = (

343 self.atoms.get_potential_energy(force_consistent=True)

344 + self.atoms.get_kinetic_energy()

345 + self._thermostat.get_thermostat_energy()

346 + self._barostat.get_barostat_energy()

347 + self._p_eps * self._p_eps / (2 * self._barostat.W)

348 + self._pressure_au * self._get_volume()

349 )

350 return float(conserved_energy)

351

352 def _update_atoms(self) -> None:

353 self.atoms.set_positions(self._q)

354 self.atoms.set_momenta(self._p)

355 cell = self._cell0 * np.exp(self._eps)

356 # Never set scale_atoms=True

357 self.atoms.set_cell(cell, scale_atoms=False)

358

359 def _get_volume(self) -> float:

360 return self._volume0 * np.exp(3 * self._eps)

361

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

363 self._update_atoms()

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

365

366 def _get_pressure(self) -> np.ndarray:

367 self._update_atoms()

368 stress = self.atoms.get_stress(voigt=False, include_ideal_gas=True)

369 pressure = -np.trace(stress) / 3

370 return pressure

371

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

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

374 x = delta * self._p_eps / self._barostat.W

375 self._q = (

376 self._q * np.exp(x)

377 + self._p * delta / self.masses * exprel(x)

378 )

379

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

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

382 x = (1 + 1 / self._num_atoms_global) * self._p_eps * delta \

383 / self._barostat.W

384 forces = self._get_forces()

385 self._p = self._p * np.exp(-x) + delta * forces * exprel(-x)

386

387 def _integrate_q_cell(self, delta: float) -> None:

388 """Integrate exp(i * L_(epsilon, 1) * delta)"""

389 self._eps += delta * self._p_eps / self._barostat.W

390

391 def _integrate_p_cell(self, delta: float) -> None:

392 """Integrate exp(i * L_(epsilon, 2) * delta)"""

393 pressure = self._get_pressure()

394 volume = self._get_volume()

395 G = (

396 3 * volume * (pressure - self._pressure_au)

397 + np.sum(self._p**2 / self.masses) / self._num_atoms_global

398 )

399 self._p_eps += delta * G

400

401

402class IsotropicMTKBarostat:

403 """MTK barostat for isotropic volume fluctuations.

404

405 See `IsotropicMTKNPT` for the references.

406 """

407 def __init__(

408 self,

409 num_atoms_global: int,

410 masses: np.ndarray,

411 temperature_K: float,

412 pdamp: float,

413 pchain: int = 3,

414 ploop: int = 1,

415 ):

416 self._num_atoms_global = num_atoms_global

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

418 self._pdamp = pdamp

419 self._pchain = pchain

420 self._ploop = ploop

421

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

423

424 self._W = (3 * self._num_atoms_global + 3) * self._kT * self._pdamp**2

425

426 assert pchain >= 1

427 self._R = np.zeros(self._pchain)

428 self._R[0] = 9 * self._kT * self._pdamp**2

429 self._R[1:] = self._kT * self._pdamp**2

430

431 self._xi = np.zeros(self._pchain) # barostat coordinates

432 self._p_xi = np.zeros(self._pchain)

433

434 @property

435 def W(self) -> float:

436 """Virtual mass for barostat momenta `p_xi`."""

437 return self._W

438

439 def get_barostat_energy(self) -> float:

440 """Return energy-like contribution from the barostat variables."""

441 energy = (

442 + np.sum(0.5 * self._p_xi**2 / self._R)

443 + self._kT * np.sum(self._xi)

444 )

445 return float(energy)

446

447 def integrate_nhc_baro(self, p_eps: float, delta: float) -> float:

448 """Integrate exp(i * L_NHC-baro * delta)"""

449 for _ in range(self._ploop):

450 for coeff in FOURTH_ORDER_COEFFS:

451 p_eps = self._integrate_nhc_baro_loop(

452 p_eps, coeff * delta / len(FOURTH_ORDER_COEFFS)

453 )

454 return p_eps

455

456 def _integrate_nhc_baro_loop(self, p_eps: float, delta: float) -> float:

457 delta2 = delta / 2

458 delta4 = delta / 4

459

460 for j in range(self._pchain):

461 self._integrate_p_xi_j(p_eps, self._pchain - j - 1, delta2, delta4)

462 self._integrate_xi(delta)

463 p_eps = self._integrate_nhc_p_eps(p_eps, delta)

464 for j in range(self._pchain):

465 self._integrate_p_xi_j(p_eps, j, delta2, delta4)

466

467 return p_eps

468

469 def _integrate_p_xi_j(self, p_eps: float, j: int,

470 delta2: float, delta4: float) -> None:

471 if j < self._pchain - 1:

472 self._p_xi[j] *= np.exp(

473 -delta4 * self._p_xi[j + 1] / self._R[j + 1]

474 )

475

476 if j == 0:

477 g_j = p_eps ** 2 / self._W - self._kT

478 else:

479 g_j = self._p_xi[j - 1] ** 2 / self._R[j - 1] - self._kT

480 self._p_xi[j] += delta2 * g_j

481

482 if j < self._pchain - 1:

483 self._p_xi[j] *= np.exp(

484 -delta4 * self._p_xi[j + 1] / self._R[j + 1]

485 )

486

487 def _integrate_xi(self, delta) -> None:

488 self._xi += delta * self._p_xi / self._R

489

490 def _integrate_nhc_p_eps(self, p_eps: float, delta: float) -> float:

491 p_eps_new = p_eps * float(

492 np.exp(-delta * self._p_xi[0] / self._R[0])

493 )

494 return p_eps_new

Coverage for /builds/ericyuan00000/ase/ase/md/nose_hoover_chain.py: 99.49%

195 statements