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: str, 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. 

127 debug 

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

129 

130 Returns 

131 ------- 

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

133 """ 

134 local_structure = structure.copy() 

135 natoms = len(local_structure) 

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

137 

138 nepy = _setup_nepy( 

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

140 ) 

141 all_descriptors = nepy.get_descriptors() 

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

143 return descriptors_per_atom 

144 

145 

146def get_latent_space( 

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

148) -> np.ndarray: 

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

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

151 

152 Parameters 

153 ---------- 

154 structure 

155 Input structure 

156 model_filename 

157 Path to NEP model. Defaults to None. 

158 debug 

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

160 

161 Returns 

162 ------- 

163 Activation with shape `(natoms, N_neurons)` 

164 """ 

165 if model_filename is None: 

166 raise ValueError('Model is undefined') 

167 local_structure = structure.copy() 

168 natoms = len(local_structure) 

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

170 

171 nepy = _setup_nepy( 

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

173 ) 

174 

175 latent = nepy.get_latent_space() 

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

177 return latent 

178 

179 

180def get_potential_forces_and_virials( 

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

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

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

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

185 

186 Parameters 

187 ---------- 

188 structure 

189 Input structure 

190 model_filename 

191 Path to NEP model. Defaults to None. 

192 debug 

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

194 

195 Returns 

196 ------- 

197 potential with shape `(natoms,)` 

198 forces with shape `(natoms, 3)` 

199 virials with shape `(natoms, 9)` 

200 """ 

201 if model_filename is None: 

202 raise ValueError('Model is undefined') 

203 

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

205 if model_type != 'potential': 

206 raise ValueError( 

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

208 ) 

209 

210 local_structure = structure.copy() 

211 natoms = len(local_structure) 

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

213 

214 nepy = _setup_nepy( 

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

216 ) 

217 

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

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

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

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

222 

223 

224def get_polarizability( 

225 structure: Atoms, 

226 model_filename: Optional[str] = None, 

227 debug: bool = False, 

228) -> np.ndarray: 

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

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

231 polarizability. 

232 

233 Parameters 

234 ---------- 

235 structure 

236 Input structure 

237 model_filename 

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

239 debug 

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

241 

242 Returns 

243 ------- 

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

245 """ 

246 if model_filename is None: 

247 raise ValueError('Model is undefined') 

248 

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

250 if model_type != 'polarizability': 

251 raise ValueError( 

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

253 ) 

254 

255 local_structure = structure.copy() 

256 natoms = len(local_structure) 

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

258 

259 nepy = _setup_nepy( 

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

261 ) 

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

263 pol = nepy.get_polarizability() 

264 polarizability = np.array([ 

265 [pol[0], pol[3], pol[5]], 

266 [pol[3], pol[1], pol[4]], 

267 [pol[5], pol[4], pol[2]] 

268 ]) 

269 

270 return polarizability 

271 

272 

273def get_dipole( 

274 structure: Atoms, 

275 model_filename: Optional[str] = None, 

276 debug: bool = False, 

277) -> np.ndarray: 

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

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

280 

281 Parameters 

282 ---------- 

283 structure 

284 Input structure 

285 model_filename 

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

287 debug 

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

289 

290 Returns 

291 ------- 

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

293 """ 

294 if model_filename is None: 

295 raise ValueError('Model is undefined') 

296 

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

298 if model_type != 'dipole': 

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

300 

301 local_structure = structure.copy() 

302 natoms = len(local_structure) 

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

304 

305 nepy = _setup_nepy( 

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

307 ) 

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

309 

310 return dipole 

311 

312 

313def get_dipole_gradient( 

314 structure: Atoms, 

315 model_filename: Optional[str] = None, 

316 backend: str = 'c++', 

317 method: str = 'central difference', 

318 displacement: float = 0.01, 

319 charge: float = 1.0, 

320 nep_command: str = 'nep', 

321 debug: bool = False, 

322) -> np.ndarray: 

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

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

325 

326 Parameters 

327 ---------- 

328 structure 

329 Input structure 

330 model_filename 

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

332 backend 

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

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

335 Defaults to ``'c++'``. 

336 method 

337 Method for computing gradient with finite differences. 

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

339 Defaults to 'central difference' 

340 displacement 

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

342 charge 

343 System charge in units of the elemental charge. 

344 Used for correcting the dipoles before computing the gradient. 

345 Defaults to ``1.0``. 

346 nep_command 

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

348 debug 

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

350 

351 

352 Returns 

353 ------- 

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

355 """ 

356 if model_filename is None: 

357 raise ValueError('Model is undefined') 

358 

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

360 if model_type != 'dipole': 

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

362 

363 local_structure = structure.copy() 

364 

365 if backend == 'c++': 

366 dipole_gradient = _dipole_gradient_cpp( 

367 local_structure, 

368 model_filename, 

369 displacement=displacement, 

370 method=method, 

371 charge=charge, 

372 debug=debug, 

373 ) 

374 elif backend == 'python': 

375 dipole_gradient = _dipole_gradient_python( 

376 local_structure, 

377 model_filename, 

378 displacement=displacement, 

379 charge=charge, 

380 method=method, 

381 ) 

382 elif backend == 'nep': 

383 dipole_gradient = _dipole_gradient_nep( 

384 local_structure, 

385 model_filename, 

386 displacement=displacement, 

387 method=method, 

388 charge=charge, 

389 nep_command=nep_command, 

390 ) 

391 else: 

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

393 return dipole_gradient 

394 

395 

396def _dipole_gradient_cpp( 

397 structure: Atoms, 

398 model_filename: str, 

399 method: str = 'central difference', 

400 displacement: float = 0.01, 

401 charge: float = 1.0, 

402 debug: bool = False, 

403) -> np.ndarray: 

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

405 

406 Parameters 

407 ---------- 

408 structure 

409 Input structure 

410 model_filename 

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

412 method 

413 Method for computing gradient with finite differences. 

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

415 Defaults to ``'central difference'`` 

416 displacement 

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

418 charge 

419 System charge in units of the elemental charge. 

420 Used for correcting the dipoles before computing the gradient. 

421 Defaults to ``1.0``. 

422 

423 Returns 

424 ------- 

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

426 """ 

427 if displacement <= 0: 

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

429 

430 implemented_methods = { 

431 'forward difference': 0, 

432 'central difference': 1, 

433 'second order central difference': 2, 

434 } 

435 

436 if method not in implemented_methods.keys(): 

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

438 

439 local_structure = structure.copy() 

440 natoms = len(local_structure) 

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

442 nepy = _setup_nepy( 

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

444 ) 

445 dipole_gradient = np.array( 

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

447 ).reshape(natoms, 3, 3) 

448 return dipole_gradient 

449 

450 

451def _dipole_gradient_python( 

452 structure: Atoms, 

453 model_filename: str, 

454 method: str = 'central difference', 

455 displacement: float = 0.01, 

456 charge: float = 1.0, 

457) -> np.ndarray: 

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

459 

460 Parameters 

461 ---------- 

462 structure 

463 Input structure 

464 model_filename 

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

466 method 

467 Method for computing gradient with finite differences. 

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

469 Defaults to ``'central difference'`` 

470 displacement 

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

472 charge 

473 System charge in units of the elemental charge. 

474 Used for correcting the dipoles before computing the gradient. 

475 Defaults to ``1.0``. 

476 

477 Returns 

478 ------- 

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

480 """ 

481 if displacement <= 0: 

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

483 

484 N = len(structure) 

485 if method == 'forward difference': 

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

487 d = ( 

488 get_dipole(structure, model_filename) 

489 + charge * structure.get_center_of_mass() 

490 ) 

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

492 for atom in range(N): 

493 for cartesian in range(3): 

494 copy = structure.copy() 

495 positions = copy.get_positions() 

496 positions[atom, cartesian] += displacement 

497 copy.set_positions(positions) 

498 d_forward[atom, cartesian, :] = ( 

499 get_dipole(copy, model_filename) 

500 + charge * copy.get_center_of_mass() 

501 ) 

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

503 

504 elif method == 'central difference': 

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

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

507 for atom in range(N): 

508 for cartesian in range(3): 

509 # Forward displacements 

510 copy_forward = structure.copy() 

511 positions_forward = copy_forward.get_positions() 

512 positions_forward[atom, cartesian] += displacement 

513 copy_forward.set_positions(positions_forward) 

514 d_forward[atom, cartesian, :] = ( 

515 get_dipole(copy_forward, model_filename) 

516 + charge * copy_forward.get_center_of_mass() 

517 ) 

518 # Backwards displacement 

519 copy_backward = structure.copy() 

520 positions_backward = copy_backward.get_positions() 

521 positions_backward[atom, cartesian] -= displacement 

522 copy_backward.set_positions(positions_backward) 

523 d_backward[atom, cartesian, :] = ( 

524 get_dipole(copy_backward, model_filename) 

525 + charge * copy_backward.get_center_of_mass() 

526 ) 

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

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

529 # Coefficients from 

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

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

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

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

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

535 for atom in range(N): 

536 for cartesian in range(3): 

537 copy = structure.copy() 

538 positions = copy.get_positions() 

539 # Forward displacements 

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

541 copy.set_positions(positions) 

542 d_forward_one_h[atom, cartesian, :] = ( 

543 get_dipole(copy, model_filename) 

544 + charge * copy.get_center_of_mass() 

545 ) 

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

547 copy.set_positions(positions) 

548 d_forward_two_h[atom, cartesian, :] = ( 

549 get_dipole(copy, model_filename) 

550 + charge * copy.get_center_of_mass() 

551 ) 

552 # Backwards displacement 

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

554 copy.set_positions(positions) 

555 d_backward_one_h[atom, cartesian, :] = ( 

556 get_dipole(copy, model_filename) 

557 + charge * copy.get_center_of_mass() 

558 ) 

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

560 copy.set_positions(positions) 

561 d_backward_two_h[atom, cartesian, :] = ( 

562 get_dipole(copy, model_filename) 

563 + charge * copy.get_center_of_mass() 

564 ) 

565 c0 = -1.0 / 12.0 

566 c1 = 2.0 / 3.0 

567 gradient = ( 

568 c0 * d_forward_two_h 

569 + c1 * d_forward_one_h 

570 - c1 * d_backward_one_h 

571 - c0 * d_backward_two_h 

572 ) / displacement 

573 else: 

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

575 return gradient 

576 

577 

578def _dipole_gradient_nep( 

579 structure: Atoms, 

580 model_filename: str, 

581 method: str = 'central difference', 

582 displacement: float = 0.01, 

583 charge: float = 1.0, 

584 nep_command: str = 'nep', 

585) -> np.ndarray: 

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

587 

588 Parameters 

589 ---------- 

590 structure 

591 Input structure 

592 model_filename 

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

594 method 

595 Method for computing gradient with finite differences. 

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

597 Defaults to ``'central difference'`` 

598 displacement 

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

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

601 which might be due to rounding errors. 

602 charge 

603 System charge in units of the elemental charge. 

604 Used for correcting the dipoles before computing the gradient. 

605 Defaults to 1.0. 

606 nep_command 

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

608 

609 

610 Returns 

611 ------- 

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

613 """ 

614 if displacement <= 0: 

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

616 

617 if displacement < 1e-2: 

618 warnings.warn( 

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

620 ) 

621 

622 N = len(structure) 

623 if method == 'forward difference': 

624 structure = _set_dummy_energy_forces(structure) 

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

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

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

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

629 for atom in range(N): 

630 for cartesian in range(3): 

631 copy = structure.copy() 

632 positions = copy.get_positions() 

633 positions[atom, cartesian] += displacement 

634 copy.set_positions(positions) 

635 copy = _set_dummy_energy_forces(copy) 

636 structures.append(copy) 

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

638 charge * copy.get_center_of_mass() 

639 ) 

640 

641 dipoles = ( 

642 _predict_dipole_batch(structures, model_filename, nep_command) * N 

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

644 dipoles += corrections 

645 

646 d = dipoles[0, :] 

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

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

649 

650 elif method == 'central difference': 

651 structures_forward = [] # will hold 3N structures 

652 structures_backward = [] # will hold 3N structures 

653 

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

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

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

657 

658 for atom in range(N): 

659 for cartesian in range(3): 

660 # Forward displacements 

661 copy_forward = structure.copy() 

662 positions_forward = copy_forward.get_positions() 

663 positions_forward[atom, cartesian] += displacement 

664 copy_forward.set_positions(positions_forward) 

665 copy_forward = _set_dummy_energy_forces(copy_forward) 

666 structures_forward.append(copy_forward) 

667 corrections_forward[atom * 3 + cartesian] = ( 

668 charge * copy_forward.get_center_of_mass() 

669 ) 

670 

671 # Backwards displacement 

672 copy_backward = structure.copy() 

673 positions_backward = copy_backward.get_positions() 

674 positions_backward[atom, cartesian] -= displacement 

675 copy_backward.set_positions(positions_backward) 

676 copy_backward = _set_dummy_energy_forces(copy_backward) 

677 structures_backward.append(copy_backward) 

678 corrections_backward[atom * 3 + cartesian] = ( 

679 charge * copy_backward.get_center_of_mass() 

680 ) 

681 

682 structures = structures_forward + structures_backward 

683 dipoles = ( 

684 _predict_dipole_batch(structures, model_filename, nep_command) * N 

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

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

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

688 

689 d_forward += corrections_forward 

690 d_backward += corrections_backward 

691 

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

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

694 

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

696 else: 

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

698 return gradient 

699 

700 

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

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

703 

704 Parameters 

705 ---------- 

706 structure 

707 Input structure 

708 

709 

710 Returns 

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

712 """ 

713 from ase.calculators.singlepoint import SinglePointCalculator 

714 

715 copy = structure.copy() 

716 N = len(copy) 

717 energy = 0 

718 

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

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

721 copy.calc = dummy 

722 return copy 

723 

724 

725def _predict_dipole_batch( 

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

727) -> np.ndarray: 

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

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

730 

731 Parameters 

732 ---------- 

733 structure 

734 Input structures 

735 model_filename 

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

737 nep_command 

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

739 

740 

741 Returns 

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

743 """ 

744 import shutil 

745 from os.path import join as join_path 

746 from subprocess import run 

747 from tempfile import TemporaryDirectory 

748 

749 from calorine.nep import read_model, write_nepfile, write_structures 

750 

751 with TemporaryDirectory() as directory: 

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

753 model = read_model(model_filename) 

754 

755 parameters = dict( 

756 prediction=1, 

757 mode=1, 

758 version=model.version, 

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

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

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

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

763 neuron=model.n_neuron, 

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

765 ) 

766 

767 write_nepfile(parameters, directory) 

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

769 write_structures(file, structures) 

770 

771 # Execute nep 

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

773 completed.check_returncode() 

774 

775 # Read results 

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

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

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

779 return dipoles[:, :3] 

780 

781 

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

783 """ 

784 Verifies that the selected components are ok. 

785 """ 

786 allowed_components = { 

787 'x': 0, 

788 'y': 1, 

789 'z': 2, 

790 } 

791 

792 # Check if chosen components are ok 

793 if component == 'full': 

794 components_to_compute = [0, 1, 2] 

795 else: 

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

797 components_to_compute = [] 

798 for c in components_list: 

799 if c in allowed_components.keys(): 

800 components_to_compute.append(allowed_components[c]) 

801 elif c == 'full': 

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

803 else: 

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

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

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

807 return components_to_compute 

808 

809 

810def get_polarizability_gradient( 

811 structure: Atoms, 

812 model_filename: Optional[str] = None, 

813 displacement: float = 0.01, 

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

815 debug: bool = False, 

816) -> np.ndarray: 

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

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

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

820 method with a C++ backend. 

821 

822 Parameters 

823 ---------- 

824 structure 

825 Input structure. 

826 model_filename 

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

828 displacement 

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

830 component 

831 Component or components of the polarizability tensor that the gradient 

832 should be computed for. 

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

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

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

836 Multiple components may be specified. 

837 Defaults to ``full``. 

838 debug 

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

840 

841 

842 Returns 

843 ------- 

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

845 is the number of components chosen. 

846 """ 

847 if model_filename is None: 

848 raise ValueError('Model is undefined') 

849 

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

851 if model_type != 'polarizability': 

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

853 

854 components_to_compute = _check_components_polarizability_gradient(component) 

855 

856 local_structure = structure.copy() 

857 polarizability_gradient = _polarizability_gradient_cpp( 

858 local_structure, 

859 model_filename, 

860 displacement=displacement, 

861 components=components_to_compute, 

862 debug=debug, 

863 ) 

864 return polarizability_gradient 

865 

866 

867def _polarizability_gradient_to_3x3(natoms, pg): 

868 """ 

869 Converts a polarizability gradient tensor with 

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

871 The 6 items in the polarizability gradient have 

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

873 """ 

874 gpumd_to_3x3_indices = np.array([ 

875 [0, 3, 5], 

876 [3, 1, 4], 

877 [5, 4, 2] 

878 ]) 

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

880 for i in range(3): 

881 for j in range(3): 

882 index = gpumd_to_3x3_indices[i, j] 

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

884 return polgrad_3x3 

885 

886 

887def _polarizability_gradient_cpp( 

888 structure: Atoms, 

889 model_filename: str, 

890 displacement: float, 

891 components: List[int],