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

1import contextlib 

2import os 

3import warnings 

4from pathlib import Path 

5from typing import List, Optional, Tuple, Union 

6 

7import numpy as np 

8from ase import Atoms 

9 

10import _nepy 

11from calorine.nep.model import _get_nep_contents 

12 

13 

14def _get_atomic_properties( 

15 structure: Atoms, 

16) -> Tuple[List[float], List[str], List[float]]: 

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

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

19 

20 Parameters 

21 ---------- 

22 structure 

23 Atoms object representing the structure. 

24 

25 Returns 

26 ------- 

27 List[float] 

28 Cell vectors 

29 List[str] 

30 Atomic species 

31 List[float] 

32 Atom positions 

33 List[float] 

34 Atom masses 

35 """ 

36 if structure.cell.rank == 0: 

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

38 set_default_cell(structure) 

39 

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

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

42 

43 symbols = structure.get_chemical_symbols() 

44 positions = list( 

45 structure.get_positions().T.flatten() 

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

47 masses = structure.get_masses() 

48 return cell, symbols, positions, masses 

49 

50 

51def _setup_nepy( 

52 model_filename: str, 

53 natoms: int, 

54 cell: List[float], 

55 symbols: List[str], 

56 positions: List[float], 

57 masses: List[float], 

58 debug: bool, 

59) -> _nepy.NEPY: 

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

61 

62 Parameters 

63 ---------- 

64 model_filename 

65 Path to model. 

66 natoms: 

67 Number of atoms in the structure. 

68 cell: 

69 Cell vectors. 

70 symbols: 

71 Atom species. 

72 positions: 

73 Atom positions. 

74 masses: 

75 Atom masses. 

76 debug: 

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

78 

79 Returns 

80 ------- 

81 NEPY 

82 NEPY interface 

83 """ 

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

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

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

87 

88 # Disable output from C++ code by default 

89 if debug: 

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

91 else: 

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

93 with contextlib.redirect_stdout(f): 

94 with contextlib.redirect_stderr(f): 

95 nepy = _nepy.NEPY( 

96 model_filename, natoms, cell, symbols, positions, masses 

97 ) 

98 return nepy 

99 

100 

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

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

103 

104 Parameters 

105 ---------- 

106 structure 

107 Structure to add box to 

108 box_length 

109 Cubic box side length in Å, by default 100 

110 """ 

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

112 structure.center() 

113 

114 

115def get_descriptors( 

116 structure: Atoms, model_filename: Optional[str] = None, debug: bool = False 

117) -> np.ndarray: 

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

119 can additionally be provided to get the NEP3 model specific descriptors. 

120 

121 Parameters 

122 ---------- 

123 structure 

124 Input structure 

125 model_filename 

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

127 debug 

128 Flag to toggle debug mode. Makes the generated dummy NEP2 model available 

129 in a local tmp directory, as well as prints GPUMD output. Defaults to ``False``. 

130 

131 Returns 

132 ------- 

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

134 """ 

135 local_structure = structure.copy() 

136 natoms = len(local_structure) 

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

138 

139 nepy = _setup_nepy( 

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

141 ) 

142 all_descriptors = nepy.get_descriptors() 

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

144 return descriptors_per_atom 

145 

146 

147def get_latent_space( 

148 structure: Atoms, model_filename: Union[str, None] = None, debug: bool = False 

149) -> np.ndarray: 

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

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

152 

153 Parameters 

154 ---------- 

155 structure 

156 Input structure 

157 model_filename 

158 Path to NEP model. Defaults to None. 

159 debug 

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

161 

162 Returns 

163 ------- 

164 Activation with shape `(natoms, N_neurons)` 

165 """ 

166 if model_filename is None: 

167 raise ValueError('Model is undefined') 

168 local_structure = structure.copy() 

169 natoms = len(local_structure) 

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

171 

172 nepy = _setup_nepy( 

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

174 ) 

175 

176 latent = nepy.get_latent_space() 

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

178 return latent 

179 

180 

181def get_potential_forces_and_virials( 

182 structure: Atoms, model_filename: Optional[str] = None, debug: bool = False 

183) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: 

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

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

186 

187 Parameters 

188 ---------- 

189 structure 

190 Input structure 

191 model_filename 

192 Path to NEP model. Defaults to None. 

193 debug 

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

195 

196 Returns 

197 ------- 

198 potential with shape `(natoms,)` 

199 forces with shape `(natoms, 3)` 

200 virials with shape `(natoms, 9)` 

201 """ 

202 if model_filename is None: 

203 raise ValueError('Model is undefined') 

204 

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

206 if model_type != 'potential': 

207 raise ValueError( 

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

209 ) 

210 

211 local_structure = structure.copy() 

212 natoms = len(local_structure) 

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

214 

215 nepy = _setup_nepy( 

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

217 ) 

218 

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

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

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

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

223 

224 

225def get_polarizability( 

226 structure: Atoms, 

227 model_filename: Optional[str] = None, 

228 debug: bool = False, 

229) -> np.ndarray: 

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

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

232 polarizability. 

233 

234 Parameters 

235 ---------- 

236 structure 

237 Input structure 

238 model_filename 

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

240 debug 

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

242 

243 Returns 

244 ------- 

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

246 """ 

247 if model_filename is None: 

248 raise ValueError('Model is undefined') 

249 

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

251 if model_type != 'polarizability': 

252 raise ValueError( 

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

254 ) 

255 

256 local_structure = structure.copy() 

257 natoms = len(local_structure) 

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

259 

260 nepy = _setup_nepy( 

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

262 ) 

263 # Components are in order xx yy zz xy yz zx. 

264 pol = nepy.get_polarizability() 

265 polarizability = np.array([ 

266 [pol[0], pol[3], pol[5]], 

267 [pol[3], pol[1], pol[4]], 

268 [pol[5], pol[4], pol[2]] 

269 ]) 

270 

271 return polarizability 

272 

273 

274def get_dipole( 

275 structure: Atoms, 

276 model_filename: Optional[str] = None, 

277 debug: bool = False, 

278) -> np.ndarray: 

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

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

281 

282 Parameters 

283 ---------- 

284 structure 

285 Input structure 

286 model_filename 

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

288 debug 

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

290 

291 Returns 

292 ------- 

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

294 """ 

295 if model_filename is None: 

296 raise ValueError('Model is undefined') 

297 

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

299 if model_type != 'dipole': 

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

301 

302 local_structure = structure.copy() 

303 natoms = len(local_structure) 

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

305 

306 nepy = _setup_nepy( 

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

308 ) 

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

310 

311 return dipole 

312 

313 

314def get_dipole_gradient( 

315 structure: Atoms, 

316 model_filename: Optional[str] = None, 

317 backend: str = 'c++', 

318 method: str = 'central difference', 

319 displacement: float = 0.01, 

320 charge: float = 1.0, 

321 nep_command: str = 'nep', 

322 debug: bool = False, 

323) -> np.ndarray: 

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

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

326 

327 Parameters 

328 ---------- 

329 structure 

330 Input structure 

331 model_filename 

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

333 backend 

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

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

336 Defaults to ``'c++'``. 

337 method 

338 Method for computing gradient with finite differences. 

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

340 Defaults to 'central difference' 

341 displacement 

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

343 charge 

344 System charge in units of the elemental charge. 

345 Used for correcting the dipoles before computing the gradient. 

346 Defaults to ``1.0``. 

347 nep_command 

348 Command for running the NEP executable. Defaults to ``'nep'``. 

349 debug 

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

351 

352 

353 Returns 

354 ------- 

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

356 """ 

357 if model_filename is None: 

358 raise ValueError('Model is undefined') 

359 

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

361 if model_type != 'dipole': 

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

363 

364 local_structure = structure.copy() 

365 

366 if backend == 'c++': 

367 dipole_gradient = _dipole_gradient_cpp( 

368 local_structure, 

369 model_filename, 

370 displacement=displacement, 

371 method=method, 

372 charge=charge, 

373 debug=debug, 

374 ) 

375 elif backend == 'python': 

376 dipole_gradient = _dipole_gradient_python( 

377 local_structure, 

378 model_filename, 

379 displacement=displacement, 

380 charge=charge, 

381 method=method, 

382 ) 

383 elif backend == 'nep': 

384 dipole_gradient = _dipole_gradient_nep( 

385 local_structure, 

386 model_filename, 

387 displacement=displacement, 

388 method=method, 

389 charge=charge, 

390 nep_command=nep_command, 

391 ) 

392 else: 

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

394 return dipole_gradient 

395 

396 

397def _dipole_gradient_cpp( 

398 structure: Atoms, 

399 model_filename: str, 

400 method: str = 'central difference', 

401 displacement: float = 0.01, 

402 charge: float = 1.0, 

403 debug: bool = False, 

404) -> np.ndarray: 

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

406 

407 Parameters 

408 ---------- 

409 structure 

410 Input structure 

411 model_filename 

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

413 method 

414 Method for computing gradient with finite differences. 

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

416 Defaults to ``'central difference'`` 

417 displacement 

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

419 charge 

420 System charge in units of the elemental charge. 

421 Used for correcting the dipoles before computing the gradient. 

422 Defaults to ``1.0``. 

423 

424 Returns 

425 ------- 

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

427 """ 

428 if displacement <= 0: 

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

430 

431 implemented_methods = { 

432 'forward difference': 0, 

433 'central difference': 1, 

434 'second order central difference': 2, 

435 } 

436 

437 if method not in implemented_methods.keys(): 

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

439 

440 local_structure = structure.copy() 

441 natoms = len(local_structure) 

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

443 nepy = _setup_nepy( 

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

445 ) 

446 dipole_gradient = np.array( 

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

448 ).reshape(natoms, 3, 3) 

449 return dipole_gradient 

450 

451 

452def _dipole_gradient_python( 

453 structure: Atoms, 

454 model_filename: str, 

455 method: str = 'central difference', 

456 displacement: float = 0.01, 

457 charge: float = 1.0, 

458) -> np.ndarray: 

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

460 

461 Parameters 

462 ---------- 

463 structure 

464 Input structure 

465 model_filename 

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

467 method 

468 Method for computing gradient with finite differences. 

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

470 Defaults to ``'central difference'`` 

471 displacement 

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

473 charge 

474 System charge in units of the elemental charge. 

475 Used for correcting the dipoles before computing the gradient. 

476 Defaults to ``1.0``. 

477 

478 Returns 

479 ------- 

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

481 """ 

482 if displacement <= 0: 

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

484 

485 N = len(structure) 

486 if method == 'forward difference': 

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

488 d = ( 

489 get_dipole(structure, model_filename) 

490 + charge * structure.get_center_of_mass() 

491 ) 

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

493 for atom in range(N): 

494 for cartesian in range(3): 

495 copy = structure.copy() 

496 positions = copy.get_positions() 

497 positions[atom, cartesian] += displacement 

498 copy.set_positions(positions) 

499 d_forward[atom, cartesian, :] = ( 

500 get_dipole(copy, model_filename) 

501 + charge * copy.get_center_of_mass() 

502 ) 

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

504 

505 elif method == 'central difference': 

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

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

508 for atom in range(N): 

509 for cartesian in range(3): 

510 # Forward displacements 

511 copy_forward = structure.copy() 

512 positions_forward = copy_forward.get_positions() 

513 positions_forward[atom, cartesian] += displacement 

514 copy_forward.set_positions(positions_forward) 

515 d_forward[atom, cartesian, :] = ( 

516 get_dipole(copy_forward, model_filename) 

517 + charge * copy_forward.get_center_of_mass() 

518 ) 

519 # Backwards displacement 

520 copy_backward = structure.copy() 

521 positions_backward = copy_backward.get_positions() 

522 positions_backward[atom, cartesian] -= displacement 

523 copy_backward.set_positions(positions_backward) 

524 d_backward[atom, cartesian, :] = ( 

525 get_dipole(copy_backward, model_filename) 

526 + charge * copy_backward.get_center_of_mass() 

527 ) 

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

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

530 # Coefficients from 

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

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

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

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

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

536 for atom in range(N): 

537 for cartesian in range(3): 

538 copy = structure.copy() 

539 positions = copy.get_positions() 

540 # Forward displacements 

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

542 copy.set_positions(positions) 

543 d_forward_one_h[atom, cartesian, :] = ( 

544 get_dipole(copy, model_filename) 

545 + charge * copy.get_center_of_mass() 

546 ) 

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

548 copy.set_positions(positions) 

549 d_forward_two_h[atom, cartesian, :] = ( 

550 get_dipole(copy, model_filename) 

551 + charge * copy.get_center_of_mass() 

552 ) 

553 # Backwards displacement 

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

555 copy.set_positions(positions) 

556 d_backward_one_h[atom, cartesian, :] = ( 

557 get_dipole(copy, model_filename) 

558 + charge * copy.get_center_of_mass() 

559 ) 

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

561 copy.set_positions(positions) 

562 d_backward_two_h[atom, cartesian, :] = ( 

563 get_dipole(copy, model_filename) 

564 + charge * copy.get_center_of_mass() 

565 ) 

566 c0 = -1.0 / 12.0 

567 c1 = 2.0 / 3.0 

568 gradient = ( 

569 c0 * d_forward_two_h 

570 + c1 * d_forward_one_h 

571 - c1 * d_backward_one_h 

572 - c0 * d_backward_two_h 

573 ) / displacement 

574 else: 

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

576 return gradient 

577 

578 

579def _dipole_gradient_nep( 

580 structure: Atoms, 

581 model_filename: str, 

582 method: str = 'central difference', 

583 displacement: float = 0.01, 

584 charge: float = 1.0, 

585 nep_command: str = 'nep', 

586) -> np.ndarray: 

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

588 

589 Parameters 

590 ---------- 

591 structure 

592 Input structure 

593 model_filename 

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

595 method 

596 Method for computing gradient with finite differences. 

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

598 Defaults to ``'central difference'`` 

599 displacement 

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

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

602 which might be due to rounding errors. 

603 charge 

604 System charge in units of the elemental charge. 

605 Used for correcting the dipoles before computing the gradient. 

606 Defaults to 1.0. 

607 nep_command 

608 Command for running the NEP executable. Defaults to ``'nep'``. 

609 

610 

611 Returns 

612 ------- 

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

614 """ 

615 if displacement <= 0: 

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

617 

618 if displacement < 1e-2: 

619 warnings.warn( 

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

621 ) 

622 

623 N = len(structure) 

624 if method == 'forward difference': 

625 structure = _set_dummy_energy_forces(structure) 

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

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

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

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

630 for atom in range(N): 

631 for cartesian in range(3): 

632 copy = structure.copy() 

633 positions = copy.get_positions() 

634 positions[atom, cartesian] += displacement 

635 copy.set_positions(positions) 

636 copy = _set_dummy_energy_forces(copy) 

637 structures.append(copy) 

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

639 charge * copy.get_center_of_mass() 

640 ) 

641 

642 dipoles = ( 

643 _predict_dipole_batch(structures, model_filename, nep_command) * N 

644 ) # dipole/atom, shape (3N+1, 3) 

645 dipoles += corrections 

646 

647 d = dipoles[0, :] 

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

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

650 

651 elif method == 'central difference': 

652 structures_forward = [] # will hold 3N structures 

653 structures_backward = [] # will hold 3N structures 

654 

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

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

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

658 

659 for atom in range(N): 

660 for cartesian in range(3): 

661 # Forward displacements 

662 copy_forward = structure.copy() 

663 positions_forward = copy_forward.get_positions() 

664 positions_forward[atom, cartesian] += displacement 

665 copy_forward.set_positions(positions_forward) 

666 copy_forward = _set_dummy_energy_forces(copy_forward) 

667 structures_forward.append(copy_forward) 

668 corrections_forward[atom * 3 + cartesian] = ( 

669 charge * copy_forward.get_center_of_mass() 

670 ) 

671 

672 # Backwards displacement 

673 copy_backward = structure.copy() 

674 positions_backward = copy_backward.get_positions() 

675 positions_backward[atom, cartesian] -= displacement 

676 copy_backward.set_positions(positions_backward) 

677 copy_backward = _set_dummy_energy_forces(copy_backward) 

678 structures_backward.append(copy_backward) 

679 corrections_backward[atom * 3 + cartesian] = ( 

680 charge * copy_backward.get_center_of_mass() 

681 ) 

682 

683 structures = structures_forward + structures_backward 

684 dipoles = ( 

685 _predict_dipole_batch(structures, model_filename, nep_command) * N 

686 ) # dipole/atom, shape (6N, 3) 

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

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

689 

690 d_forward += corrections_forward 

691 d_backward += corrections_backward 

692 

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

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

695 

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

697 else: 

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

699 return gradient 

700 

701 

702def _set_dummy_energy_forces(structure: Atoms) -> Atoms: 

703 """Sets the energies and forces of structure to zero. 

704 

705 Parameters 

706 ---------- 

707 structure 

708 Input structure 

709 

710 

711 Returns 

712 ------- Copy of structure, with SinglePointCalculator with zero energy and force. 

713 """ 

714 from ase.calculators.singlepoint import SinglePointCalculator 

715 

716 copy = structure.copy() 

717 N = len(copy) 

718 energy = 0 

719 

720 forces = np.zeros((N, 3)) 

721 dummy = SinglePointCalculator(copy, **{'energy': energy, 'forces': forces}) 

722 copy.calc = dummy 

723 return copy 

724 

725 

726def _predict_dipole_batch( 

727 structures: List[Atoms], model_filename: str, nep_command: str = 'nep' 

728) -> np.ndarray: 

729 """Predicts dipoles for a set of structures using the NEP executable 

730 Note that the units are in (dipole units)/atom. 

731 

732 Parameters 

733 ---------- 

734 structure 

735 Input structures 

736 model_filename 

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

738 nep_command 

739 Command for running the NEP executable. Defaults to ``'nep'``. 

740 

741 

742 Returns 

743 ------- Predicted dipoles, with shape (len(structures), 3). 

744 """ 

745 import shutil 

746 from os.path import join as join_path 

747 from subprocess import run 

748 from tempfile import TemporaryDirectory 

749 

750 from calorine.nep import read_model, write_nepfile, write_structures 

751 

752 with TemporaryDirectory() as directory: 

753 shutil.copy2(model_filename, join_path(directory, 'nep.txt')) 

754 model = read_model(model_filename) 

755 

756 parameters = dict( 

757 prediction=1, 

758 mode=1, 

759 version=model.version, 

760 n_max=[model.n_max_radial, model.n_max_angular], 

761 cutoff=[model.radial_cutoff, model.angular_cutoff], 

762 basis_size=[model.n_basis_radial, model.n_basis_radial], 

763 l_max=[model.l_max_3b, model.l_max_4b, model.l_max_5b], 

764 neuron=model.n_neuron, 

765 type=[len(model.types), *model.types], 

766 ) 

767 

768 write_nepfile(parameters, directory) 

769 file = join_path(directory, 'train.xyz') 

770 write_structures(file, structures) 

771 

772 # Execute nep 

773 completed = run([nep_command], cwd=directory, capture_output=True) 

774 completed.check_returncode() 

775 

776 # Read results 

777 dipoles = np.loadtxt(join_path(directory, 'dipole_train.out')) 

778 if len(dipoles.shape) == 1: 

779 dipoles = dipoles.reshape(1, -1) 

780 return dipoles[:, :3] 

781 

782 

783def _check_components_polarizability_gradient(component: Union[str, List[str]]) -> List[int]: 

784 """ 

785 Verifies that the selected components are ok. 

786 """ 

787 allowed_components = { 

788 'x': 0, 

789 'y': 1, 

790 'z': 2, 

791 } 

792 

793 # Check if chosen components are ok 

794 if component == 'full': 

795 components_to_compute = [0, 1, 2] 

796 else: 

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

798 components_to_compute = [] 

799 for c in components_list: 

800 if c in allowed_components.keys(): 

801 components_to_compute.append(allowed_components[c]) 

802 elif c == 'full': 

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

804 else: 

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

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

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

808 return components_to_compute 

809 

810 

811def get_polarizability_gradient( 

812 structure: Atoms, 

813 model_filename: Optional[str] = None, 

814 displacement: float = 0.01, 

815 component: Union[str, List[str]] = 'full', 

816 debug: bool = False, 

817) -> np.ndarray: 

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

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

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

821 method with a C++ backend. 

822 

823 Parameters 

824 ---------- 

825 structure 

826 Input structure. 

827 model_filename 

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

829 displacement 

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

831 component 

832 Component or components of the polarizability tensor that the gradient 

833 should be computed for. 

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

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

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

837 Multiple components may be specified. 

838 Defaults to ``full``. 

839 debug 

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

841 

842 

843 Returns 

844 ------- 

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

846 is the number of components chosen. 

847 """ 

848 if model_filename is None: 

849 raise ValueError('Model is undefined') 

850 

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

852 if model_type != 'polarizability': 

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

854 

855 components_to_compute = _check_components_polarizability_gradient(component) 

856 

857 local_structure = structure.copy() 

858 polarizability_gradient = _polarizability_gradient_cpp( 

859 local_structure, 

860 model_filename, 

861 displacement=displacement, 

862 components=components_to_compute, 

863 debug=debug, 

864 ) 

865 return polarizability_gradient 

866 

867 

868def _polarizability_gradient_to_3x3(natoms, pg): 

869 """ 

870 Converts a polarizability gradient tensor with 

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

872 The 6 items in the polarizability gradient have 

873 the following order from GPUMD: xx, yy, zz, xy, yz, zx. 

874 """ 

875 gpumd_to_3x3_indices = np.array([ 

876 [0, 3, 5], 

877 [3, 1, 4], 

878 [5, 4, 2] 

879 ]) 

880 polgrad_3x3 = np.zeros((natoms, 3, 3, 3)) 

881 for i in range(3): 

882 for j in range(3): 

883 index = gpumd_to_3x3_indices[i, j] 

884 polgrad_3x3[:, :, i, j] = pg[:, :, index] 

885 return polgrad_3x3 

886 

887 

888def _polarizability_gradient_cpp( 

889 structure: Atoms, 

890 model_filename: str, 

891 displacement: float, 

892 components: List[int], 

893 debug: bool = False, 

894) -> np.ndarray: 

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