Coverage for calorine/nep/nep.py: 99%
337 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-23 12:50 +0000
1import contextlib
2import os
3import warnings
4from collections import Counter
5from pathlib import Path
6from typing import Dict, List, Optional, Tuple, Union
8import numpy as np
9from ase import Atoms
10from ase.calculators.singlepoint import SinglePointCalculator
12import _nepy
13from calorine.env import calorine_getenv
14from calorine.nep.model import Model, _get_nep_contents
15from calorine.nep.tensor_conventions import reduced6_to_full_3x3
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.
24 Parameters
25 ----------
26 structure
27 Atoms object representing the structure.
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)
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]]
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
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.
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.
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')
92 # Disable output from C++ code by default
93 if debug:
94 nepy = _nepy.NEPY(model_filename, natoms, cell, symbols, positions, masses)
95 else:
96 with open(os.devnull, 'w') as f:
97 with contextlib.redirect_stdout(f):
98 with contextlib.redirect_stderr(f):
99 nepy = _nepy.NEPY(
100 model_filename, natoms, cell, symbols, positions, masses
101 )
102 return nepy
105def set_default_cell(structure: Atoms, box_length: float = 100):
106 """Adds a cubic box to an Atoms object. Atoms object is edited in-place.
108 Parameters
109 ----------
110 structure
111 Structure to add box to
112 box_length
113 Cubic box side length in Å, by default 100
114 """
115 structure.set_cell([[box_length, 0, 0], [0, box_length, 0], [0, 0, box_length]])
116 structure.center()
119def get_descriptors(
120 structure: Atoms, model_filename: str, debug: bool = False
121) -> np.ndarray:
122 """Calculates the NEP descriptors for a given structure. A NEP model defined by a nep.txt
123 can additionally be provided to get the NEP3 model specific descriptors.
125 Parameters
126 ----------
127 structure
128 Input structure
129 model_filename
130 Path to NEP model in ``nep.txt`` format.
131 debug
132 Flag to toggle debug mode. Prints GPUMD output. Defaults to ``False``.
134 Returns
135 -------
136 Descriptors for the supplied structure, with shape (number_of_atoms, descriptor components)
137 """
138 local_structure = structure.copy()
139 natoms = len(local_structure)
140 cell, symbols, positions, masses = _get_atomic_properties(local_structure)
142 nepy = _setup_nepy(
143 model_filename, natoms, cell, symbols, positions, masses, debug
144 )
145 all_descriptors = nepy.get_descriptors()
146 descriptors_per_atom = np.array(all_descriptors).reshape(-1, natoms).T
147 return descriptors_per_atom
150def get_latent_space(
151 structure: Atoms, model_filename: Union[str, None] = None, debug: bool = False
152) -> np.ndarray:
153 """Calculates the latent space representation of a structure, i.e, the activiations in
154 the hidden layer. A NEP model defined by a `nep.txt` file needs to be provided.
156 Parameters
157 ----------
158 structure
159 Input structure
160 model_filename
161 Path to NEP model. Defaults to None.
162 debug
163 Flag to toggle debug mode. Prints GPUMD output. Defaults to False.
165 Returns
166 -------
167 Activation with shape `(natoms, N_neurons)`
168 """
169 if model_filename is None:
170 raise ValueError('Model is undefined')
171 local_structure = structure.copy()
172 natoms = len(local_structure)
173 cell, symbols, positions, masses = _get_atomic_properties(local_structure)
175 nepy = _setup_nepy(
176 model_filename, natoms, cell, symbols, positions, masses, debug
177 )
179 latent = nepy.get_latent_space()
180 latent = np.array(latent).reshape(-1, natoms).T
181 return latent
184def get_potential_forces_and_virials(
185 structure: Atoms, model_filename: Optional[str] = None, debug: bool = False
186) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
187 """Calculates the per-atom potential, forces and virials for a given structure.
188 A NEP model defined by a `nep.txt` file needs to be provided.
190 Parameters
191 ----------
192 structure
193 Input structure
194 model_filename
195 Path to NEP model. Defaults to None.
196 debug
197 Flag to toggle debug mode. Prints GPUMD output. Defaults to False.
199 Returns
200 -------
201 potential with shape `(natoms,)`
202 forces with shape `(natoms, 3)`
203 virials with shape `(natoms, 9)`
204 """
205 if model_filename is None:
206 raise ValueError('Model is undefined')
208 model_type = _get_nep_contents(model_filename)[0]['model_type']
209 if model_type != 'potential':
210 raise ValueError(
211 'A NEP model trained for predicting energies and forces must be used.'
212 )
214 local_structure = structure.copy()
215 natoms = len(local_structure)
216 cell, symbols, positions, masses = _get_atomic_properties(local_structure)
218 nepy = _setup_nepy(
219 model_filename, natoms, cell, symbols, positions, masses, debug
220 )
222 energies, forces, virials = nepy.get_potential_forces_and_virials()
223 forces_per_atom = np.array(forces).reshape(-1, natoms).T
224 # Raw, row-major full-3x3 per atom [xx,xy,xz,yx,yy,yz,zx,zy,zz];
225 # documented in src/nepy/nep.h and matching the accumulation pattern in
226 # src/nepy/nep.cpp.
227 virials_per_atom = np.array(virials).reshape(-1, natoms).T
228 return np.array(energies), forces_per_atom, virials_per_atom
231def get_polarizability(
232 structure: Atoms,
233 model_filename: Optional[str] = None,
234 debug: bool = False,
235) -> np.ndarray:
236 """Calculates the polarizability tensor for a given structure. A NEP model defined
237 by a ``nep.txt`` file needs to be provided. The model must be trained to predict the
238 polarizability.
240 Parameters
241 ----------
242 structure
243 Input structure
244 model_filename
245 Path to NEP model in ``nep.txt`` format. Defaults to ``None``.
246 debug
247 Flag to toggle debug mode. Prints GPUMD output. Defaults to ``False``.
249 Returns
250 -------
251 polarizability with shape ``(3, 3)``
252 """
253 if model_filename is None:
254 raise ValueError('Model is undefined')
256 model_type = _get_nep_contents(model_filename)[0]['model_type']
257 if model_type != 'polarizability':
258 raise ValueError(
259 'A NEP model trained for predicting polarizability must be used.'
260 )
262 local_structure = structure.copy()
263 natoms = len(local_structure)
264 cell, symbols, positions, masses = _get_atomic_properties(local_structure)
266 nepy = _setup_nepy(
267 model_filename, natoms, cell, symbols, positions, masses, debug
268 )
269 # Components are in NEP_REDUCED6_ORDER (xx,yy,zz,xy,yz,zx); see
270 # `find_polarizability` in src/nepy/nep.cpp.
271 pol = nepy.get_polarizability()
272 polarizability = reduced6_to_full_3x3(pol)
274 return polarizability
277def get_dipole(
278 structure: Atoms,
279 model_filename: Optional[str] = None,
280 debug: bool = False,
281) -> np.ndarray:
282 """Calculates the dipole for a given structure. A NEP model defined by a
283 ``nep.txt`` file needs to be provided.
285 Parameters
286 ----------
287 structure
288 Input structure
289 model_filename
290 Path to NEP model in ``nep.txt`` format. Defaults to ``None``.
291 debug
292 Flag to toggle debug mode. Prints GPUMD output. Defaults to ``False``.
294 Returns
295 -------
296 dipole with shape ``(3,)``
297 """
298 if model_filename is None:
299 raise ValueError('Model is undefined')
301 model_type = _get_nep_contents(model_filename)[0]['model_type']
302 if model_type != 'dipole':
303 raise ValueError('A NEP model trained for predicting dipoles must be used.')
305 local_structure = structure.copy()
306 natoms = len(local_structure)
307 cell, symbols, positions, masses = _get_atomic_properties(local_structure)
309 nepy = _setup_nepy(
310 model_filename, natoms, cell, symbols, positions, masses, debug
311 )
312 dipole = np.array(nepy.get_dipole())
314 return dipole
317def get_dipole_gradient(
318 structure: Atoms,
319 model_filename: Optional[str] = None,
320 backend: str = 'c++',
321 method: str = 'central difference',
322 displacement: float = 0.01,
323 charge: float = 1.0,
324 nep_command: Optional[str] = None,
325 debug: bool = False,
326) -> np.ndarray:
327 """Calculates the dipole gradient for a given structure using finite differences.
328 A NEP model defined by a `nep.txt` file needs to be provided.
330 Parameters
331 ----------
332 structure
333 Input structure
334 model_filename
335 Path to NEP model in ``nep.txt`` format. Defaults to ``None``.
336 backend
337 Backend to use for computing dipole gradient with finite differences.
338 One of ``'c++'`` (CPU), ``'python'`` (CPU) and ``'nep'`` (GPU).
339 Defaults to ``'c++'``.
340 method
341 Method for computing gradient with finite differences.
342 One of 'forward difference' and 'central difference'.
343 Defaults to 'central difference'
344 displacement
345 Displacement in Å to use for finite differences. Defaults to ``0.01``.
346 charge
347 System charge in units of the elemental charge.
348 Used for correcting the dipoles before computing the gradient.
349 Defaults to ``1.0``.
350 nep_command
351 Command for running the NEP executable.
352 Default: ``nep``, or the value of the ``CALORINE_NEP_COMMAND``
353 environment variable if set.
354 debug
355 Flag to toggle debug mode. Prints GPUMD output (if applicable). Defaults to ``False``.
358 Returns
359 -------
360 dipole gradient with shape ``(N, 3, 3)``
361 """
362 if model_filename is None:
363 raise ValueError('Model is undefined')
365 model_type = _get_nep_contents(model_filename)[0]['model_type']
366 if model_type != 'dipole':
367 raise ValueError('A NEP model trained for predicting dipoles must be used.')
369 local_structure = structure.copy()
371 if backend == 'c++':
372 dipole_gradient = _dipole_gradient_cpp(
373 local_structure,
374 model_filename,
375 displacement=displacement,
376 method=method,
377 charge=charge,
378 debug=debug,
379 )
380 elif backend == 'python':
381 dipole_gradient = _dipole_gradient_python(
382 local_structure,
383 model_filename,
384 displacement=displacement,
385 charge=charge,
386 method=method,
387 )
388 elif backend == 'nep':
389 dipole_gradient = _dipole_gradient_nep(
390 local_structure,
391 model_filename,
392 displacement=displacement,
393 method=method,
394 charge=charge,
395 nep_command=nep_command,
396 )
397 else:
398 raise ValueError(f'Invalid backend {backend}')
399 return dipole_gradient
402def _dipole_gradient_cpp(
403 structure: Atoms,
404 model_filename: str,
405 method: str = 'central difference',
406 displacement: float = 0.01,
407 charge: float = 1.0,
408 debug: bool = False,
409) -> np.ndarray:
410 """Calculates the dipole gradient with finite differences, using NEP_CPU.
412 Parameters
413 ----------
414 structure
415 Input structure
416 model_filename
417 Path to NEP model in ``nep.txt`` format.
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``.
429 Returns
430 -------
431 dipole gradient with shape ``(N, 3, 3)``
432 """
433 if displacement <= 0:
434 raise ValueError('Displacement must be > 0 Å')
436 implemented_methods = {
437 'forward difference': 0,
438 'central difference': 1,
439 'second order central difference': 2,
440 }
442 if method not in implemented_methods.keys():
443 raise ValueError(f'Invalid method {method} for calculating gradient')
445 local_structure = structure.copy()
446 natoms = len(local_structure)
447 cell, symbols, positions, masses = _get_atomic_properties(local_structure)
448 nepy = _setup_nepy(
449 model_filename, natoms, cell, symbols, positions, masses, debug
450 )
451 dipole_gradient = np.array(
452 nepy.get_dipole_gradient(displacement, implemented_methods[method], charge)
453 ).reshape(natoms, 3, 3)
454 return dipole_gradient
457def _dipole_gradient_python(
458 structure: Atoms,
459 model_filename: str,
460 method: str = 'central difference',
461 displacement: float = 0.01,
462 charge: float = 1.0,
463) -> np.ndarray:
464 """Calculates the dipole gradient with finite differences, using the Python and get_dipole().
466 Parameters
467 ----------
468 structure
469 Input structure
470 model_filename
471 Path to NEP model in ``nep.txt`` format.
472 method
473 Method for computing gradient with finite differences.
474 One of ``'forward difference'`` and ``'central difference'``.
475 Defaults to ``'central difference'``
476 displacement
477 Displacement in Å to use for finite differences. Defaults to ``0.01``.
478 charge
479 System charge in units of the elemental charge.
480 Used for correcting the dipoles before computing the gradient.
481 Defaults to ``1.0``.
483 Returns
484 -------
485 dipole gradient with shape ``(N, 3, 3)``
486 """
487 if displacement <= 0:
488 raise ValueError('Displacement must be > 0 Å')
490 N = len(structure)
491 if method == 'forward difference':
492 # Correct all dipole by the permanent dipole, charge * center of mass
493 d = (
494 get_dipole(structure, model_filename)
495 + charge * structure.get_center_of_mass()
496 )
497 d_forward = np.zeros((N, 3, 3))
498 for atom in range(N):
499 for cartesian in range(3):
500 copy = structure.copy()
501 positions = copy.get_positions()
502 positions[atom, cartesian] += displacement
503 copy.set_positions(positions)
504 d_forward[atom, cartesian, :] = (
505 get_dipole(copy, model_filename)
506 + charge * copy.get_center_of_mass()
507 )
508 gradient = (d_forward - d[None, None, :]) / displacement
510 elif method == 'central difference':
511 d_forward = np.zeros((N, 3, 3))
512 d_backward = np.zeros((N, 3, 3))
513 for atom in range(N):
514 for cartesian in range(3):
515 # Forward displacements
516 copy_forward = structure.copy()
517 positions_forward = copy_forward.get_positions()
518 positions_forward[atom, cartesian] += displacement
519 copy_forward.set_positions(positions_forward)
520 d_forward[atom, cartesian, :] = (
521 get_dipole(copy_forward, model_filename)
522 + charge * copy_forward.get_center_of_mass()
523 )
524 # Backwards displacement
525 copy_backward = structure.copy()
526 positions_backward = copy_backward.get_positions()
527 positions_backward[atom, cartesian] -= displacement
528 copy_backward.set_positions(positions_backward)
529 d_backward[atom, cartesian, :] = (
530 get_dipole(copy_backward, model_filename)
531 + charge * copy_backward.get_center_of_mass()
532 )
533 gradient = (d_forward - d_backward) / (2 * displacement)
534 elif method == 'second order central difference':
535 # Coefficients from
536 # https://en.wikipedia.org/wiki/Finite_difference_coefficient#Central_finite_difference
537 d_forward_one_h = np.zeros((N, 3, 3))
538 d_forward_two_h = np.zeros((N, 3, 3))
539 d_backward_one_h = np.zeros((N, 3, 3))
540 d_backward_two_h = np.zeros((N, 3, 3))
541 for atom in range(N):
542 for cartesian in range(3):
543 copy = structure.copy()
544 positions = copy.get_positions()
545 # Forward displacements
546 positions[atom, cartesian] += displacement # + h
547 copy.set_positions(positions)
548 d_forward_one_h[atom, cartesian, :] = (
549 get_dipole(copy, model_filename)
550 + charge * copy.get_center_of_mass()
551 )
552 positions[atom, cartesian] += displacement # + 2h total
553 copy.set_positions(positions)
554 d_forward_two_h[atom, cartesian, :] = (
555 get_dipole(copy, model_filename)
556 + charge * copy.get_center_of_mass()
557 )
558 # Backwards displacement
559 positions[atom, cartesian] -= 3 * displacement # 2h - 3h = -h
560 copy.set_positions(positions)
561 d_backward_one_h[atom, cartesian, :] = (
562 get_dipole(copy, model_filename)
563 + charge * copy.get_center_of_mass()
564 )
565 positions[atom, cartesian] -= displacement # - 2h total
566 copy.set_positions(positions)
567 d_backward_two_h[atom, cartesian, :] = (
568 get_dipole(copy, model_filename)
569 + charge * copy.get_center_of_mass()
570 )
571 c0 = -1.0 / 12.0
572 c1 = 2.0 / 3.0
573 gradient = (
574 c0 * d_forward_two_h
575 + c1 * d_forward_one_h
576 - c1 * d_backward_one_h
577 - c0 * d_backward_two_h
578 ) / displacement
579 else:
580 raise ValueError(f'Invalid method {method} for calculating gradient')
581 return gradient
584def _dipole_gradient_nep(
585 structure: Atoms,
586 model_filename: str,
587 method: str = 'central difference',
588 displacement: float = 0.01,
589 charge: float = 1.0,
590 nep_command: Optional[str] = None,
591) -> np.ndarray:
592 """Calculates the dipole gradient with finite differences, using the NEP executable.
594 Parameters
595 ----------
596 structure
597 Input structure
598 model_filename
599 Path to NEP model in ``nep.txt`` format.
600 method
601 Method for computing gradient with finite differences.
602 One of ``'forward difference'`` and ``'central difference'``.
603 Defaults to ``'central difference'``
604 displacement
605 Displacement in Å to use for finite differences. Defaults to 0.01 Å.
606 Note that results are possibly unreliable for displacemen < 0.01,
607 which might be due to rounding errors.
608 charge
609 System charge in units of the elemental charge.
610 Used for correcting the dipoles before computing the gradient.
611 Defaults to 1.0.
612 nep_command
613 Command for running the NEP executable.
614 Default: ``nep``, or the value of the ``CALORINE_NEP_COMMAND``
615 environment variable if set.
618 Returns
619 -------
620 dipole gradient with shape ``(N, 3, 3)``
621 """
622 if displacement <= 0:
623 raise ValueError('Displacement must be > 0 Å')
625 if displacement < 1e-2:
626 warnings.warn(
627 'Dipole gradients with nep are unstable for displacements < 0.01 Å.'
628 )
630 N = len(structure)
631 if method == 'forward difference':
632 structure = _set_dummy_energy_forces(structure)
633 structures = [structure] # will hold 3N+1 structures
634 # Correct for the constant dipole, by adding charge * center of mass
635 corrections = np.zeros((3 * N + 1, 3))
636 corrections[0] = charge * structure.get_center_of_mass()
637 for atom in range(N):
638 for cartesian in range(3):
639 copy = structure.copy()
640 positions = copy.get_positions()
641 positions[atom, cartesian] += displacement
642 copy.set_positions(positions)
643 copy = _set_dummy_energy_forces(copy)
644 structures.append(copy)
645 corrections[1 + atom * 3 + cartesian] = (
646 charge * copy.get_center_of_mass()
647 )
649 dipoles = (
650 _predict_dipole_batch(structures, model_filename, nep_command) * N
651 ) # dipole/atom, shape (3N+1, 3)
652 dipoles += corrections
654 d = dipoles[0, :]
655 d_forward = dipoles[1:].reshape(N, 3, 3)
656 gradient = (d_forward - d[None, None, :]) / displacement
658 elif method == 'central difference':
659 structures_forward = [] # will hold 3N structures
660 structures_backward = [] # will hold 3N structures
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))
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 copy_forward = _set_dummy_energy_forces(copy_forward)
674 structures_forward.append(copy_forward)
675 corrections_forward[atom * 3 + cartesian] = (
676 charge * copy_forward.get_center_of_mass()
677 )
679 # Backwards displacement
680 copy_backward = structure.copy()
681 positions_backward = copy_backward.get_positions()
682 positions_backward[atom, cartesian] -= displacement
683 copy_backward.set_positions(positions_backward)
684 copy_backward = _set_dummy_energy_forces(copy_backward)
685 structures_backward.append(copy_backward)
686 corrections_backward[atom * 3 + cartesian] = (
687 charge * copy_backward.get_center_of_mass()
688 )
690 structures = structures_forward + structures_backward
691 dipoles = (
692 _predict_dipole_batch(structures, model_filename, nep_command) * N
693 ) # dipole/atom, shape (6N, 3)
694 d_forward = dipoles[: 3 * N, :]
695 d_backward = dipoles[3 * N :, :]
697 d_forward += corrections_forward
698 d_backward += corrections_backward
700 d_forward = d_forward.reshape(N, 3, 3)
701 d_backward = d_backward.reshape(N, 3, 3)
703 gradient = (d_forward - d_backward) / (2 * displacement)
704 else:
705 raise ValueError(f'Invalid method {method} for calculating gradient')
706 return gradient
709def _set_dummy_energy_forces(structure: Atoms) -> Atoms:
710 """Sets the energies and forces of structure to zero.
712 Parameters
713 ----------
714 structure
715 Input structure
718 Returns
719 ------- Copy of structure, with SinglePointCalculator with zero energy and force.
720 """
721 from ase.calculators.singlepoint import SinglePointCalculator
723 copy = structure.copy()
724 N = len(copy)
725 energy = 0
727 forces = np.zeros((N, 3))
728 dummy = SinglePointCalculator(copy, **{'energy': energy, 'forces': forces})
729 copy.calc = dummy
730 return copy
733def _predict_dipole_batch(
734 structures: List[Atoms], model_filename: str, nep_command: Optional[str] = None
735) -> np.ndarray:
736 """Predicts dipoles for a set of structures using the NEP executable
737 Note that the units are in (dipole units)/atom.
739 Parameters
740 ----------
741 structure
742 Input structures
743 model_filename
744 Path to NEP model in ``nep.txt`` format.
745 nep_command
746 Command for running the NEP executable.
747 Default: ``nep``, or the value of the ``CALORINE_NEP_COMMAND``
748 environment variable if set.
751 Returns
752 ------- Predicted dipoles, with shape (len(structures), 3).
753 """
754 import shutil
755 from os.path import join as join_path
756 from subprocess import run
757 from tempfile import TemporaryDirectory
759 from calorine.nep import read_model, write_nepfile, write_structures
761 with TemporaryDirectory() as directory:
762 shutil.copy2(model_filename, join_path(directory, 'nep.txt'))
763 model = read_model(model_filename)
765 parameters = dict(
766 prediction=1,
767 mode=1,
768 version=model.version,
769 n_max=[model.n_max_radial, model.n_max_angular],
770 type=[len(model.types), *model.types],
771 cutoff=[model.radial_cutoff, model.angular_cutoff],
772 basis_size=[model.n_basis_radial, model.n_basis_radial],
773 l_max=[model.l_max_3b, model.l_max_4b, model.l_max_5b],
774 neuron=model.n_neuron,
775 )
777 write_nepfile(parameters, directory)
778 file = join_path(directory, 'train.xyz')
779 with warnings.catch_warnings():
780 # this function only ever predicts dipoles; the input stress of the
781 # structures is irrelevant and usually absent, so the resulting
782 # warning is expected
783 warnings.filterwarnings(
784 'ignore', message='Failed to retrieve stresses for structure',
785 category=UserWarning)
786 write_structures(file, structures)
788 # Execute nep
789 nep_command = nep_command or calorine_getenv('NEP_COMMAND')
790 completed = run([nep_command], cwd=directory, capture_output=True)
791 completed.check_returncode()
793 # Read results
794 dipoles = np.loadtxt(join_path(directory, 'dipole_train.out'))
795 if len(dipoles.shape) == 1:
796 dipoles = dipoles.reshape(1, -1)
797 return dipoles[:, :3]
800def _check_components_polarizability_gradient(component: Union[str, List[str]]) -> List[int]:
801 """
802 Verifies that the selected components are ok.
803 """
804 allowed_components = {
805 'x': 0,
806 'y': 1,
807 'z': 2,
808 }
810 # Check if chosen components are ok
811 if component == 'full':
812 components_to_compute = [0, 1, 2]
813 else:
814 components_list = [component] if isinstance(component, str) else component
815 components_to_compute = []
816 for c in components_list:
817 if c in allowed_components.keys():
818 components_to_compute.append(allowed_components[c])
819 elif c == 'full':
820 raise ValueError('Write ``component="full"`` to get all components.')
821 else:
822 raise ValueError(f'Invalid component {c}')
823 assert len(components_list) == len(components_to_compute), \
824 'Number of components to compute does not match.'
825 return components_to_compute
828def get_polarizability_gradient(
829 structure: Atoms,
830 model_filename: Optional[str] = None,
831 displacement: float = 0.01,
832 component: Union[str, List[str]] = 'full',
833 debug: bool = False,
834) -> np.ndarray:
835 """Calculates the dipole gradient for a given structure using finite differences.
836 A NEP model defined by a ``nep.txt`` file needs to be provided.
837 This function computes the derivatives using the second-order central difference
838 method with a C++ backend.
840 Parameters
841 ----------
842 structure
843 Input structure.
844 model_filename
845 Path to NEP model in ``nep.txt`` format. Defaults to ``None``.
846 displacement
847 Displacement in Å to use for finite differences. Defaults to ``0.01``.
848 component
849 Component or components of the polarizability tensor that the gradient
850 should be computed for.
851 The following components are available: ``x`, ``y``, ``z``, ``full``.
852 Option ``full`` computes the derivative whilst moving the atoms in each Cartesian
853 direction, which yields a tensor of shape ``(N, 3, 6)``.
854 Multiple components may be specified.
855 Defaults to ``full``.
856 debug
857 Flag to toggle debug mode. Prints GPUMD output (if applicable). Defaults to ``False``.
860 Returns
861 -------
862 polarizability gradient with shape ``(N, C, 6)`` where ``C``
863 is the number of components chosen.
864 """
865 if model_filename is None:
866 raise ValueError('Model is undefined')
868 model_type = _get_nep_contents(model_filename)[0]['model_type']
869 if model_type != 'polarizability':
870 raise ValueError('A NEP model trained for predicting polarizability must be used.')
872 components_to_compute = _check_components_polarizability_gradient(component)
874 local_structure = structure.copy()
875 polarizability_gradient = _polarizability_gradient_cpp(
876 local_structure,
877 model_filename,
878 displacement=displacement,
879 components=components_to_compute,
880 debug=debug,
881 )
882 return polarizability_gradient
885def _polarizability_gradient_to_3x3(pg):
886 """
887 Converts a polarizability gradient tensor with
888 shape (Natoms, 3, 6) to (Natoms, 3, 3, 3).
889 The 6 items in the polarizability gradient are in NEP_REDUCED6_ORDER
890 (xx, yy, zz, xy, yz, zx), the same convention as `get_polarizability`:
891 src/nepy/nep.cpp computes the gradient via finite differences of
892 repeated `find_polarizability` calls, so it inherits that function's
893 component order directly.
894 """
895 return reduced6_to_full_3x3(pg)
898def _polarizability_gradient_cpp(
899 structure: Atoms,
900 model_filename: str,
901 displacement: float,
902 components: List[int],
903 debug: bool = False,
904) -> np.ndarray:
905 """Calculates the polarizability gradient with finite differences, using NEP_CPU.
907 Parameters
908 ----------
909 structure
910 Input structure.
911 model_filename
912 Path to NEP model in ``nep.txt`` format.
913 displacement
914 Displacement in Å to use for finite differences. Defaults to ``0.01``.
915 components
916 List of components to compute. Integer values from 0 to 2 that corresponds
917 to indices in the ordered list [x, y, z]. Index 1 corresponds to the
918 derivative with regards to only the y positions of the atoms, and so forth.
920 Returns
921 -------
922 dipole gradient with shape ``(N, len(components), 3, 3)``
923 """
924 if displacement <= 0:
925 raise ValueError('Displacement must be > 0 Å')
926 # TODO possibly use components later to only move atoms in one cartesian direction
927 local_structure = structure.copy()
928 natoms = len(local_structure)
929 cell, symbols, positions, masses = _get_atomic_properties(local_structure)
930 nepy = _setup_nepy(
931 model_filename, natoms, cell, symbols, positions, masses, debug
932 )
933 pg = np.array(
934 nepy.get_polarizability_gradient(displacement, components)
935 ).reshape(natoms, 3, 6)
936 # Convert to 3x3
937 polarizability_gradient = _polarizability_gradient_to_3x3(pg)
938 return polarizability_gradient[:, components, :, :] # Only return the relevant components
941def determine_energy_offsets(
942 structures: List[Atoms],
943 reference: Union[str, Model, None] = None,
944 command: Optional[str] = None,
945) -> Dict[str, float]:
946 """Determines a per-species energy offset (reference energy shift) ``mu_i`` for a set
947 of structures, by fitting ``E_target - E_reference = sum_i n_i * mu_i`` via ordinary
948 least squares, where ``n_i`` is the number of atoms of species ``i`` in a structure.
949 The fit never includes a constant term, since a structure-size-independent constant
950 would distort per-atom energy comparisons between structures of different sizes.
952 If ``reference`` is ``None``, the offsets are fitted directly against the target
953 energies of ``structures`` themselves, which is useful, for example, to obtain
954 isolated-atom-style reference energies before training a model from scratch. If
955 ``reference`` is given, either as a path to a NEP model in ``nep.txt`` format or as a
956 :class:`Model <calorine.nep.model.Model>` object, the offsets are instead fitted
957 against the residual between the target energies and the energies predicted by
958 ``reference``, e.g., to correct for an energy-scale mismatch when fine-tuning a
959 reference model such as NEP89 on data from a different reference method.
961 Parameters
962 ----------
963 structures
964 Structures with target energies attached, e.g., via an ASE calculator.
965 reference
966 Path to a NEP model in ``nep.txt`` format, or a :class:`Model
967 <calorine.nep.model.Model>` object, to fit the offsets against. If ``None``,
968 the offsets are fitted directly against the target energies of ``structures``.
969 command
970 Command used to invoke the ``nep`` executable when ``reference`` is given.
971 Default: ``nep``, or the value of the ``CALORINE_NEP_COMMAND`` environment
972 variable if set.
974 Returns
975 -------
976 Per-species energy offsets in eV, keyed by chemical symbol.
977 """
978 target_energies = np.array([atoms.get_potential_energy() for atoms in structures])
979 if reference is None:
980 reference_energies = np.zeros(len(structures))
981 else:
982 from calorine.tools import batch_predict_properties
983 predicted_structures = batch_predict_properties(structures, reference, command=command)
984 reference_energies = np.array(
985 [atoms.get_potential_energy() for atoms in predicted_structures])
987 species = sorted({s for atoms in structures for s in atoms.get_chemical_symbols()})
988 X = np.zeros((len(structures), len(species)))
989 for i, atoms in enumerate(structures):
990 counts = Counter(atoms.get_chemical_symbols())
991 for j, s in enumerate(species):
992 X[i, j] = counts.get(s, 0)
993 y = target_energies - reference_energies
994 coeffs, *_ = np.linalg.lstsq(X, y, rcond=None)
995 return dict(zip(species, coeffs))
998def apply_energy_offsets(structures: List[Atoms], offsets: Dict[str, float]) -> List[Atoms]:
999 """Applies a per-species energy offset, as determined by
1000 :func:`determine_energy_offsets <calorine.nep.determine_energy_offsets>`, to the
1001 target energies of a set of structures. Forces and stresses, if present, are passed
1002 through unchanged.
1004 Parameters
1005 ----------
1006 structures
1007 Structures with target energies attached, e.g., via an ASE calculator.
1008 offsets
1009 Per-species energy offsets in eV, keyed by chemical symbol, as returned by
1010 :func:`determine_energy_offsets <calorine.nep.determine_energy_offsets>`.
1012 Returns
1013 -------
1014 A new list of :class:`Atoms <ase.Atoms>` objects, in the same order as
1015 :attr:`structures`, with corrected target energies. The input :attr:`structures`
1016 are not modified.
1017 """
1018 missing_species = {
1019 s for atoms in structures for s in atoms.get_chemical_symbols()} - offsets.keys()
1020 if missing_species:
1021 raise ValueError(
1022 f'offsets is missing the following species found in structures: '
1023 f'{sorted(missing_species)}')
1025 corrected = []
1026 for atoms in structures:
1027 new_atoms = atoms.copy()
1028 results = atoms.calc.results if atoms.calc is not None else {}
1029 properties = {}
1030 if 'energy' in results: 1030 ↛ 1034line 1030 didn't jump to line 1034 because the condition on line 1030 was always true
1031 counts = Counter(atoms.get_chemical_symbols())
1032 shift = sum(counts.get(s, 0) * offsets[s] for s in counts)
1033 properties['energy'] = atoms.get_potential_energy() - shift
1034 if 'forces' in results:
1035 properties['forces'] = atoms.get_forces()
1036 if 'stress' in results:
1037 properties['stress'] = atoms.get_stress()
1038 new_atoms.calc = SinglePointCalculator(new_atoms, **properties)
1039 corrected.append(new_atoms)
1040 return corrected