Coverage for ase / md / nose_hoover_chain.py: 98.88%

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 https://arxiv.org/abs/2603.24061. 

39 

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

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

42 """ 

43 

44 def __init__( 

45 self, 

46 atoms: Atoms, 

47 timestep: float, 

48 temperature_K: float, 

49 tdamp: float, 

50 tchain: int = 3, 

51 tloop: int = 1, 

52 **kwargs, 

53 ): 

54 """ 

55 Parameters 

56 ---------- 

57 atoms: ase.Atoms 

58 The atoms object. 

59 timestep: float 

60 The time step in ASE time units. 

61 temperature_K: float 

62 The target temperature in K. 

63 tdamp: float 

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

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

66 tchain: int 

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

68 tloop: int 

69 The number of sub-steps in thermostat integration. 

70 **kwargs : dict, optional 

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

72 base class. 

73 """ 

74 super().__init__( 

75 atoms=atoms, 

76 timestep=timestep, 

77 **kwargs, 

78 ) 

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

80 

81 num_atoms = self.atoms.get_global_number_of_atoms() 

82 self._thermostat = NoseHooverChainThermostat( 

83 num_atoms_global=num_atoms, 

84 masses=self.masses, 

85 temperature_K=temperature_K, 

86 tdamp=tdamp, 

87 tchain=tchain, 

88 tloop=tloop, 

89 ) 

90 

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

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

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

94 

95 def step(self) -> None: 

96 dt2 = self.dt / 2 

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

98 self._integrate_p(dt2) 

99 self._integrate_q(self.dt) 

100 self._integrate_p(dt2) 

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

102 

103 self._update_atoms() 

104 

105 def get_conserved_energy(self) -> float: 

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

107 

108 This method is mainly used for testing. 

109 """ 

110 conserved_energy = ( 

111 self.atoms.get_potential_energy(force_consistent=True) 

112 + self.atoms.get_kinetic_energy() 

113 + self._thermostat.get_thermostat_energy() 

114 ) 

115 return float(conserved_energy) 

116 

117 def _update_atoms(self) -> None: 

118 self.atoms.set_positions(self._q) 

119 self.atoms.set_momenta(self._p) 

120 

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

122 self._update_atoms() 

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

124 

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

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

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

128 

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

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

131 forces = self._get_forces() 

132 self._p += delta * forces 

133 

134 

135class NoseHooverChainThermostat: 

136 """Nose-Hoover chain style thermostats. 

137 

138 See `NoseHooverChainNVT` for the references. 

139 """ 

140 def __init__( 

141 self, 

142 num_atoms_global: int, 

143 masses: np.ndarray, 

144 temperature_K: float, 

145 tdamp: float, 

146 tchain: int = 3, 

147 tloop: int = 1, 

148 ): 

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

150 self._num_atoms_global = num_atoms_global 

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

152 self._tdamp = tdamp 

153 self._tchain = tchain 

154 self._tloop = tloop 

155 

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

157 

158 assert tchain >= 1 

159 self._Q = np.zeros(tchain) 

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

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

162 

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

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

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

166 

167 def get_thermostat_energy(self) -> float: 

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

169 energy = ( 

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

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

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

173 ) 

174 return float(energy) 

175 

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

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

178 for _ in range(self._tloop): 

179 for coeff in FOURTH_ORDER_COEFFS: 

180 p = self._integrate_nhc_loop( 

181 p, coeff * delta / self._tloop 

182 ) 

183 

184 return p 

185 

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

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

188 if j < self._tchain - 1: 

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

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

191 ) 

192 

193 if j == 0: 

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

195 - 3 * self._num_atoms_global * self._kT 

196 else: 

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

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

199 

200 if j < self._tchain - 1: 

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

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

203 ) 

204 

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

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

207 

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

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

210 

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

212 delta2 = delta / 2 

213 delta4 = delta / 4 

214 

215 for j in range(self._tchain): 

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

217 self._integrate_eta(delta) 

218 self._integrate_nhc_p(p, delta) 

219 for j in range(self._tchain): 

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

221 

222 return p 

223 

224 

225class IsotropicMTKNPT(MolecularDynamics): 

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

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

228 

229 See also :class:`NoseHooverChainNVT` for the references. 

230 The factorization of the Liouville operator is the same as Reference [1]. 

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 / self._ploop 

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 :class:`NoseHooverChainNVT` for the references. 

499 The factorization of the Liouville operator is the same as Reference [1]. 

500 

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

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

503 """ 

504 def __init__( 

505 self, 

506 atoms: Atoms, 

507 timestep: float, 

508 temperature_K: float, 

509 pressure_au: float, 

510 tdamp: float, 

511 pdamp: float, 

512 tchain: int = 3, 

513 pchain: int = 3, 

514 tloop: int = 1, 

515 ploop: int = 1, 

516 **kwargs, 

517 ): 

518 """ 

519 Parameters 

520 ---------- 

521 atoms: ase.Atoms 

522 The atoms object. 

523 timestep: float 

524 The time step in ASE time units. 

525 temperature_K: float 

526 The target temperature in K. 

527 pressure_au: float 

528 The external pressure in eV/Ang^3. 

529 tdamp: float 

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

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

532 pdamp: float 

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

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

535 tchain: int 

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

537 pchain: int 

538 The number of barostat variables in the MTK barostat. 

539 tloop: int 

540 The number of sub-steps in thermostat integration. 

541 ploop: int 

542 The number of sub-steps in barostat integration. 

543 **kwargs : dict, optional 

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

545 base class. 

546 """ 

547 super().__init__( 

548 atoms=atoms, 

549 timestep=timestep, 

550 **kwargs, 

551 ) 

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

553 

554 if len(atoms.constraints) > 0: 

555 raise NotImplementedError( 

556 "Current implementation does not support constraints" 

557 ) 

558 

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

560 self._thermostat = NoseHooverChainThermostat( 

561 num_atoms_global=self._num_atoms_global, 

562 masses=self.masses, 

563 temperature_K=temperature_K, 

564 tdamp=tdamp, 

565 tchain=tchain, 

566 tloop=tloop, 

567 ) 

568 

569 self._barostat = MTKBarostat( 

570 num_atoms_global=self._num_atoms_global, 

571 temperature_K=temperature_K, 

572 pdamp=pdamp, 

573 pchain=pchain, 

574 ploop=ploop, 

575 mask=self.mask, 

576 ) 

577 

578 self._temperature_K = temperature_K 

579 self._pressure_au = pressure_au 

580 

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

582 

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

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

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

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

587 

588 self._init_cell_momenta() 

589 

590 @property 

591 def mask(self) -> tuple[bool, bool, bool] | None: 

592 return None 

593 

594 @mask.setter 

595 def mask(self, mask: tuple[bool, bool, bool]) -> None: 

596 raise AttributeError() 

597 

598 def _init_cell_momenta(self) -> None: 

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

600 

601 def step(self) -> None: 

602 dt2 = self.dt / 2 

603 

604 self._integrate_p_cell_by_barostat(dt2) 

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

606 self._integrate_p_cell(dt2) 

607 self._integrate_p(dt2) 

608 self._integrate_q(self.dt) 

609 self._integrate_q_cell(self.dt) 

610 self._integrate_p(dt2) 

611 self._integrate_p_cell(dt2) 

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

613 self._integrate_p_cell_by_barostat(dt2) 

614 

615 self._update_atoms() 

616 

617 def get_conserved_energy(self) -> float: 

618 conserved_energy = ( 

619 self.atoms.get_total_energy() 

620 + self._thermostat.get_thermostat_energy() 

621 + self._barostat.get_barostat_energy() 

622 + self._get_cell_kinetic_energy() 

623 + self._pressure_au * self._get_volume() 

624 ) 

625 return float(conserved_energy) 

626 

627 def _update_atoms(self) -> None: 

628 self.atoms.set_positions(self._q) 

629 self.atoms.set_momenta(self._p) 

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

631 

632 def _get_volume(self) -> float: 

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

634 

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

636 self._update_atoms() 

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

638 

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

640 self._update_atoms() 

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

642 return -stress 

643 

644 def _get_cell_kinetic_energy(self) -> float: 

645 return float(np.sum(self._p_g**2) / (2 * self._barostat.W)) 

646 

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

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

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

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

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

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

653 sol = ( 

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

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

656 eigvals * delta / self._barostat.W 

657 )[None, :] 

658 ) # (num_atoms, 3-eigvec) 

659 self._q = sol @ U.T 

660 

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

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

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

664 

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

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

667 kappas = eigvals \ 

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

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

670 sol = ( 

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

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

673 -kappas * delta / self._barostat.W 

674 )[None, :] 

675 ) # (num_atoms, 3-eigvec) 

676 self._p = sol @ U.T 

677 

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

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

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

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

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

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

684 sol = n * np.exp( 

685 eigvals * delta / self._barostat.W 

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

687 self._h = sol @ U.T 

688 

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

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

691 stress = self._get_stress() 

692 volume = self._get_volume() 

693 particle_dof = 3 * self._num_atoms_global 

694 kinetic_term = np.sum(self._p**2 / self.masses) / particle_dof 

695 pv_tensor = volume * (stress - self._pressure_au * np.eye(3)) 

696 G = pv_tensor + kinetic_term * np.eye(3) 

697 self._p_g += delta * G 

698 

699 def _integrate_p_cell_by_barostat(self, delta: float) -> None: 

700 self._p_g = self._barostat.integrate_nhc_baro(self._p_g, delta) 

701 

702 

703class MaskedMTKNPT(MTKNPT): 

704 """Isothermal-isobaric molecular dynamics with cell fluctuations along 

705 specified crystallographic axes by Martyna-Tobias-Klein (MTK) method [1]. 

706 

707 The core equations are detailed in the implementation note available at 

708 https://arxiv.org/abs/2603.24061. 

709 

710 See also :class:`NoseHooverChainNVT` for the references. 

711 

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

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

714 """ 

715 def __init__( 

716 self, 

717 *args, 

718 mask: tuple[bool, bool, bool], 

719 **kwargs, 

720 ): 

721 """ 

722 Parameters 

723 ---------- 

724 mask: tuple[bool, bool, bool] 

725 Boolean mask for (a, b, c) axis fluctuations. ``True`` enables 

726 fluctuations along an axis and ``False`` disables them. For 

727 example, ``mask=(False, False, True)`` allows only c-axis 

728 fluctuations. 

729 *args 

730 Positional arguments passed to :class:`MTKNPT`. 

731 **kwargs 

732 Keyword arguments passed to :class:`MTKNPT`. 

733 """ 

734 self.mask = mask 

735 super().__init__(*args, **kwargs) 

736 

737 def _init_cell_momenta(self) -> None: 

738 # Additional variables for masked MTK 

739 self._h0_basis = self._h / np.linalg.norm(self._h, axis=1)[:, None] 

740 self._p_c = np.zeros(3) # Cell momenta for axis 

741 

742 @property 

743 def mask(self) -> tuple[bool, bool, bool]: 

744 return self._mask 

745 

746 @mask.setter 

747 def mask(self, mask: tuple[bool, bool, bool]) -> None: 

748 self._mask = mask 

749 

750 @property 

751 def _p_g(self) -> np.ndarray: 

752 # equivalent to sum(pc * np.outer(ec, ec)) 

753 return (self._p_c * self._h0_basis.T) @ self._h0_basis 

754 

755 @_p_g.setter 

756 def _p_g(self, value: np.ndarray) -> None: 

757 raise AttributeError() 

758 

759 def _get_cell_kinetic_energy(self) -> float: 

760 return float(np.sum(self._p_c**2) / (2 * self._barostat.W)) 

761 

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

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

764 stress = self._get_stress() 

765 volume = self._get_volume() 

766 particle_dof = 3 * self._num_atoms_global 

767 kinetic_term = np.sum(self._p**2 / self.masses) / particle_dof 

768 pv_tensor = volume * (stress - self._pressure_au * np.eye(3)) 

769 gc = np.diag( 

770 self._h0_basis @ pv_tensor @ self._h0_basis.T 

771 ) + kinetic_term 

772 self._p_c[self.mask] += delta * gc[self.mask] 

773 

774 def _integrate_p_cell_by_barostat(self, delta: float) -> None: 

775 self._p_c = self._barostat.integrate_nhc_baro(self._p_c, delta) 

776 

777 

778class MTKBarostat: 

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

780