Coverage for calorine/nep/nep.py: 100%

1import contextlib 

2import os 

3import signal 

4import subprocess 

5import warnings 

6from collections import Counter 

7from pathlib import Path 

8 

9import numpy as np 

10from ase import Atoms 

11from ase.calculators.singlepoint import SinglePointCalculator 

12 

13import _nepy 

14from calorine.nep.model import Model, _get_nep_contents, _get_nep_version, _require_nep4 

15from calorine.nep.tensor_conventions import reduced6_to_full_3x3 

16 

17 

18def _get_atomic_properties( 

19 structure: Atoms, 

20) -> tuple[list[float], list[str], list[float]]: 

21 """Fetches cell, symbols and positions for a structure. Since NEP_CPU requires a cell, if the 

22 structure has no cell a default cubic cell with a side of 100 Å will be used. 

23 

24 Parameters 

25 ---------- 

26 structure 

27 Atoms object representing the structure. 

28 

29 Returns 

30 ------- 

31 list[float] 

32 Cell vectors 

33 list[str] 

34 Atomic species 

35 list[float] 

36 Atom positions 

37 list[float] 

38 Atom masses 

39 """ 

40 if structure.cell.rank == 0: 

41 warnings.warn('Using default unit cell (cubic with side 100 Å).') 

42 set_default_cell(structure) 

43 

44 c = structure.get_cell(complete=True).flatten() 

45 cell = [c[0], c[3], c[6], c[1], c[4], c[7], c[2], c[5], c[8]] 

46 

47 symbols = structure.get_chemical_symbols() 

48 positions = list( 

49 structure.get_positions().T.flatten() 

50 ) # [x1, ..., xN, y1, ... yN,...] 

51 masses = structure.get_masses() 

52 return cell, symbols, positions, masses 

53 

54 

55def _setup_nepy( 

56 model_filename: str, 

57 natoms: int, 

58 cell: list[float], 

59 symbols: list[str], 

60 positions: list[float], 

61 masses: list[float], 

62 debug: bool, 

63) -> _nepy.NEPY: 

64 """Sets up an instance of the NEPY pybind11 interface to NEP_CPU. 

65 

66 Parameters 

67 ---------- 

68 model_filename 

69 Path to model. 

70 natoms: 

71 Number of atoms in the structure. 

72 cell: 

73 Cell vectors. 

74 symbols: 

75 Atom species. 

76 positions: 

77 Atom positions. 

78 masses: 

79 Atom masses. 

80 debug: 

81 Flag to control if the output from NEP_CPU will be printed. 

82 

83 Returns 

84 ------- 

85 NEPY 

86 NEPY interface 

87 """ 

88 # Ensure that `model_filename` exists to avoid segfault in pybind11 code 

89 if not os.path.isfile(model_filename): 

90 raise ValueError(f'{Path(model_filename)} does not exist') 

91 

92 nep_version = _get_nep_version(model_filename) 

93 _require_nep4(nep_version) 

94 

95 # Disable output from C++ code by default 

96 if debug: 

97 nepy = _nepy.NEPY(model_filename, natoms, cell, symbols, positions, masses) 

98 else: 

99 with open(os.devnull, 'w') as f: 

100 with contextlib.redirect_stdout(f): 

101 with contextlib.redirect_stderr(f): 

102 nepy = _nepy.NEPY( 

103 model_filename, natoms, cell, symbols, positions, masses 

104 ) 

105 return nepy 

106 

107 

108def set_default_cell(structure: Atoms, box_length: float = 100): 

109 """Adds a cubic box to an Atoms object. Atoms object is edited in-place. 

110 

111 Parameters 

112 ---------- 

113 structure 

114 Structure to add box to 

115 box_length 

116 Cubic box side length in Å, by default 100 

117 """ 

118 structure.set_cell([[box_length, 0, 0], [0, box_length, 0], [0, 0, box_length]]) 

119 structure.center() 

120 

121 

122def get_descriptors( 

123 structure: Atoms, model_filename: str, debug: bool = False 

124) -> np.ndarray: 

125 """Calculates the NEP descriptors for a given structure. A NEP model defined by a nep.txt 

126 can additionally be provided to get the model specific descriptors. 

127 

128 Parameters 

129 ---------- 

130 structure 

131 Input structure 

132 model_filename 

133 Path to NEP model in ``nep.txt`` format. 

134 debug 

135 Flag to toggle debug mode. Prints GPUMD output. Defaults to ``False``. 

136 

137 Returns 

138 ------- 

139 Descriptors for the supplied structure, with shape (number_of_atoms, descriptor components) 

140 """ 

141 local_structure = structure.copy() 

142 natoms = len(local_structure) 

143 cell, symbols, positions, masses = _get_atomic_properties(local_structure) 

144 

145 nepy = _setup_nepy( 

146 model_filename, natoms, cell, symbols, positions, masses, debug 

147 ) 

148 all_descriptors = nepy.get_descriptors() 

149 descriptors_per_atom = np.array(all_descriptors).reshape(-1, natoms).T 

150 return descriptors_per_atom 

151 

152 

153def get_latent_space( 

154 structure: Atoms, model_filename: str | None = None, debug: bool = False 

155) -> np.ndarray: 

156 """Calculates the latent space representation of a structure, i.e, the activiations in 

157 the hidden layer. A NEP model defined by a `nep.txt` file needs to be provided. 

158 

159 Parameters 

160 ---------- 

161 structure 

162 Input structure 

163 model_filename 

164 Path to NEP model. Defaults to None. 

165 debug 

166 Flag to toggle debug mode. Prints GPUMD output. Defaults to False. 

167 

168 Returns 

169 ------- 

170 Activation with shape `(natoms, N_neurons)` 

171 """ 

172 if model_filename is None: 

173 raise ValueError('Model is undefined') 

174 local_structure = structure.copy() 

175 natoms = len(local_structure) 

176 cell, symbols, positions, masses = _get_atomic_properties(local_structure) 

177 

178 nepy = _setup_nepy( 

179 model_filename, natoms, cell, symbols, positions, masses, debug 

180 ) 

181 

182 latent = nepy.get_latent_space() 

183 latent = np.array(latent).reshape(-1, natoms).T 

184 return latent 

185 

186 

187def get_potential_forces_and_virials( 

188 structure: Atoms, model_filename: str | None = None, debug: bool = False 

189) -> tuple[np.ndarray, np.ndarray, np.ndarray]: 

190 """Calculates the per-atom potential, forces and virials for a given structure. 

191 A NEP model defined by a `nep.txt` file needs to be provided. 

192 

193 Parameters 

194 ---------- 

195 structure 

196 Input structure 

197 model_filename 

198 Path to NEP model. Defaults to None. 

199 debug 

200 Flag to toggle debug mode. Prints GPUMD output. Defaults to False. 

201 

202 Returns 

203 ------- 

204 potential with shape `(natoms,)` 

205 forces with shape `(natoms, 3)` 

206 virials with shape `(natoms, 9)` 

207 """ 

208 if model_filename is None: 

209 raise ValueError('Model is undefined') 

210 

211 model_type = _get_nep_contents(model_filename)[0]['model_type'] 

212 if model_type != 'potential': 

213 raise ValueError( 

214 'A NEP model trained for predicting energies and forces must be used.' 

215 ) 

216 

217 local_structure = structure.copy() 

218 natoms = len(local_structure) 

219 cell, symbols, positions, masses = _get_atomic_properties(local_structure) 

220 

221 nepy = _setup_nepy( 

222 model_filename, natoms, cell, symbols, positions, masses, debug 

223 ) 

224 

225 energies, forces, virials = nepy.get_potential_forces_and_virials() 

226 forces_per_atom = np.array(forces).reshape(-1, natoms).T 

227 # Raw, row-major full-3x3 per atom [xx,xy,xz,yx,yy,yz,zx,zy,zz]; 

228 # documented in src/nepy/nep.h and matching the accumulation pattern in 

229 # src/nepy/nep.cpp. 

230 virials_per_atom = np.array(virials).reshape(-1, natoms).T 

231 return np.array(energies), forces_per_atom, virials_per_atom 

232 

233 

234def get_polarizability( 

235 structure: Atoms, 

236 model_filename: str | None = None, 

237 debug: bool = False, 

238) -> np.ndarray: 

239 """Calculates the polarizability tensor for a given structure. A NEP model defined 

240 by a ``nep.txt`` file needs to be provided. The model must be trained to predict the 

241 polarizability. 

242 

243 Parameters 

244 ---------- 

245 structure 

246 Input structure 

247 model_filename 

248 Path to NEP model in ``nep.txt`` format. Defaults to ``None``. 

249 debug 

250 Flag to toggle debug mode. Prints GPUMD output. Defaults to ``False``. 

251 

252 Returns 

253 ------- 

254 polarizability with shape ``(3, 3)`` 

255 """ 

256 if model_filename is None: 

257 raise ValueError('Model is undefined') 

258 

259 model_type = _get_nep_contents(model_filename)[0]['model_type'] 

260 if model_type != 'polarizability': 

261 raise ValueError( 

262 'A NEP model trained for predicting polarizability must be used.' 

263 ) 

264 

265 local_structure = structure.copy() 

266 natoms = len(local_structure) 

267 cell, symbols, positions, masses = _get_atomic_properties(local_structure) 

268 

269 nepy = _setup_nepy( 

270 model_filename, natoms, cell, symbols, positions, masses, debug 

271 ) 

272 # Components are in NEP_REDUCED6_ORDER (xx,yy,zz,xy,yz,zx); see 

273 # `find_polarizability` in src/nepy/nep.cpp. 

274 pol = nepy.get_polarizability() 

275 polarizability = reduced6_to_full_3x3(pol) 

276 

277 return polarizability 

278 

279 

280def get_dipole( 

281 structure: Atoms, 

282 model_filename: str | None = None, 

283 debug: bool = False, 

284) -> np.ndarray: 

285 """Calculates the dipole for a given structure. A NEP model defined by a 

286 ``nep.txt`` file needs to be provided. 

287 

288 Parameters 

289 ---------- 

290 structure 

291 Input structure 

292 model_filename 

293 Path to NEP model in ``nep.txt`` format. Defaults to ``None``. 

294 debug 

295 Flag to toggle debug mode. Prints GPUMD output. Defaults to ``False``. 

296 

297 Returns 

298 ------- 

299 dipole with shape ``(3,)`` 

300 """ 

301 if model_filename is None: 

302 raise ValueError('Model is undefined') 

303 

304 model_type = _get_nep_contents(model_filename)[0]['model_type'] 

305 if model_type != 'dipole': 

306 raise ValueError('A NEP model trained for predicting dipoles must be used.') 

307 

308 local_structure = structure.copy() 

309 natoms = len(local_structure) 

310 cell, symbols, positions, masses = _get_atomic_properties(local_structure) 

311 

312 nepy = _setup_nepy( 

313 model_filename, natoms, cell, symbols, positions, masses, debug 

314 ) 

315 dipole = np.array(nepy.get_dipole()) 

316 

317 return dipole 

318 

319 

320def get_dipole_gradient( 

321 structure: Atoms, 

322 model_filename: str | None = None, 

323 backend: str = 'c++', 

324 method: str = 'central difference', 

325 displacement: float = 0.01, 

326 charge: float = 1.0, 

327 nep_command: str | None = None, 

328 debug: bool = False, 

329) -> np.ndarray: 

330 """Calculates the dipole gradient for a given structure using finite differences. 

331 A NEP model defined by a `nep.txt` file needs to be provided. 

332 

333 Parameters 

334 ---------- 

335 structure 

336 Input structure 

337 model_filename 

338 Path to NEP model in ``nep.txt`` format. Defaults to ``None``. 

339 backend 

340 Backend to use for computing dipole gradient with finite differences. 

341 One of ``'c++'`` (CPU), ``'python'`` (CPU) and ``'nep'`` (GPU). 

342 Defaults to ``'c++'``. 

343 method 

344 Method for computing gradient with finite differences. 

345 One of 'forward difference' and 'central difference'. 

346 Defaults to 'central difference' 

347 displacement 

348 Displacement in Å to use for finite differences. Defaults to ``0.01``. 

349 charge 

350 System charge in units of the elemental charge. 

351 Used for correcting the dipoles before computing the gradient. 

352 Defaults to ``1.0``. 

353 nep_command 

354 Command for running the NEP executable. 

355 Default: ``nep``, or the value of the ``CALORINE_NEP_COMMAND`` 

356 environment variable if set. 

357 debug 

358 Flag to toggle debug mode. Prints GPUMD output (if applicable). Defaults to ``False``. 

359 

360 

361 Returns 

362 ------- 

363 dipole gradient with shape ``(N, 3, 3)`` 

364 """ 

365 if model_filename is None: 

366 raise ValueError('Model is undefined') 

367 

368 model_type = _get_nep_contents(model_filename)[0]['model_type'] 

369 if model_type != 'dipole': 

370 raise ValueError('A NEP model trained for predicting dipoles must be used.') 

371 

372 local_structure = structure.copy() 

373 

374 if backend == 'c++': 

375 dipole_gradient = _dipole_gradient_cpp( 

376 local_structure, 

377 model_filename, 

378 displacement=displacement, 

379 method=method, 

380 charge=charge, 

381 debug=debug, 

382 ) 

383 elif backend == 'python': 

384 dipole_gradient = _dipole_gradient_python( 

385 local_structure, 

386 model_filename, 

387 displacement=displacement, 

388 charge=charge, 

389 method=method, 

390 ) 

391 elif backend == 'nep': 

392 dipole_gradient = _dipole_gradient_nep( 

393 local_structure, 

394 model_filename, 

395 displacement=displacement, 

396 method=method, 

397 charge=charge, 

398 nep_command=nep_command, 

399 ) 

400 else: 

401 raise ValueError(f'Invalid backend {backend}') 

402 return dipole_gradient 

403 

404 

405def _dipole_gradient_cpp( 

406 structure: Atoms, 

407 model_filename: str, 

408 method: str = 'central difference', 

409 displacement: float = 0.01, 

410 charge: float = 1.0, 

411 debug: bool = False, 

412) -> np.ndarray: 

413 """Calculates the dipole gradient with finite differences, using NEP_CPU. 

414 

415 Parameters 

416 ---------- 

417 structure 

418 Input structure 

419 model_filename 

420 Path to NEP model in ``nep.txt`` format. 

421 method 

422 Method for computing gradient with finite differences. 

423 One of ``'forward difference'`` and ``'central difference'``. 

424 Defaults to ``'central difference'`` 

425 displacement 

426 Displacement in Å to use for finite differences. Defaults to ``0.01``. 

427 charge 

428 System charge in units of the elemental charge. 

429 Used for correcting the dipoles before computing the gradient. 

430 Defaults to ``1.0``. 

431 

432 Returns 

433 ------- 

434 dipole gradient with shape ``(N, 3, 3)`` 

435 """ 

436 if displacement <= 0: 

437 raise ValueError('Displacement must be > 0 Å') 

438 

439 implemented_methods = { 

440 'forward difference': 0, 

441 'central difference': 1, 

442 'second order central difference': 2, 

443 } 

444 

445 if method not in implemented_methods.keys(): 

446 raise ValueError(f'Invalid method {method} for calculating gradient') 

447 

448 local_structure = structure.copy() 

449 natoms = len(local_structure) 

450 cell, symbols, positions, masses = _get_atomic_properties(local_structure) 

451 nepy = _setup_nepy( 

452 model_filename, natoms, cell, symbols, positions, masses, debug 

453 ) 

454 dipole_gradient = np.array( 

455 nepy.get_dipole_gradient(displacement, implemented_methods[method], charge) 

456 ).reshape(natoms, 3, 3) 

457 return dipole_gradient 

458 

459 

460def _dipole_gradient_python( 

461 structure: Atoms, 

462 model_filename: str, 

463 method: str = 'central difference', 

464 displacement: float = 0.01, 

465 charge: float = 1.0, 

466) -> np.ndarray: 

467 """Calculates the dipole gradient with finite differences, using the Python and get_dipole(). 

468 

469 Parameters 

470 ---------- 

471 structure 

472 Input structure 

473 model_filename 

474 Path to NEP model in ``nep.txt`` format. 

475 method 

476 Method for computing gradient with finite differences. 

477 One of ``'forward difference'`` and ``'central difference'``. 

478 Defaults to ``'central difference'`` 

479 displacement 

480 Displacement in Å to use for finite differences. Defaults to ``0.01``. 

481 charge 

482 System charge in units of the elemental charge. 

483 Used for correcting the dipoles before computing the gradient. 

484 Defaults to ``1.0``. 

485 

486 Returns 

487 ------- 

488 dipole gradient with shape ``(N, 3, 3)`` 

489 """ 

490 if displacement <= 0: 

491 raise ValueError('Displacement must be > 0 Å') 

492 

493 N = len(structure) 

494 if method == 'forward difference': 

495 # Correct all dipole by the permanent dipole, charge * center of mass 

496 d = ( 

497 get_dipole(structure, model_filename) 

498 + charge * structure.get_center_of_mass() 

499 ) 

500 d_forward = np.zeros((N, 3, 3)) 

501 for atom in range(N): 

502 for cartesian in range(3): 

503 copy = structure.copy() 

504 positions = copy.get_positions() 

505 positions[atom, cartesian] += displacement 

506 copy.set_positions(positions) 

507 d_forward[atom, cartesian, :] = ( 

508 get_dipole(copy, model_filename) 

509 + charge * copy.get_center_of_mass() 

510 ) 

511 gradient = (d_forward - d[None, None, :]) / displacement 

512 

513 elif method == 'central difference': 

514 d_forward = np.zeros((N, 3, 3)) 

515 d_backward = np.zeros((N, 3, 3)) 

516 for atom in range(N): 

517 for cartesian in range(3): 

518 # Forward displacements 

519 copy_forward = structure.copy() 

520 positions_forward = copy_forward.get_positions() 

521 positions_forward[atom, cartesian] += displacement 

522 copy_forward.set_positions(positions_forward) 

523 d_forward[atom, cartesian, :] = ( 

524 get_dipole(copy_forward, model_filename) 

525 + charge * copy_forward.get_center_of_mass() 

526 ) 

527 # Backwards displacement 

528 copy_backward = structure.copy() 

529 positions_backward = copy_backward.get_positions() 

530 positions_backward[atom, cartesian] -= displacement 

531 copy_backward.set_positions(positions_backward) 

532 d_backward[atom, cartesian, :] = ( 

533 get_dipole(copy_backward, model_filename) 

534 + charge * copy_backward.get_center_of_mass() 

535 ) 

536 gradient = (d_forward - d_backward) / (2 * displacement) 

537 elif method == 'second order central difference': 

538 # Coefficients from 

539 # https://en.wikipedia.org/wiki/Finite_difference_coefficient#Central_finite_difference 

540 d_forward_one_h = np.zeros((N, 3, 3)) 

541 d_forward_two_h = np.zeros((N, 3, 3)) 

542 d_backward_one_h = np.zeros((N, 3, 3)) 

543 d_backward_two_h = np.zeros((N, 3, 3)) 

544 for atom in range(N): 

545 for cartesian in range(3): 

546 copy = structure.copy() 

547 positions = copy.get_positions() 

548 # Forward displacements 

549 positions[atom, cartesian] += displacement # + h 

550 copy.set_positions(positions) 

551 d_forward_one_h[atom, cartesian, :] = ( 

552 get_dipole(copy, model_filename) 

553 + charge * copy.get_center_of_mass() 

554 ) 

555 positions[atom, cartesian] += displacement # + 2h total 

556 copy.set_positions(positions) 

557 d_forward_two_h[atom, cartesian, :] = ( 

558 get_dipole(copy, model_filename) 

559 + charge * copy.get_center_of_mass() 

560 ) 

561 # Backwards displacement 

562 positions[atom, cartesian] -= 3 * displacement # 2h - 3h = -h 

563 copy.set_positions(positions) 

564 d_backward_one_h[atom, cartesian, :] = ( 

565 get_dipole(copy, model_filename) 

566 + charge * copy.get_center_of_mass() 

567 ) 

568 positions[atom, cartesian] -= displacement # - 2h total 

569 copy.set_positions(positions) 

570 d_backward_two_h[atom, cartesian, :] = ( 

571 get_dipole(copy, model_filename) 

572 + charge * copy.get_center_of_mass() 

573 ) 

574 c0 = -1.0 / 12.0 

575 c1 = 2.0 / 3.0 

576 gradient = ( 

577 c0 * d_forward_two_h 

578 + c1 * d_forward_one_h 

579 - c1 * d_backward_one_h 

580 - c0 * d_backward_two_h 

581 ) / displacement 

582 else: 

583 raise ValueError(f'Invalid method {method} for calculating gradient') 

584 return gradient 

585 

586 

587def _dipole_gradient_nep( 

588 structure: Atoms, 

589 model_filename: str, 

590 method: str = 'central difference', 

591 displacement: float = 0.01, 

592 charge: float = 1.0, 

593 nep_command: str | None = None, 

594) -> np.ndarray: 

595 """Calculates the dipole gradient with finite differences, using the NEP executable. 

596 

597 Parameters 

598 ---------- 

599 structure 

600 Input structure 

601 model_filename 

602 Path to NEP model in ``nep.txt`` format. 

603 method 

604 Method for computing gradient with finite differences. 

605 One of ``'forward difference'`` and ``'central difference'``. 

606 Defaults to ``'central difference'`` 

607 displacement 

608 Displacement in Å to use for finite differences. Defaults to 0.01 Å. 

609 Note that results are possibly unreliable for displacemen < 0.01, 

610 which might be due to rounding errors. 

611 charge 

612 System charge in units of the elemental charge. 

613 Used for correcting the dipoles before computing the gradient. 

614 Defaults to 1.0. 

615 nep_command 

616 Command for running the NEP executable. 

617 Default: ``nep``, or the value of the ``CALORINE_NEP_COMMAND`` 

618 environment variable if set. 

619 

620 

621 Returns 

622 ------- 

623 dipole gradient with shape ``(N, 3, 3)`` 

624 """ 

625 if displacement <= 0: 

626 raise ValueError('Displacement must be > 0 Å') 

627 

628 if displacement < 1e-2: 

629 warnings.warn( 

630 'Dipole gradients with nep are unstable for displacements < 0.01 Å.' 

631 ) 

632 

633 N = len(structure) 

634 if method == 'forward difference': 

635 structures = [structure] # will hold 3N+1 structures 

636 # Correct for the constant dipole, by adding charge * center of mass 

637 corrections = np.zeros((3 * N + 1, 3)) 

638 corrections[0] = charge * structure.get_center_of_mass() 

639 for atom in range(N): 

640 for cartesian in range(3): 

641 copy = structure.copy() 

642 positions = copy.get_positions() 

643 positions[atom, cartesian] += displacement 

644 copy.set_positions(positions) 

645 structures.append(copy) 

646 corrections[1 + atom * 3 + cartesian] = ( 

647 charge * copy.get_center_of_mass() 

648 ) 

649 

650 dipoles = _predict_dipole_batch( 

651 structures, model_filename, nep_command) # shape (3N+1, 3) 

652 dipoles += corrections 

653 

654 d = dipoles[0, :] 

655 d_forward = dipoles[1:].reshape(N, 3, 3) 

656 gradient = (d_forward - d[None, None, :]) / displacement 

657 

658 elif method == 'central difference': 

659 structures_forward = [] # will hold 3N structures 

660 structures_backward = [] # will hold 3N structures 

661 

662 # Correct for the constant dipole, by adding charge * center of mass 

663 corrections_forward = np.zeros((3 * N, 3)) 

664 corrections_backward = np.zeros((3 * N, 3)) 

665 

666 for atom in range(N): 

667 for cartesian in range(3): 

668 # Forward displacements 

669 copy_forward = structure.copy() 

670 positions_forward = copy_forward.get_positions() 

671 positions_forward[atom, cartesian] += displacement 

672 copy_forward.set_positions(positions_forward) 

673 structures_forward.append(copy_forward) 

674 corrections_forward[atom * 3 + cartesian] = ( 

675 charge * copy_forward.get_center_of_mass() 

676 ) 

677 

678 # Backwards displacement 

679 copy_backward = structure.copy() 

680 positions_backward = copy_backward.get_positions() 

681 positions_backward[atom, cartesian] -= displacement 

682 copy_backward.set_positions(positions_backward) 

683 structures_backward.append(copy_backward) 

684 corrections_backward[atom * 3 + cartesian] = ( 

685 charge * copy_backward.get_center_of_mass() 

686 ) 

687 

688 structures = structures_forward + structures_backward 

689 dipoles = _predict_dipole_batch( 

690 structures, model_filename, nep_command) # shape (6N, 3) 

691 d_forward = dipoles[: 3 * N, :] 

692 d_backward = dipoles[3 * N :, :] 

693 

694 d_forward += corrections_forward 

695 d_backward += corrections_backward 

696 

697 d_forward = d_forward.reshape(N, 3, 3) 

698 d_backward = d_backward.reshape(N, 3, 3) 

699 

700 gradient = (d_forward - d_backward) / (2 * displacement) 

701 else: 

702 raise ValueError(f'Invalid method {method} for calculating gradient') 

703 return gradient 

704 

705 

706def _run_nep_prediction(command: str, directory: str) -> None: 

707 """Runs the NEP executable in prediction mode in a directory that holds its input files. 

708 

709 Parameters 

710 ---------- 

711 command 

712 Command for running the NEP executable. 

713 directory 

714 Directory holding ``nep.in``, ``nep.txt`` and ``train.xyz``. 

715 

716 Raises 

717 ------ 

718 RuntimeError 

719 If the executable cannot be started or ends with an error, carrying what it printed 

720 and its exit status or the signal that ended it. 

721 """ 

722 try: 

723 completed = subprocess.run([command], cwd=directory, capture_output=True) 

724 except OSError as e: 

725 raise RuntimeError(f'Failed to run `{command}` in prediction mode: {e}') from e 

726 if completed.returncode != 0: 

727 if completed.returncode < 0: 

728 try: 

729 name = signal.Signals(-completed.returncode).name 

730 except ValueError: 

731 # `signal.Signals` names only some of the real-time signals 

732 name = f'{-completed.returncode}' 

733 ending = f'killed by signal {name}' 

734 else: 

735 ending = f'exit status {completed.returncode}' 

736 # `nep` prints the keywords on which nep.in and nep.txt disagree on standard output 

737 # and only the closing error message on standard error 

738 output = ''.join(stream.decode(errors='replace') 

739 for stream in (completed.stdout, completed.stderr)) 

740 raise RuntimeError( 

741 f'Failed to run `{command}` in prediction mode ({ending}):\n{output}') 

742 

743 

744def _predict_dipole_batch( 

745 structures: list[Atoms], model_filename: str, nep_command: str | None = None 

746) -> np.ndarray: 

747 """Predicts the dipoles of a set of structures using the NEP executable. 

748 

749 Parameters 

750 ---------- 

751 structures 

752 Input structures 

753 model_filename 

754 Path to NEP model in ``nep.txt`` format. 

755 nep_command 

756 Command for running the NEP executable. 

757 Default: ``nep``, or the value of the ``CALORINE_NEP_COMMAND`` 

758 environment variable if set. 

759 

760 Returns 

761 ------- 

762 Predicted dipoles, with shape ``(len(structures), 3)``. 

763 """ 

764 # calorine.tools.prediction imports this module, so the import is made at call time 

765 from calorine.tools.prediction import batch_predict_properties 

766 

767 predicted = batch_predict_properties(structures, model_filename, command=nep_command) 

768 return np.array([structure.get_dipole_moment() for structure in predicted]) 

769 

770 

771def _check_components_polarizability_gradient(component: str | list[str]) -> list[int]: 

772 """ 

773 Verifies that the selected components are ok. 

774 """ 

775 allowed_components = { 

776 'x': 0, 

777 'y': 1, 

778 'z': 2, 

779 } 

780 

781 # Check if chosen components are ok 

782 if component == 'full': 

783 components_to_compute = [0, 1, 2] 

784 else: 

785 components_list = [component] if isinstance(component, str) else component 

786 components_to_compute = [] 

787 for c in components_list: 

788 if c in allowed_components.keys(): 

789 components_to_compute.append(allowed_components[c]) 

790 elif c == 'full': 

791 raise ValueError('Write ``component="full"`` to get all components.') 

792 else: 

793 raise ValueError(f'Invalid component {c}') 

794 assert len(components_list) == len(components_to_compute), \ 

795 'Number of components to compute does not match.' 

796 return components_to_compute 

797 

798 

799def get_polarizability_gradient( 

800 structure: Atoms, 

801 model_filename: str | None = None, 

802 displacement: float = 0.01, 

803 component: str | list[str] = 'full', 

804 debug: bool = False, 

805) -> np.ndarray: 

806 """Calculates the dipole gradient for a given structure using finite differences. 

807 A NEP model defined by a ``nep.txt`` file needs to be provided. 

808 This function computes the derivatives using the second-order central difference 

809 method with a C++ backend. 

810 

811 Parameters 

812 ---------- 

813 structure 

814 Input structure. 

815 model_filename 

816 Path to NEP model in ``nep.txt`` format. Defaults to ``None``. 

817 displacement 

818 Displacement in Å to use for finite differences. Defaults to ``0.01``. 

819 component 

820 Component or components of the polarizability tensor that the gradient 

821 should be computed for. 

822 The following components are available: ``x`, ``y``, ``z``, ``full``. 

823 Option ``full`` computes the derivative whilst moving the atoms in each Cartesian 

824 direction, which yields a tensor of shape ``(N, 3, 6)``. 

825 Multiple components may be specified. 

826 Defaults to ``full``. 

827 debug 

828 Flag to toggle debug mode. Prints GPUMD output (if applicable). Defaults to ``False``. 

829 

830 

831 Returns 

832 ------- 

833 polarizability gradient with shape ``(N, C, 6)`` where ``C`` 

834 is the number of components chosen. 

835 """ 

836 if model_filename is None: 

837 raise ValueError('Model is undefined') 

838 

839 model_type = _get_nep_contents(model_filename)[0]['model_type'] 

840 if model_type != 'polarizability': 

841 raise ValueError('A NEP model trained for predicting polarizability must be used.') 

842 

843 components_to_compute = _check_components_polarizability_gradient(component) 

844 

845 local_structure = structure.copy() 

846 polarizability_gradient = _polarizability_gradient_cpp( 

847 local_structure, 

848 model_filename, 

849 displacement=displacement, 

850 components=components_to_compute, 

851 debug=debug, 

852 ) 

853 return polarizability_gradient 

854 

855 

856def _polarizability_gradient_to_3x3(pg): 

857 """ 

858 Converts a polarizability gradient tensor with 

859 shape (Natoms, 3, 6) to (Natoms, 3, 3, 3). 

860 The 6 items in the polarizability gradient are in NEP_REDUCED6_ORDER 

861 (xx, yy, zz, xy, yz, zx), the same convention as `get_polarizability`: 

862 src/nepy/nep.cpp computes the gradient via finite differences of 

863 repeated `find_polarizability` calls, so it inherits that function's 

864 component order directly. 

865 """ 

866 return reduced6_to_full_3x3(pg) 

867 

868 

869def _polarizability_gradient_cpp( 

870 structure: Atoms, 

871 model_filename: str, 

872 displacement: float, 

873 components: list[int], 

874 debug: bool = False, 

875) -> np.ndarray: 

876 """Calculates the polarizability gradient with finite differences, using NEP_CPU. 

877 

878 Parameters 

879 ---------- 

880 structure 

881 Input structure. 

882 model_filename 

883 Path to NEP model in ``nep.txt`` format. 

884 displacement 

885 Displacement in Å to use for finite differences. Defaults to ``0.01``. 

886 components 

887 List of components to compute. Integer values from 0 to 2 that corresponds 

888 to indices in the ordered list [x, y, z]. Index 1 corresponds to the 

889 derivative with regards to only the y positions of the atoms, and so forth. 

890 

891 Returns 

892 ------- 

893 dipole gradient with shape ``(N, len(components), 3, 3)`` 

894 """ 

895 if displacement <= 0: 

896 raise ValueError('Displacement must be > 0 Å') 

897 # TODO possibly use components later to only move atoms in one cartesian direction 

898 local_structure = structure.copy() 

899 natoms = len(local_structure) 

900 cell, symbols, positions, masses = _get_atomic_properties(local_structure) 

901 nepy = _setup_nepy( 

902 model_filename, natoms, cell, symbols, positions, masses, debug 

903 ) 

904 pg = np.array( 

905 nepy.get_polarizability_gradient(displacement, components) 

906 ).reshape(natoms, 3, 6) 

907 # Convert to 3x3 

908 polarizability_gradient = _polarizability_gradient_to_3x3(pg) 

909 return polarizability_gradient[:, components, :, :] # Only return the relevant components 

910 

911 

912def determine_energy_offsets(