Coverage for calorine/calculators/cpunep.py: 100%
212 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-10-08 12:52 +0000
1import contextlib
2import os
3from tempfile import TemporaryFile
4from typing import Any
6import numpy as np
7from ase import Atoms
8from ase.calculators.calculator import Calculator, all_changes, compare_atoms
9from ase.stress import full_3x3_to_voigt_6_stress
11import _nepy
12from calorine.nep.model import _get_nep_contents
13from calorine.nep.nep import _check_components_polarizability_gradient, \
14 _polarizability_gradient_to_3x3
15from calorine.nep.tensor_conventions import reduced6_to_full_3x3
16from ._dftd3 import DFTD3_KEYS, parse_dftd3
19class CPUNEP(Calculator):
20 """This class provides an ASE calculator for `nep_cpu`,
21 the in-memory CPU implementation of GPUMD.
23 Parameters
24 ----------
25 model_filename : str
26 Path to file in ``nep.txt`` format with model parameters
27 atoms : Atoms
28 Atoms to attach the calculator to
29 label : str
30 Label for this calclator
31 debug : bool, optional
32 Flag to toggle debug mode. Prints GPUMD output. Defaults to False.
33 dftd3 : dict, optional
34 Settings for the DFT-D3 dispersion correction with Becke-Johnson damping, added to the
35 energy, forces and virials of a potential model.
36 The keys are ``functional``, ``potential_cutoff`` and ``coordination_number_cutoff``,
37 the cutoffs in Å, for example ``dict(functional='pbe', potential_cutoff=12,
38 coordination_number_cutoff=6)``.
39 The coordination number cutoff must not exceed the potential cutoff.
40 An unsupported functional raises :class:`ValueError` when the calculator is created.
41 The energy is truncated at the potential cutoff without a switching function, as in
42 GPUMD.
43 Energy and forces change by a step as a pair crosses it, about 1e-4 eV per pair at
44 12 Å, and a molecular dynamics run does not conserve the energy exactly.
45 Place the potential cutoff between neighbor shells.
46 A cutoff on a shell includes or excludes the whole shell depending on rounding, which
47 makes the energy of a crystal depend on the supercell it is computed in.
48 Defaults to None.
50 Raises
51 ------
52 FileNotFoundError
53 Raises :class:`FileNotFoundError` if :attr:`model_filename` does not point to a valid file.
54 ValueError
55 Raises :class:`ValueError` atoms are not defined when trying to get energies and forces.
56 Example
57 -------
59 >>> # xdoctest: +SKIP
60 >>> calc = CPUNEP('nep.txt')
61 >>> atoms.calc = calc
62 >>> atoms.get_potential_energy()
63 """
65 base_implemented_properties = [
66 'energy',
67 'energies',
68 'forces',
69 'stress',
70 'stresses',
71 ]
72 debug = False
73 nepy = None
74 natoms = None
75 _nepy_atoms = None
77 def __init__(
78 self,
79 model_filename: str,
80 atoms: Atoms | None = None,
81 label: str | None = None,
82 debug: bool = False,
83 dftd3: dict[str, Any] | None = None,
84 ):
85 self.debug = debug
86 self._dftd3 = parse_dftd3(dftd3)
87 if self._dftd3 is not None and self._dftd3[2] > self._dftd3[1]:
88 raise ValueError(f'The coordination number cutoff ({self._dftd3[2]} Å) must not '
89 f'exceed the potential cutoff ({self._dftd3[1]} Å).')
91 if not os.path.exists(model_filename):
92 raise FileNotFoundError(f'{model_filename} does not exist.')
93 self.model_filename = str(model_filename)
95 # Get model type from first row in nep.txt
96 header, _ = _get_nep_contents(self.model_filename)
97 self.model_type = header['model_type']
98 self.supported_species = set(header['types'])
99 self.nep_version = header['version']
101 if self._dftd3 is not None:
102 if self.model_type != 'potential':
103 raise ValueError('The DFT-D3 dispersion correction is available for a potential '
104 f'model, and this model is of type {self.model_type}.')
105 # The extension holds the parameter table, and with it the supported names.
106 _nepy.check_dftd3_functional(self._dftd3[0])
108 # Set implemented properties -- not use class-level property
109 # to avoid leaking state between calculator instances.
110 self.implemented_properties = list(self.base_implemented_properties)
111 if 'charge' in self.model_type:
112 # Only available for charge models
113 self.implemented_properties.extend(
114 ['charges', 'born_effective_charges'])
115 elif self.model_type == 'dipole':
116 # Only available for dipole models
117 self.implemented_properties = ['dipole']
118 elif self.model_type == 'polarizability':
119 # Only available for polarizability models
120 self.implemented_properties = ['polarizability']
122 # Initialize atoms, results and nepy - note that this is also done in Calculator.__init__()
123 if atoms is not None:
124 self.set_atoms(atoms)
125 parameters = {'model_filename': model_filename}
126 Calculator.__init__(self, label=label, atoms=atoms, **parameters)
127 if atoms is not None:
128 self._setup_nepy()
130 def __str__(self) -> str:
131 def indent(s: str, i: int) -> str:
132 s = '\n'.join([i * ' ' + line for line in s.split('\n')])
133 return s
135 parameters = '\n'.join(
136 [f'{key}: {value}' for key, value in self.todict().items()]
137 )
138 parameters = indent(parameters, 4)
139 using_debug = '\nIn debug mode' if self.debug else ''
141 s = f'{self.__class__.__name__}\n{parameters}{using_debug}'
142 return s
144 def todict(self, skip_default: bool = True) -> dict[str, Any]:
145 """Return the parameters of the calculator, including the DFT-D3 settings in use."""
146 parameters = Calculator.todict(self, skip_default=skip_default)
147 parameters.pop('dftd3', None)
148 if self._dftd3 is not None:
149 parameters['dftd3'] = self.dftd3
150 return parameters
152 def set(self, **kwargs):
153 """Set parameters as defined by the ASE calculator interface.
155 Raises ``ValueError`` for ``dftd3`` settings other than those in use, since they are
156 fixed when the calculator is created.
157 """
158 if 'dftd3' in kwargs and parse_dftd3(kwargs.pop('dftd3')) != self._dftd3:
159 raise ValueError('The DFT-D3 settings are fixed when the calculator is created.')
160 return Calculator.set(self, **kwargs)
162 @property
163 def dftd3(self) -> dict[str, Any] | None:
164 """Settings for the DFT-D3 dispersion correction, fixed when the calculator is created.
165 None if no correction is applied.
166 """
167 if self._dftd3 is None:
168 return None
169 return dict(zip(DFTD3_KEYS, self._dftd3))
171 def _setup_nepy(self):
172 """
173 Creates an instance of the NEPY class and attaches it to the calculator object.
174 The output from `nep.cpp` is only written to STDOUT if debug == True
175 """
176 self._validate_atoms()
178 natoms = len(self.atoms)
179 self.natoms = natoms
180 c = self.atoms.get_cell(complete=True).flatten()
181 cell = [c[0], c[3], c[6], c[1], c[4], c[7], c[2], c[5], c[8]]
182 symbols = self.atoms.get_chemical_symbols()
183 positions = list(
184 self.atoms.get_positions().T.flatten()
185 ) # [x1, ..., xN, y1, ... yN,...]
186 masses = self.atoms.get_masses()
188 # Disable output from C++ code by default
189 if self.debug:
190 self.nepy = _nepy.NEPY(
191 self.model_filename, self.natoms, cell, symbols, positions, masses
192 )
193 else:
194 with TemporaryFile('w') as f:
195 with contextlib.redirect_stdout(f):
196 self.nepy = _nepy.NEPY(
197 self.model_filename,
198 self.natoms,
199 cell,
200 symbols,
201 positions,
202 masses,
203 )
204 self._nepy_atoms = self.atoms.copy()
206 def _check_species(self, atoms: Atoms):
207 """Checks that a structure only contains species covered by the model.
209 Parameters
210 ----------
211 atoms : Atoms
212 Structure to check
213 """
214 species_in_atoms_object = set(np.unique(atoms.get_chemical_symbols()))
215 if not species_in_atoms_object.issubset(self.supported_species):
216 raise ValueError('Structure contains species that are not supported by the NEP model.')
218 def _validate_atoms(self):
219 """Checks that the attached structure can be handled by the model."""
220 if self.atoms is None:
221 raise ValueError('Atoms must be defined when calculating properties.')
222 if self.atoms.cell.rank == 0:
223 raise ValueError('Atoms must have a defined cell.')
224 self._check_species(self.atoms)
226 def _sync_nepy(self):
227 """Brings the NEPY object in sync with the attached structure.
229 The structure that was last handed to NEPY is kept in
230 :attr:`_nepy_atoms`. The difference between that snapshot and
231 :attr:`atoms` determines what needs to be updated. Deriving the changes
232 here rather than from the :attr:`system_changes` argument of
233 :func:`calculate` also covers the gradient methods, which take no
234 structure argument, as well as structures that are modified in place.
236 A NEPY object is bound to a fixed number of atoms, since its setters
237 write into buffers that they cannot resize. It is therefore rebuilt
238 whenever the number of atoms changes.
239 """
240 self._validate_atoms()
242 if (
243 self.nepy is None
244 or self._nepy_atoms is None
245 or len(self._nepy_atoms) != len(self.atoms)
246 ):
247 self._setup_nepy()
248 self.results = {}
249 return
251 # `nep_cpu` always treats the cell as periodic and does not consume the
252 # initial charges or magnetic moments.
253 changes = compare_atoms(
254 self._nepy_atoms,
255 self.atoms,
256 tol=0,
257 excluded_properties={'pbc', 'initial_charges', 'initial_magmoms'},
258 )
259 # ASE does not track the masses, which NEPY uses for the center of mass
260 # correction of the dipole gradient.
261 masses_changed = not np.array_equal(
262 self._nepy_atoms.get_masses(), self.atoms.get_masses()
263 )
264 if not changes and not masses_changed:
265 return
267 if 'numbers' in changes:
268 self._update_symbols()
269 if 'numbers' in changes or masses_changed:
270 self._update_masses()
271 if 'positions' in changes:
272 self._update_positions()
273 if 'cell' in changes:
274 self._update_cell()
276 self._nepy_atoms = self.atoms.copy()
277 # The cached results belong to the structure that was replaced.
278 self.results = {}
280 def set_atoms(self, atoms: Atoms):
281 """Updates the Atoms object.
283 Parameters
284 ----------
285 atoms : Atoms
286 Atoms to attach the calculator to
287 """
288 self._check_species(atoms)
289 self.atoms = atoms
290 self.results = {}
291 self.nepy = None
292 self._nepy_atoms = None
294 def _update_symbols(self):
295 """Update atom symbols in NEPY."""
296 symbols = self.atoms.get_chemical_symbols()
297 self.nepy.set_symbols(symbols)
299 def _update_masses(self):
300 """Update atom masses in NEPY"""
301 masses = self.atoms.get_masses()
302 self.nepy.set_masses(masses)
304 def _update_cell(self):
305 """Update cell parameters in NEPY."""
306 c = self.atoms.get_cell(complete=True).flatten()
307 cell = [c[0], c[3], c[6], c[1], c[4], c[7], c[2], c[5], c[8]]
308 self.nepy.set_cell(cell)
310 def _update_positions(self):
311 """Update atom positions in NEPY."""
312 positions = list(
313 self.atoms.get_positions().T.flatten()
314 ) # [x1, ..., xN, y1, ... yN,...]
315 self.nepy.set_positions(positions)
317 def calculate(
318 self,
319 atoms: Atoms = None,
320 properties: list[str] = None,
321 system_changes: list[str] = all_changes,
322 ):
323 """Calculate energy, per atom energies, forces, stress and dipole.
325 Parameters
326 ----------
327 atoms : Atoms, optional
328 System for which to calculate properties, by default None
329 properties : list[str], optional
330 Properties to calculate, by default None
331 system_changes : list[str], optional
332 Changes to the system since last call, by default all_changes.
333 Accepted for compatibility with the ASE calculator interface and
334 not used, since :func:`_sync_nepy` derives the changes itself.
335 """
336 if properties is None:
337 properties = self.implemented_properties
339 Calculator.calculate(self, atoms, properties, system_changes)
341 self._sync_nepy()
342 natoms = len(self.atoms)
344 if 'dipole' in properties:
345 dipole = np.array(self.nepy.get_dipole())
346 self.results['dipole'] = dipole
347 elif 'polarizability' in properties:
348 # NEPY.get_polarizability() returns components in NEP_REDUCED6_ORDER
349 # (xx,yy,zz,xy,yz,zx); see `find_polarizability` in src/nepy/nep.cpp.
350 pol = np.array(self.nepy.get_polarizability())
351 polarizability = reduced6_to_full_3x3(pol)
352 self.results['polarizability'] = polarizability
353 elif 'descriptors' in properties:
354 descriptors = np.array(self.nepy.get_descriptors())
355 descriptors_per_atom = descriptors.reshape(-1, natoms).T
356 self.results['descriptors'] = descriptors_per_atom
357 else:
358 if 'charge' in self.model_type:
359 energies, forces, virials, charges, becs = \
360 self.nepy.get_potential_forces_virials_and_charges()
361 elif self._dftd3 is not None:
362 energies, forces, virials = \
363 self.nepy.get_potential_forces_and_virials_with_dftd3(*self._dftd3)
364 else:
365 energies, forces, virials = self.nepy.get_potential_forces_and_virials()
367 energies_per_atom = np.array(energies)
368 energy = energies_per_atom.sum()
369 forces_per_atom = np.array(forces).reshape(-1, natoms).T
370 # NEPY's per-atom virial is a raw, row-major full-3x3 tensor
371 # [xx,xy,xz,yx,yy,yz,zx,zy,zz]; documented in src/nepy/nep.h and
372 # matching the accumulation pattern in src/nepy/nep.cpp. Given
373 # that layout, `.reshape((3,3))` below is already correctly
374 # oriented, and ASE's own `full_3x3_to_voigt_6_stress` produces
375 # a correct ASE-Voigt-6 `stress` -- independently confirmed
376 # empirically via isolated shear strains, each producing a
377 # stress response concentrated at the expected component.
378 virials_per_atom = np.array(virials).reshape(-1, natoms).T
379 stresses_per_atom = virials_per_atom / self.atoms.get_volume()
380 stress = -(np.sum(virials_per_atom, axis=0) / self.atoms.get_volume()).reshape((3, 3))
381 stress = full_3x3_to_voigt_6_stress(stress)
383 self.results['energy'] = energy
384 self.results['energies'] = energies_per_atom
385 self.results['forces'] = forces_per_atom
386 self.results['stress'] = stress
387 self.results['stresses'] = stresses_per_atom
389 if 'charge' in self.model_type:
390 charges_per_atom = np.array(charges)
391 # Row-major full-3x3 per atom, same convention as the virial
392 # above (BEC is not symmetric, so no reduced-6 form applies);
393 # see src/nepy/nep.cpp. Matches GPUMD's own `dump_xyz`/
394 # `bec_*.out` convention (src/force/nep_charge.cu,
395 # src/measure/dump_xyz.cu -- no reindexing applied to BEC).
396 becs_per_atom = np.array(becs).reshape(-1, natoms).T
398 self.results['charges'] = charges_per_atom
399 self.results['born_effective_charges'] = becs_per_atom
401 def get_dipole_gradient(
402 self,
403 displacement: float = 0.01,
404 method: str = 'central difference',
405 charge: float = 1.0,
406 atoms: Atoms = None,
407 ):
408 """Calculates the dipole gradient using finite differences.
410 Parameters
411 ----------
412 displacement
413 Displacement in Å to use for finite differences. Defaults to 0.01 Å.
414 method
415 Method for computing gradient with finite differences.
416 One of 'forward difference' and 'central difference'.
417 Defaults to 'central difference'
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 atoms
423 System for which to compute the gradient. A copy is attached to the
424 calculator. Defaults to ``None``, in which case the structure that
425 is already attached is used.
427 Returns
428 -------
429 dipole gradient with shape `(N, 3, 3)` where ``N`` are the number of atoms.
430 """
431 if 'dipole' not in self.implemented_properties:
432 raise ValueError('Dipole gradients are only defined for dipole NEP models.')
434 if displacement <= 0:
435 raise ValueError('displacement must be > 0 Å')
437 implemented_methods = {
438 'forward difference': 0,
439 'central difference': 1,
440 'second order central difference': 2,
441 }
443 if method not in implemented_methods.keys():
444 raise ValueError(f'Invalid method {method} for calculating gradient')
446 if atoms is not None:
447 self.atoms = atoms.copy()
448 self._sync_nepy()
450 dipole_gradient = np.array(
451 self.nepy.get_dipole_gradient(
452 displacement, implemented_methods[method], charge
453 )
454 ).reshape(len(self.atoms), 3, 3)
455 return dipole_gradient
457 def get_polarizability(
458 self,
459 atoms: Atoms = None,
460 properties: list[str] = None,
461 system_changes: list[str] = all_changes,
462 ) -> np.ndarray:
463 """Calculates the polarizability tensor for the current structure.
464 The model must have been trained to predict the polarizability.
465 This is a wrapper function for :func:`calculate`.
467 Parameters
468 ----------
469 atoms : Atoms, optional
470 System for which to calculate properties, by default None
471 properties : list[str], optional
472 Properties to calculate, by default None
473 system_changes : list[str], optional
474 Changes to the system since last call, by default all_changes
476 Returns
477 -------
478 polarizability with shape ``(3, 3)``
479 """
480 if properties is None:
481 properties = self.implemented_properties
483 if 'polarizability' not in properties:
484 raise ValueError('Polarizability is only defined for polarizability NEP models.')
485 self.calculate(atoms, properties, system_changes)
486 return self.results['polarizability']
488 def get_polarizability_gradient(
489 self,
490 displacement: float = 0.01,
491 component: str | list[str] = 'full',
492 atoms: Atoms = None,
493 ) -> np.ndarray:
494 """Calculates the dipole gradient for a given structure using finite differences.
495 This function computes the derivatives using the second-order central difference
496 method with a C++ backend.
498 Parameters
499 ----------
500 displacement
501 Displacement in Å to use for finite differences. Defaults to ``0.01``.
502 component
503 Component or components of the polarizability tensor that the gradient
504 should be computed for.
505 The following components are available: `x`, `y`, `z`, `full`
506 Option ``full`` computes the derivative whilst moving the atoms in each Cartesian
507 direction, which yields a tensor of shape ``(N, 3, 3, 3)``,
508 where ``N`` is the number of atoms.
509 Multiple components may be specified.
510 Defaults to ``full``.
511 atoms
512 System for which to compute the gradient. A copy is attached to the
513 calculator. Defaults to ``None``, in which case the structure that
514 is already attached is used.
516 Returns
517 -------
518 polarizability gradient with shape ``(N, C, 3, 3)`` with ``C`` components chosen.
519 """
520 if 'polarizability' not in self.implemented_properties:
521 raise ValueError('Polarizability gradients are only defined'
522 ' for polarizability NEP models.')
524 if displacement <= 0:
525 raise ValueError('displacement must be > 0 Å')
527 if atoms is not None:
528 self.atoms = atoms.copy()
529 self._sync_nepy()
531 component_array = _check_components_polarizability_gradient(component)
533 pg = np.array(
534 self.nepy.get_polarizability_gradient(
535 displacement, component_array
536 )
537 ).reshape(len(self.atoms), 3, 6)
538 polarizability_gradient = _polarizability_gradient_to_3x3(pg)
539 return polarizability_gradient[:, component_array, :, :]
541 def get_descriptors(
542 self,
543 atoms: Atoms = None,
544 properties: list[str] = None,
545 system_changes: list[str] = all_changes,
546 ) -> np.ndarray:
547 """Calculates the descriptor tensor for the current structure.
548 This is a wrapper function for :func:`calculate`.
550 Parameters
551 ----------
552 atoms : Atoms, optional
553 System for which to calculate properties, by default None
554 properties : list[str], optional
555 Properties to calculate, by default None
556 system_changes : list[str], optional
557 Changes to the system since last call, by default all_changes
559 Returns
560 -------
561 descriptors with shape ``(number_of_atoms, descriptor_components)``
562 """
563 self.calculate(atoms, ['descriptors'], system_changes)
564 return self.results['descriptors']
566 def get_born_effective_charges(
567 self,
568 atoms: Atoms = None,
569 properties: list[str] = None,
570 system_changes: list[str] = all_changes,
571 ) -> np.ndarray:
572 """Calculates (if needed) and returns the Born effective charges.
573 Note that this requires a qNEP model.
575 Parameters
576 ----------
577 atoms
578 System for which to calculate properties, by default `None`.
579 properties
580 Properties to calculate, by default `None`.
581 system_changes
582 Changes to the system since last call, by default all_changes.
583 """
584 if 'born_effective_charges' not in self.implemented_properties:
585 raise ValueError(
586 'This model does not support the calculation of Born effective charges.')
587 self.calculate(atoms, properties, system_changes)
588 return self.results['born_effective_charges']