Coverage for /builds/ase/ase/ase/md/nose_hoover_chain.py: 99.36%

1# fmt: off 

2 

3from __future__ import annotations 

4 

5import numpy as np 

6from scipy.special import exprel 

7 

8import ase.units 

9from ase import Atoms 

10from ase.md.md import MolecularDynamics 

11 

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] 

20 

21 

22class NoseHooverChainNVT(MolecularDynamics): 

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

24 

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

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

27 

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). 

35 

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. 

40 

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

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

43 """ 

44 

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) 

81 

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 ) 

91 

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

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

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

95 

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 temperature_K=temperature_K, 

302 pdamp=pdamp, 

303 pchain=pchain, 

304 ploop=ploop, 

305 ) 

306 

307 self._temperature_K = temperature_K 

308 self._pressure_au = pressure_au 

309 

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

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

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

313 

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

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

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

317 self._eps = 0.0 # volume 

318 self._p_eps = 0.0 # volume momenta 

319 

320 def step(self) -> None: 

321 dt2 = self.dt / 2 

322 

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

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

325 self._integrate_p_cell(dt2) 

326 self._integrate_p(dt2) 

327 self._integrate_q(self.dt) 

328 self._integrate_q_cell(self.dt) 

329 self._integrate_p(dt2) 

330 self._integrate_p_cell(dt2) 

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

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

333 

334 self._update_atoms() 

335 

336 def get_conserved_energy(self) -> float: 

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

338 

339 This method is mainly used for testing. 

340 """ 

341 conserved_energy = ( 

342 self.atoms.get_potential_energy(force_consistent=True) 

343 + self.atoms.get_kinetic_energy() 

344 + self._thermostat.get_thermostat_energy() 

345 + self._barostat.get_barostat_energy() 

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

347 + self._pressure_au * self._get_volume() 

348 ) 

349 return float(conserved_energy) 

350 

351 def _update_atoms(self) -> None: 

352 self.atoms.set_positions(self._q) 

353 self.atoms.set_momenta(self._p) 

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

355 # Never set scale_atoms=True 

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

357 

358 def _get_volume(self) -> float: 

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

360 

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

362 self._update_atoms() 

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

364 

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

366 self._update_atoms() 

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

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

369 return pressure 

370 

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

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

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

374 self._q = ( 

375 self._q * np.exp(x) 

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

377 ) 

378 

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

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

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

382 / self._barostat.W 

383 forces = self._get_forces() 

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

385 

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

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

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

389 

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

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

392 pressure = self._get_pressure() 

393 volume = self._get_volume() 

394 G = ( 

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

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

397 ) 

398 self._p_eps += delta * G 

399 

400 

401class IsotropicMTKBarostat: 

402 """MTK barostat for isotropic volume fluctuations. 

403 

404 See `IsotropicMTKNPT` for the references. 

405 """ 

406 def __init__( 

407 self, 

408 num_atoms_global: int, 

409 temperature_K: float, 

410 pdamp: float, 

411 pchain: int = 3, 

412 ploop: int = 1, 

413 ): 

414 self._num_atoms_global = num_atoms_global 

415 self._pdamp = pdamp 

416 self._pchain = pchain 

417 self._ploop = ploop 

418 

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

420 

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

422 

423 assert pchain >= 1 

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

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

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

427 

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

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

430 

431 @property 

432 def W(self) -> float: 

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

434 return self._W 

435 

436 def get_barostat_energy(self) -> float: 

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

438 energy = ( 

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

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

441 ) 

442 return float(energy) 

443 

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

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

446 for _ in range(self._ploop): 

447 for coeff in FOURTH_ORDER_COEFFS: 

448 p_eps = self._integrate_nhc_baro_loop( 

449 p_eps, coeff * delta / len(FOURTH_ORDER_COEFFS) 

450 ) 

451 return p_eps 

452 

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

454 delta2 = delta / 2 

455 delta4 = delta / 4 

456 

457 for j in range(self._pchain): 

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

459 self._integrate_xi(delta) 

460 p_eps = self._integrate_nhc_p_eps(p_eps, delta) 

461 for j in range(self._pchain): 

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

463 

464 return p_eps 

465 

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

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

468 if j < self._pchain - 1: 

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

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

471 ) 

472 

473 if j == 0: 

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

475 else: 

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

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

478 

479 if j < self._pchain - 1: 

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

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

482 ) 

483 

484 def _integrate_xi(self, delta: float) -> None: 

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

486 

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

488 p_eps_new = p_eps * float( 

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

490 ) 

491 return p_eps_new 

492 

493 

494class MTKNPT(MolecularDynamics): 

495 """Isothermal-isobaric molecular dynamics with volume-and-cell fluctuations 

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

497 

498 See also `NoseHooverChainNVT` for the references. 

499 

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

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

502 """ 

503 def __init__( 

504 self, 

505 atoms: Atoms, 

506 timestep: float, 

507 temperature_K: float, 

508 pressure_au: float, 

509 tdamp: float, 

510 pdamp: float, 

511 tchain: int = 3, 

512 pchain: int = 3, 

513 tloop: int = 1, 

514 ploop: int = 1, 

515 **kwargs, 

516 ): 

517 """ 

518 Parameters 

519 ---------- 

520 atoms: ase.Atoms 

521 The atoms object. 

522 timestep: float 

523 The time step in ASE time units. 

524 temperature_K: float 

525 The target temperature in K. 

526 pressure_au: float 

527 The external pressure in eV/Ang^3. 

528 tdamp: float 

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

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

531 pdamp: float 

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

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

534 tchain: int 

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

536 pchain: int 

537 The number of barostat variables in the MTK barostat. 

538 tloop: int 

539 The number of sub-steps in thermostat integration. 

540 ploop: int 

541 The number of sub-steps in barostat integration. 

542 **kwargs : dict, optional 

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

544 base class. 

545 """ 

546 super().__init__( 

547 atoms=atoms, 

548 timestep=timestep, 

549 **kwargs, 

550 ) 

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

552 

553 if len(atoms.constraints) > 0: 

554 raise NotImplementedError( 

555 "Current implementation does not support constraints" 

556 ) 

557 

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

559 self._thermostat = NoseHooverChainThermostat( 

560 num_atoms_global=self._num_atoms_global, 

561 masses=self.masses, 

562 temperature_K=temperature_K, 

563 tdamp=tdamp, 

564 tchain=tchain, 

565 tloop=tloop, 

566 ) 

567 self._barostat = MTKBarostat( 

568 num_atoms_global=self._num_atoms_global, 

569 temperature_K=temperature_K, 

570 pdamp=pdamp, 

571 pchain=pchain, 

572 ploop=ploop, 

573 ) 

574 

575 self._temperature_K = temperature_K 

576 self._pressure_au = pressure_au 

577 

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

579 

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

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

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

583 self._h = np.array(self.atoms.get_cell()) # cell 

584 self._p_g = np.zeros((3, 3)) # cell momenta 

585 

586 def step(self) -> None: 

587 dt2 = self.dt / 2 

588 

589 self._p_g = self._barostat.integrate_nhc_baro(self._p_g, dt2) 

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

591 self._integrate_p_cell(dt2) 

592 self._integrate_p(dt2) 

593 self._integrate_q(self.dt) 

594 self._integrate_q_cell(self.dt) 

595 self._integrate_p(dt2) 

596 self._integrate_p_cell(dt2) 

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

598 self._p_g = self._barostat.integrate_nhc_baro(self._p_g, dt2) 

599 

600 self._update_atoms() 

601 

602 def get_conserved_energy(self) -> float: 

603 conserved_energy = ( 

604 self.atoms.get_total_energy() 

605 + self._thermostat.get_thermostat_energy() 

606 + self._barostat.get_barostat_energy() 

607 + np.trace(self._p_g.T @ self._p_g) / (2 * self._barostat.W) 

608 + self._pressure_au * self._get_volume() 

609 ) 

610 return float(conserved_energy) 

611 

612 def _update_atoms(self) -> None: 

613 self.atoms.set_positions(self._q) 

614 self.atoms.set_momenta(self._p) 

615 self.atoms.set_cell(self._h, scale_atoms=False) 

616 

617 def _get_volume(self) -> float: 

618 return np.abs(np.linalg.det(self._h)) 

619 

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

621 self._update_atoms() 

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

623 

624 def _get_stress(self) -> np.ndarray: 

625 self._update_atoms() 

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

627 return -stress 

628 

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

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

631 # eigvals: (3-eigvec), U: (3-xyz, 3-eigvec) 

632 eigvals, U = np.linalg.eigh(self._p_g) 

633 x = self._q @ U # (num_atoms, 3-eigvec) 

634 y = self._p @ U # (num_atoms, 3-eigvec) 

635 sol = ( 

636 x * np.exp(eigvals * delta / self._barostat.W)[None, :] 

637 + delta * y / self.masses * exprel( 

638 eigvals * delta / self._barostat.W 

639 )[None, :] 

640 ) # (num_atoms, 3-eigvec) 

641 self._q = sol @ U.T 

642 

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

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

645 forces = self._get_forces() # (num_atoms, 3-xyz) 

646 

647 # eigvals: (3-eigvec), U: (3-xyz, 3-eigvec) 

648 eigvals, U = np.linalg.eigh(self._p_g) 

649 kappas = eigvals \ 

650 + np.trace(self._p_g) / (3 * self._num_atoms_global) # (3-eigvec) 

651 y = self._p @ U # (num_atoms, 3-eigvec) 

652 sol = ( 

653 y * np.exp(-kappas * delta / self._barostat.W)[None, :] 

654 + delta * (forces @ U) * exprel( 

655 -kappas * delta / self._barostat.W 

656 )[None, :] 

657 ) # (num_atoms, 3-eigvec) 

658 self._p = sol @ U.T 

659 

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

661 """Integrate exp(i * L_(g, 1) * delta)""" 

662 # U @ np.diag(eigvals) @ U.T = self._p_g 

663 # eigvals: (3-eigvec), U: (3-xyz, 3-eigvec) 

664 eigvals, U = np.linalg.eigh(self._p_g) 

665 n = self._h @ U # (3-axis, 3-eigvec) 

666 sol = n * np.exp( 

667 eigvals * delta / self._barostat.W 

668 )[None, :] # (3-axis, 3-eigvec) 

669 self._h = sol @ U.T 

670 

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

672 """Integrate exp(i * L_(g, 2) * delta)""" 

673 stress = self._get_stress() 

674 G = ( 

675 self._get_volume() * (stress - self._pressure_au * np.eye(3)) 

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

677 * np.eye(3) 

678 ) 

679 self._p_g += delta * G 

680 

681 

682class MTKBarostat: 

683 """MTK barostat for volume-and-cell fluctuations. 

684 

685 See `MTKNPT` for the references. 

686 """ 

687 def __init__( 

688 self, 

689 num_atoms_global: int, 

690 temperature_K: float, 

691 pdamp: float, 

692 pchain: int = 3, 

693 ploop: int = 1, 

694 ): 

695 self._num_atoms_global = num_atoms_global 

696 self._pdamp = pdamp 

697 self._pchain = pchain 

698 self._ploop = ploop 

699 

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

701 

702 self._W = (self._num_atoms_global + 1) * self._kT * self._pdamp**2 

703 

704 assert pchain >= 1 

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

706 cell_dof = 9 # TODO: 

707 self._R[0] = cell_dof * self._kT * self._pdamp**2 

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

709 

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

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

712 

713 @property 

714 def W(self) -> float: 

715 return self._W 

716 

717 def get_barostat_energy(self) -> float: 

718 energy = ( 

719 np.sum(self._p_xi**2 / self._R) / 2 

720 + 9 * self._kT * self._xi[0] 

721 + self._kT * np.sum(self._xi[1:]) 

722 ) 

723 return float(energy) 

724 

725 def integrate_nhc_baro(self, p_g: np.ndarray, delta: float) -> np.ndarray: 

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

727 for _ in range(self._ploop): 

728 for coeff in FOURTH_ORDER_COEFFS: 

729 p_g = self._integrate_nhc_baro_loop( 

730 p_g, coeff * delta / len(FOURTH_ORDER_COEFFS) 

731 ) 

732 return p_g 

733 

734 def _integrate_nhc_baro_loop( 

735 self, p_g: np.ndarray, delta: float 

736 ) -> np.ndarray: 

737 delta2 = delta / 2 

738 delta4 = delta / 4 

739 

740 for j in range(self._pchain): 

741 self._integrate_p_xi_j(p_g, self._pchain - j - 1, delta2, delta4) 

742 self._integrate_xi(delta) 

743 self._integrate_nhc_p_eps(p_g, delta) 

744 for j in range(self._pchain): 

745 self._integrate_p_xi_j(p_g, j, delta2, delta4) 

746 

747 return p_g 

748 

749 def _integrate_p_xi_j( 

750 self, p_g: np.ndarray, j: int, delta2: float, delta4: float 

751 ) -> None: 

752 if j < self._pchain - 1: 

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

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

755 ) 

756 

757 if j == 0: 

758 # TODO: do we need to substitute 9 with cell_dof? 

759 g_j = np.trace(p_g.T @ p_g) / self._W - 9 * self._kT 

760 else: 

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

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

763 

764 if j < self._pchain - 1: 

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

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