Coverage for calorine/nep/nep.py: 100%
294 statements
« prev ^ index » next coverage.py v7.2.5, created at 2024-10-17 08:55 +0000
1import contextlib
2import os
3import shutil
4import tempfile
5import warnings
6from pathlib import Path
7from typing import List, Optional, Tuple, Union
9import numpy as np
10from ase import Atoms
12import _nepy
13from calorine.nep.model import _get_nep_contents
16def _create_dummy_nep2(model_filename: str, symbols: List[str]) -> None:
17 """Create a dummy NEP2 model, i.e., a model for which there are no descriptor parameters.
18 This is to be used when one wants descriptors not pertaining to a specific NEP3 model.
20 Parameters
21 ----------
22 model_filename
23 Path to which the NEP2 model will be saved.
24 symbols
25 Atomic elements in the configuration for which to compute descriptors.
26 """
27 unique_symbols = []
28 for sym in symbols:
29 if sym not in unique_symbols:
30 unique_symbols.append(sym)
31 with open(model_filename, 'w') as f:
32 f.write(f'nep {len(unique_symbols)} ')
33 for symbol in unique_symbols:
34 f.write(f'{symbol} ')
35 f.write('\n')
36 # Write rest of header
37 f.write('cutoff 6 4\nn_max 15 8\nl_max 4\nANN 30 0\n')
38 # Write dummy parameters
39 # The number of parameters in the network is:
40 # N_par = (N_des+2)N_neu + 1 + N_typ**2 (n_max^R + n_max^A + 2)
41 # Default settings: 1621 + N_typ**2 * 25
42 # The last 52 1:s are the normalization parameters for the descriptors.
43 for _ in range(1621 + len(unique_symbols) ** 2 * 25 + 52):
44 f.write(f' {1e0:e}\n')
47def _create_tmp_dir(debug: bool) -> str:
48 """Create temporary directory.
50 Parameters
51 ----------
52 debug
53 Flag that indicates if debugging. Use a hardcoded path when debugging.
55 Returns
56 -------
57 str
58 Path to temporary directory
59 """
60 if debug:
61 tmp_dir = './tmp_nepy/'
62 try:
63 os.mkdir(tmp_dir)
64 return tmp_dir
65 except FileExistsError:
66 raise FileExistsError('Please delete or move the conflicting directory.')
67 return tempfile.mkdtemp()
70def _clean_tmp_dir(directory: str) -> None:
71 """Remove temporary directory.
73 Parameters
74 ----------
75 directory
76 Path to temporary directory
77 """
78 shutil.rmtree(directory)
81def _get_atomic_properties(
82 structure: Atoms,
83) -> Tuple[List[float], List[str], List[float]]:
84 """Fetches cell, symbols and positions for a structure. Since NEP_CPU requires a cell, if the
85 structure has no cell a default cubic cell with a side of 100 Å will be used.
87 Parameters
88 ----------
89 structure
90 Atoms object representing the structure.
92 Returns
93 -------
94 List[float]
95 Cell vectors
96 List[str]
97 Atomic species
98 List[float]
99 Atom positions
100 List[float]
101 Atom masses
102 """
103 if structure.cell.rank == 0:
104 warnings.warn('Using default unit cell (cubic with side 100 Å).')
105 set_default_cell(structure)
107 c = structure.get_cell(complete=True).flatten()
108 cell = [c[0], c[3], c[6], c[1], c[4], c[7], c[2], c[5], c[8]]
110 symbols = structure.get_chemical_symbols()
111 positions = list(
112 structure.get_positions().T.flatten()
113 ) # [x1, ..., xN, y1, ... yN,...]
114 masses = structure.get_masses()
115 return cell, symbols, positions, masses
118def _setup_nepy(
119 model_filename: str,
120 natoms: int,
121 cell: List[float],
122 symbols: List[str],
123 positions: List[float],
124 masses: List[float],
125 debug: bool,
126) -> _nepy.NEPY:
127 """Sets up an instance of the NEPY pybind11 interface to NEP_CPU.
129 Parameters
130 ----------
131 model_filename
132 Path to model.
133 natoms:
134 Number of atoms in the structure.
135 cell:
136 Cell vectors.
137 symbols:
138 Atom species.
139 positions:
140 Atom positions.
141 masses:
142 Atom masses.
143 debug:
144 Flag to control if the output from NEP_CPU will be printed.
146 Returns
147 -------
148 NEPY
149 NEPY interface
150 """
151 # Ensure that `model_filename` exists to avoid segfault in pybind11 code
152 if not os.path.isfile(model_filename):
153 raise ValueError(f'{Path(model_filename)} does not exist')
155 # Disable output from C++ code by default
156 if debug:
157 nepy = _nepy.NEPY(model_filename, natoms, cell, symbols, positions, masses)
158 else:
159 with open(os.devnull, 'w') as f:
160 with contextlib.redirect_stdout(f):
161 with contextlib.redirect_stderr(f):
162 nepy = _nepy.NEPY(
163 model_filename, natoms, cell, symbols, positions, masses
164 )
165 return nepy
168def set_default_cell(structure: Atoms, box_length: float = 100):
169 """Adds a cubic box to an Atoms object. Atoms object is edited in-place.
171 Parameters
172 ----------
173 structure
174 Structure to add box to
175 box_length
176 Cubic box side length in Å, by default 100
177 """
178 structure.set_cell([[box_length, 0, 0], [0, box_length, 0], [0, 0, box_length]])
179 structure.center()
182def get_descriptors(
183 structure: Atoms, model_filename: Optional[str] = None, debug: bool = False
184) -> np.ndarray:
185 """Calculates the NEP descriptors for a given structure. A NEP model defined by a nep.txt
186 can additionally be provided to get the NEP3 model specific descriptors. If no model is
187 provided, a dummy NEP2 model suitable for the provided structure will be created and used.
189 Parameters
190 ----------
191 structure
192 Input structure
193 model_filename
194 Path to NEP model in ``nep.txt`` format. Defaults to ``None``.
195 debug
196 Flag to toggle debug mode. Makes the generated dummy NEP2 model available
197 in a local tmp directory, as well as prints GPUMD output. Defaults to ``False``.
199 Returns
200 -------
201 Descriptors for the supplied structure, with shape (natoms, descriptor components)
202 """
203 local_structure = structure.copy()
204 natoms = len(local_structure)
205 cell, symbols, positions, masses = _get_atomic_properties(local_structure)
207 use_dummy_nep2_potential = model_filename is None
208 if use_dummy_nep2_potential:
209 tmp_dir = _create_tmp_dir(debug)
210 model_filename = f'{tmp_dir}/nep.txt'
211 _create_dummy_nep2(model_filename, symbols)
212 else:
213 tmp_dir = None
215 nepy = _setup_nepy(
216 model_filename, natoms, cell, symbols, positions, masses, debug
217 )
218 all_descriptors = nepy.get_descriptors()
219 descriptors_per_atom = np.array(all_descriptors).reshape(-1, natoms).T
221 if use_dummy_nep2_potential and tmp_dir and not debug:
222 _clean_tmp_dir(tmp_dir)
223 if use_dummy_nep2_potential and tmp_dir and debug:
224 print(f'DEBUG - Directory containing dummy nep.in: {tmp_dir}')
225 return descriptors_per_atom
228def get_latent_space(
229 structure: Atoms, model_filename: Union[str, None] = None, debug: bool = False
230) -> np.ndarray:
231 """Calculates the latent space representation of a structure, i.e, the activiations in
232 the hidden layer. A NEP model defined by a `nep.txt` file needs to be provided.
234 Parameters
235 ----------
236 structure
237 Input structure
238 model_filename
239 Path to NEP model. Defaults to None.
240 debug
241 Flag to toggle debug mode. Prints GPUMD output. Defaults to False.
243 Returns
244 -------
245 Activation with shape `(natoms, N_neurons)`
246 """
247 if model_filename is None:
248 raise ValueError('Model must be defined!')
249 local_structure = structure.copy()
250 natoms = len(local_structure)
251 cell, symbols, positions, masses = _get_atomic_properties(local_structure)
253 nepy = _setup_nepy(
254 model_filename, natoms, cell, symbols, positions, masses, debug
255 )
257 latent = nepy.get_latent_space()
258 latent = np.array(latent).reshape(-1, natoms).T
259 return latent
262def get_potential_forces_and_virials(
263 structure: Atoms, model_filename: Optional[str] = None, debug: bool = False
264) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
265 """Calculates the per-atom potential, forces and virials for a given structure.
266 A NEP model defined by a `nep.txt` file needs to be provided.
268 Parameters
269 ----------
270 structure
271 Input structure
272 model_filename
273 Path to NEP model. Defaults to None.
274 debug
275 Flag to toggle debug mode. Prints GPUMD output. Defaults to False.
277 Returns
278 -------
279 potential with shape `(natoms,)`
280 forces with shape `(natoms, 3)`
281 virials with shape `(natoms, 9)`
282 """
283 if model_filename is None:
284 raise ValueError('Model must be defined!')
286 model_type = _get_nep_contents(model_filename)[0]['model_type']
287 if model_type != 'potential':
288 raise ValueError(
289 'A NEP model trained for predicting energies and forces must be used.'
290 )
292 local_structure = structure.copy()
293 natoms = len(local_structure)
294 cell, symbols, positions, masses = _get_atomic_properties(local_structure)
296 nepy = _setup_nepy(
297 model_filename, natoms, cell, symbols, positions, masses, debug
298 )
300 energies, forces, virials = nepy.get_potential_forces_and_virials()
301 forces_per_atom = np.array(forces).reshape(-1, natoms).T
302 virials_per_atom = np.array(virials).reshape(-1, natoms).T
303 return np.array(energies), forces_per_atom, virials_per_atom
306def get_polarizability(
307 structure: Atoms,
308 model_filename: Optional[str] = None,
309 debug: bool = False,
310) -> np.ndarray:
311 """Calculates the polarizability tensor for a given structure. A NEP model defined
312 by a ``nep.txt`` file needs to be provided. The model must be trained to predict the
313 polarizability.
315 Parameters
316 ----------
317 structure
318 Input structure
319 model_filename
320 Path to NEP model in ``nep.txt`` format. Defaults to ``None``.
321 debug
322 Flag to toggle debug mode. Prints GPUMD output. Defaults to ``False``.
324 Returns
325 -------
326 polarizability with shape ``(3, 3)``
327 """
328 if model_filename is None:
329 raise ValueError('Model must be defined!')
331 model_type = _get_nep_contents(model_filename)[0]['model_type']
332 if model_type != 'polarizability':
333 raise ValueError(
334 'A NEP model trained for predicting polarizability must be used.'
335 )
337 local_structure = structure.copy()
338 natoms = len(local_structure)
339 cell, symbols, positions, masses = _get_atomic_properties(local_structure)
341 nepy = _setup_nepy(
342 model_filename, natoms, cell, symbols, positions, masses, debug
343 )
344 # Components are in order xx yy zz xy yz zx.
345 pol = nepy.get_polarizability()
346 polarizability = np.array([
347 [pol[0], pol[3], pol[5]],
348 [pol[3], pol[1], pol[4]],
349 [pol[5], pol[4], pol[2]]
350 ])
352 return polarizability
355def get_dipole(
356 structure: Atoms,
357 model_filename: Optional[str] = None,
358 debug: bool = False,
359) -> np.ndarray:
360 """Calculates the dipole for a given structure. A NEP model defined by a
361 ``nep.txt`` file needs to be provided.
363 Parameters
364 ----------
365 structure
366 Input structure
367 model_filename
368 Path to NEP model in ``nep.txt`` format. Defaults to ``None``.
369 debug
370 Flag to toggle debug mode. Prints GPUMD output. Defaults to ``False``.
372 Returns
373 -------
374 dipole with shape ``(3,)``
375 """
376 if model_filename is None:
377 raise ValueError('Model must be defined!')
379 model_type = _get_nep_contents(model_filename)[0]['model_type']
380 if model_type != 'dipole':
381 raise ValueError('A NEP model trained for predicting dipoles must be used.')
383 local_structure = structure.copy()
384 natoms = len(local_structure)
385 cell, symbols, positions, masses = _get_atomic_properties(local_structure)
387 nepy = _setup_nepy(
388 model_filename, natoms, cell, symbols, positions, masses, debug
389 )
390 dipole = np.array(nepy.get_dipole())
392 return dipole
395def get_dipole_gradient(
396 structure: Atoms,
397 model_filename: Optional[str] = None,
398 backend: str = 'c++',
399 method: str = 'central difference',
400 displacement: float = 0.01,
401 charge: float = 1.0,
402 nep_command: str = 'nep',
403 debug: bool = False,
404) -> np.ndarray:
405 """Calculates the dipole gradient for a given structure using finite differences.
406 A NEP model defined by a `nep.txt` file needs to be provided.
408 Parameters
409 ----------
410 structure
411 Input structure
412 model_filename
413 Path to NEP model in ``nep.txt`` format. Defaults to ``None``.
414 backend
415 Backend to use for computing dipole gradient with finite differences.
416 One of ``'c++'`` (CPU), ``'python'`` (CPU) and ``'nep'`` (GPU).
417 Defaults to ``'c++'``.
418 method
419 Method for computing gradient with finite differences.
420 One of 'forward difference' and 'central difference'.
421 Defaults to 'central difference'
422 displacement
423 Displacement in Å to use for finite differences. Defaults to ``0.01``.
424 charge
425 System charge in units of the elemental charge.
426 Used for correcting the dipoles before computing the gradient.
427 Defaults to ``1.0``.
428 nep_command
429 Command for running the NEP executable. Defaults to ``'nep'``.
430 debug
431 Flag to toggle debug mode. Prints GPUMD output (if applicable). Defaults to ``False``.
434 Returns
435 -------
436 dipole gradient with shape ``(N, 3, 3)``
437 """
438 if model_filename is None:
439 raise ValueError('Model must be defined!')
441 model_type = _get_nep_contents(model_filename)[0]['model_type']
442 if model_type != 'dipole':
443 raise ValueError('A NEP model trained for predicting dipoles must be used.')
445 local_structure = structure.copy()
447 if backend == 'c++':
448 dipole_gradient = _dipole_gradient_cpp(
449 local_structure,
450 model_filename,
451 displacement=displacement,
452 method=method,
453 charge=charge,
454 debug=debug,
455 )
456 elif backend == 'python':
457 dipole_gradient = _dipole_gradient_python(
458 local_structure,
459 model_filename,
460 displacement=displacement,
461 charge=charge,
462 method=method,
463 )
464 elif backend == 'nep':
465 dipole_gradient = _dipole_gradient_nep(
466 local_structure,
467 model_filename,
468 displacement=displacement,
469 method=method,
470 charge=charge,
471 nep_command=nep_command,
472 )
473 else:
474 raise ValueError(f'Invalid backend {backend}')
475 return dipole_gradient
478def _dipole_gradient_cpp(
479 structure: Atoms,
480 model_filename: str,
481 method: str = 'central difference',
482 displacement: float = 0.01,
483 charge: float = 1.0,
484 debug: bool = False,
485) -> np.ndarray:
486 """Calculates the dipole gradient with finite differences, using NEP_CPU.
488 Parameters
489 ----------
490 structure
491 Input structure
492 model_filename
493 Path to NEP model in ``nep.txt`` format.
494 method
495 Method for computing gradient with finite differences.
496 One of ``'forward difference'`` and ``'central difference'``.
497 Defaults to ``'central difference'``
498 displacement
499 Displacement in Å to use for finite differences. Defaults to ``0.01``.
500 charge
501 System charge in units of the elemental charge.
502 Used for correcting the dipoles before computing the gradient.
503 Defaults to ``1.0``.
505 Returns
506 -------
507 dipole gradient with shape ``(N, 3, 3)``
508 """
509 if displacement <= 0:
510 raise ValueError('Displacement must be > 0 Å')
512 implemented_methods = {
513 'forward difference': 0,
514 'central difference': 1,
515 'second order central difference': 2,
516 }
518 if method not in implemented_methods.keys():
519 raise ValueError(f'Invalid method {method} for calculating gradient')
521 local_structure = structure.copy()
522 natoms = len(local_structure)
523 cell, symbols, positions, masses = _get_atomic_properties(local_structure)
524 nepy = _setup_nepy(
525 model_filename, natoms, cell, symbols, positions, masses, debug
526 )
527 dipole_gradient = np.array(
528 nepy.get_dipole_gradient(displacement, implemented_methods[method], charge)
529 ).reshape(natoms, 3, 3)
530 return dipole_gradient
533def _dipole_gradient_python(
534 structure: Atoms,
535 model_filename: str,
536 method: str = 'central difference',
537 displacement: float = 0.01,
538 charge: float = 1.0,
539) -> np.ndarray:
540 """Calculates the dipole gradient with finite differences, using the Python and get_dipole().
542 Parameters
543 ----------
544 structure
545 Input structure
546 model_filename
547 Path to NEP model in ``nep.txt`` format.
548 method
549 Method for computing gradient with finite differences.
550 One of ``'forward difference'`` and ``'central difference'``.
551 Defaults to ``'central difference'``
552 displacement
553 Displacement in Å to use for finite differences. Defaults to ``0.01``.
554 charge
555 System charge in units of the elemental charge.
556 Used for correcting the dipoles before computing the gradient.
557 Defaults to ``1.0``.
559 Returns
560 -------
561 dipole gradient with shape ``(N, 3, 3)``
562 """
563 if displacement <= 0:
564 raise ValueError('Displacement must be > 0 Å')
566 N = len(structure)
567 if method == 'forward difference':
568 # Correct all dipole by the permanent dipole, charge * center of mass
569 d = (
570 get_dipole(structure, model_filename)
571 + charge * structure.get_center_of_mass()
572 )
573 d_forward = np.zeros((N, 3, 3))
574 for atom in range(N):
575 for cartesian in range(3):
576 copy = structure.copy()
577 positions = copy.get_positions()
578 positions[atom, cartesian] += displacement
579 copy.set_positions(positions)
580 d_forward[atom, cartesian, :] = (
581 get_dipole(copy, model_filename)
582 + charge * copy.get_center_of_mass()
583 )
584 gradient = (d_forward - d[None, None, :]) / displacement
586 elif method == 'central difference':
587 d_forward = np.zeros((N, 3, 3))
588 d_backward = np.zeros((N, 3, 3))
589 for atom in range(N):
590 for cartesian in range(3):
591 # Forward displacements
592 copy_forward = structure.copy()
593 positions_forward = copy_forward.get_positions()
594 positions_forward[atom, cartesian] += displacement
595 copy_forward.set_positions(positions_forward)
596 d_forward[atom, cartesian, :] = (
597 get_dipole(copy_forward, model_filename)
598 + charge * copy_forward.get_center_of_mass()
599 )
600 # Backwards displacement
601 copy_backward = structure.copy()
602 positions_backward = copy_backward.get_positions()
603 positions_backward[atom, cartesian] -= displacement
604 copy_backward.set_positions(positions_backward)
605 d_backward[atom, cartesian, :] = (
606 get_dipole(copy_backward, model_filename)
607 + charge * copy_backward.get_center_of_mass()
608 )
609 gradient = (d_forward - d_backward) / (2 * displacement)
610 elif method == 'second order central difference':
611 # Coefficients from
612 # https://en.wikipedia.org/wiki/Finite_difference_coefficient#Central_finite_difference
613 d_forward_one_h = np.zeros((N, 3, 3))
614 d_forward_two_h = np.zeros((N, 3, 3))
615 d_backward_one_h = np.zeros((N, 3, 3))
616 d_backward_two_h = np.zeros((N, 3, 3))
617 for atom in range(N):
618 for cartesian in range(3):
619 copy = structure.copy()
620 positions = copy.get_positions()
621 # Forward displacements
622 positions[atom, cartesian] += displacement # + h
623 copy.set_positions(positions)
624 d_forward_one_h[atom, cartesian, :] = (
625 get_dipole(copy, model_filename)
626 + charge * copy.get_center_of_mass()
627 )
628 positions[atom, cartesian] += displacement # + 2h total
629 copy.set_positions(positions)
630 d_forward_two_h[atom, cartesian, :] = (
631 get_dipole(copy, model_filename)
632 + charge * copy.get_center_of_mass()
633 )
634 # Backwards displacement
635 positions[atom, cartesian] -= 3 * displacement # 2h - 3h = -h
636 copy.set_positions(positions)
637 d_backward_one_h[atom, cartesian, :] = (
638 get_dipole(copy, model_filename)
639 + charge * copy.get_center_of_mass()
640 )
641 positions[atom, cartesian] -= displacement # - 2h total
642 copy.set_positions(positions)
643 d_backward_two_h[atom, cartesian, :] = (
644 get_dipole(copy, model_filename)
645 + charge * copy.get_center_of_mass()
646 )
647 c0 = -1.0 / 12.0
648 c1 = 2.0 / 3.0
649 gradient = (
650 c0 * d_forward_two_h
651 + c1 * d_forward_one_h
652 - c1 * d_backward_one_h
653 - c0 * d_backward_two_h
654 ) / displacement
655 else:
656 raise ValueError(f'Invalid method {method} for calculating gradient')
657 return gradient
660def _dipole_gradient_nep(
661 structure: Atoms,
662 model_filename: str,
663 method: str = 'central difference',
664 displacement: float = 0.01,
665 charge: float = 1.0,
666 nep_command: str = 'nep',
667) -> np.ndarray:
668 """Calculates the dipole gradient with finite differences, using the NEP executable.
670 Parameters
671 ----------
672 structure
673 Input structure
674 model_filename
675 Path to NEP model in ``nep.txt`` format.
676 method
677 Method for computing gradient with finite differences.
678 One of ``'forward difference'`` and ``'central difference'``.
679 Defaults to ``'central difference'``
680 displacement
681 Displacement in Å to use for finite differences. Defaults to 0.01 Å.
682 Note that results are possibly unreliable for displacemen < 0.01,
683 which might be due to rounding errors.
684 charge
685 System charge in units of the elemental charge.
686 Used for correcting the dipoles before computing the gradient.
687 Defaults to 1.0.
688 nep_command
689 Command for running the NEP executable. Defaults to ``'nep'``.
692 Returns
693 -------
694 dipole gradient with shape ``(N, 3, 3)``
695 """
696 if displacement <= 0:
697 raise ValueError('Displacement must be > 0 Å')
699 if displacement < 1e-2:
700 warnings.warn(
701 'Dipole gradients with nep are unstable for displacements < 0.01 Å.'
702 )
704 N = len(structure)
705 if method == 'forward difference':
706 structure = _set_dummy_energy_forces(structure)
707 structures = [structure] # will hold 3N+1 structures
708 # Correct for the constant dipole, by adding charge * center of mass
709 corrections = np.zeros((3 * N + 1, 3))
710 corrections[0] = charge * structure.get_center_of_mass()
711 for atom in range(N):
712 for cartesian in range(3):
713 copy = structure.copy()
714 positions = copy.get_positions()
715 positions[atom, cartesian] += displacement
716 copy.set_positions(positions)
717 copy = _set_dummy_energy_forces(copy)
718 structures.append(copy)
719 corrections[1 + atom * 3 + cartesian] = (
720 charge * copy.get_center_of_mass()
721 )
723 dipoles = (
724 _predict_dipole_batch(structures, model_filename, nep_command) * N
725 ) # dipole/atom, shape (3N+1, 3)
726 dipoles += corrections
728 d = dipoles[0, :]
729 d_forward = dipoles[1:].reshape(N, 3, 3)
730 gradient = (d_forward - d[None, None, :]) / displacement
732 elif method == 'central difference':
733 structures_forward = [] # will hold 3N structures
734 structures_backward = [] # will hold 3N structures
736 # Correct for the constant dipole, by adding charge * center of mass
737 corrections_forward = np.zeros((3 * N, 3))
738 corrections_backward = np.zeros((3 * N, 3))
740 for atom in range(N):
741 for cartesian in range(3):
742 # Forward displacements
743 copy_forward = structure.copy()
744 positions_forward = copy_forward.get_positions()
745 positions_forward[atom, cartesian] += displacement
746 copy_forward.set_positions(positions_forward)
747 copy_forward = _set_dummy_energy_forces(copy_forward)
748 structures_forward.append(copy_forward)
749 corrections_forward[atom * 3 + cartesian] = (
750 charge * copy_forward.get_center_of_mass()
751 )
753 # Backwards displacement
754 copy_backward = structure.copy()
755 positions_backward = copy_backward.get_positions()
756 positions_backward[atom, cartesian] -= displacement
757 copy_backward.set_positions(positions_backward)
758 copy_backward = _set_dummy_energy_forces(copy_backward)
759 structures_backward.append(copy_backward)
760 corrections_backward[atom * 3 + cartesian] = (
761 charge * copy_backward.get_center_of_mass()
762 )
764 structures = structures_forward + structures_backward
765 dipoles = (
766 _predict_dipole_batch(structures, model_filename, nep_command) * N
767 ) # dipole/atom, shape (6N, 3)
768 d_forward = dipoles[: 3 * N, :]
769 d_backward = dipoles[3 * N :, :]
771 d_forward += corrections_forward
772 d_backward += corrections_backward
774 d_forward = d_forward.reshape(N, 3, 3)
775 d_backward = d_backward.reshape(N, 3, 3)
777 gradient = (d_forward - d_backward) / (2 * displacement)
778 else:
779 raise ValueError(f'Invalid method {method} for calculating gradient')
780 return gradient
783def _set_dummy_energy_forces(structure: Atoms) -> Atoms:
784 """Sets the energies and forces of structure to zero.
786 Parameters
787 ----------
788 structure
789 Input structure
792 Returns
793 ------- Copy of structure, with SinglePointCalculator with zero energy and force.
794 """
795 from ase.calculators.singlepoint import SinglePointCalculator
797 copy = structure.copy()
798 N = len(copy)
799 energy = 0
801 forces = np.zeros((N, 3))
802 dummy = SinglePointCalculator(copy, **{'energy': energy, 'forces': forces})
803 copy.calc = dummy
804 return copy
807def _predict_dipole_batch(
808 structures: List[Atoms], model_filename: str, nep_command: str = 'nep'
809) -> np.ndarray:
810 """Predicts dipoles for a set of structures using the NEP executable
811 Note that the units are in (dipole units)/atom.
813 Parameters
814 ----------
815 structure
816 Input structures
817 model_filename
818 Path to NEP model in ``nep.txt`` format.
819 nep_command
820 Command for running the NEP executable. Defaults to ``'nep'``.
823 Returns
824 ------- Predicted dipoles, with shape (len(structures), 3).
825 """
826 import shutil
827 from os.path import join as join_path
828 from subprocess import run
829 from tempfile import TemporaryDirectory
831 from calorine.nep import read_model, write_nepfile, write_structures
833 with TemporaryDirectory() as directory:
834 shutil.copy2(model_filename, join_path(directory, 'nep.txt'))
835 model = read_model(model_filename)
837 parameters = dict(
838 prediction=1,
839 mode=1,
840 version=model.version,
841 n_max=[model.n_max_radial, model.n_max_angular],
842 cutoff=[model.radial_cutoff, model.angular_cutoff],
843 basis_size=[model.n_basis_radial, model.n_basis_radial],
844 l_max=[model.l_max_3b, model.l_max_4b, model.l_max_5b],
845 neuron=model.n_neuron,
846 type=[len(model.types), *model.types],
847 )
849 write_nepfile(parameters, directory)
850 file = join_path(directory, 'train.xyz')
851 write_structures(file, structures)
853 # Execute nep
854 completed = run([nep_command], cwd=directory, capture_output=True)
855 completed.check_returncode()
857 # Read results
858 dipoles = np.loadtxt(join_path(directory, 'dipole_train.out'))
859 if len(dipoles.shape) == 1:
860 dipoles = dipoles.reshape(1, -1)
861 return dipoles[:, :3]