Coverage for calorine/calculators/cpunep.py: 100%
188 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-14 16:36 +0000
1from __future__ import annotations
3import contextlib
4import os
5from tempfile import TemporaryFile
6from typing import List, Union
8import numpy as np
9from ase import Atoms
10from ase.calculators.calculator import Calculator, all_changes, compare_atoms
11from ase.stress import full_3x3_to_voigt_6_stress
13import _nepy
14from calorine.nep.model import _get_nep_contents
15from calorine.nep.nep import _check_components_polarizability_gradient, \
16 _polarizability_gradient_to_3x3
17from calorine.nep.tensor_conventions import reduced6_to_full_3x3
20class CPUNEP(Calculator):
21 """This class provides an ASE calculator for `nep_cpu`,
22 the in-memory CPU implementation of GPUMD.
24 Parameters
25 ----------
26 model_filename : str
27 Path to file in ``nep.txt`` format with model parameters
28 atoms : Atoms
29 Atoms to attach the calculator to
30 label : str
31 Label for this calclator
32 debug : bool, optional
33 Flag to toggle debug mode. Prints GPUMD output. Defaults to False.
35 Raises
36 ------
37 FileNotFoundError
38 Raises :class:`FileNotFoundError` if :attr:`model_filename` does not point to a valid file.
39 ValueError
40 Raises :class:`ValueError` atoms are not defined when trying to get energies and forces.
41 Example
42 -------
44 >>> calc = CPUNEP('nep.txt')
45 >>> atoms.calc = calc
46 >>> atoms.get_potential_energy()
47 """
49 base_implemented_properties = [
50 'energy',
51 'energies',
52 'forces',
53 'stress',
54 'stresses',
55 ]
56 debug = False
57 nepy = None
58 natoms = None
59 _nepy_atoms = None
61 def __init__(
62 self,
63 model_filename: str,
64 atoms: Atoms | None = None,
65 label: str | None = None,
66 debug: bool = False,
67 ):
68 self.debug = debug
70 if not os.path.exists(model_filename):
71 raise FileNotFoundError(f'{model_filename} does not exist.')
72 self.model_filename = str(model_filename)
74 # Get model type from first row in nep.txt
75 header, _ = _get_nep_contents(self.model_filename)
76 self.model_type = header['model_type']
77 self.supported_species = set(header['types'])
78 self.nep_version = header['version']
80 # Set implemented properties -- not use class-level property
81 # to avoid leaking state between calculator instances.
82 self.implemented_properties = list(self.base_implemented_properties)
83 if 'charge' in self.model_type:
84 # Only available for charge models
85 self.implemented_properties.extend(
86 ['charges', 'born_effective_charges'])
87 elif self.model_type == 'dipole':
88 # Only available for dipole models
89 self.implemented_properties = ['dipole']
90 elif self.model_type == 'polarizability':
91 # Only available for polarizability models
92 self.implemented_properties = ['polarizability']
94 # Initialize atoms, results and nepy - note that this is also done in Calculator.__init__()
95 if atoms is not None:
96 self.set_atoms(atoms)
97 parameters = {'model_filename': model_filename}
98 Calculator.__init__(self, label=label, atoms=atoms, **parameters)
99 if atoms is not None:
100 self._setup_nepy()
102 def __str__(self) -> str:
103 def indent(s: str, i: int) -> str:
104 s = '\n'.join([i * ' ' + line for line in s.split('\n')])
105 return s
107 parameters = '\n'.join(
108 [f'{key}: {value}' for key, value in self.parameters.items()]
109 )
110 parameters = indent(parameters, 4)
111 using_debug = '\nIn debug mode' if self.debug else ''
113 s = f'{self.__class__.__name__}\n{parameters}{using_debug}'
114 return s
116 def _setup_nepy(self):
117 """
118 Creates an instance of the NEPY class and attaches it to the calculator object.
119 The output from `nep.cpp` is only written to STDOUT if debug == True
120 """
121 self._validate_atoms()
123 natoms = len(self.atoms)
124 self.natoms = natoms
125 c = self.atoms.get_cell(complete=True).flatten()
126 cell = [c[0], c[3], c[6], c[1], c[4], c[7], c[2], c[5], c[8]]
127 symbols = self.atoms.get_chemical_symbols()
128 positions = list(
129 self.atoms.get_positions().T.flatten()
130 ) # [x1, ..., xN, y1, ... yN,...]
131 masses = self.atoms.get_masses()
133 # Disable output from C++ code by default
134 if self.debug:
135 self.nepy = _nepy.NEPY(
136 self.model_filename, self.natoms, cell, symbols, positions, masses
137 )
138 else:
139 with TemporaryFile('w') as f:
140 with contextlib.redirect_stdout(f):
141 self.nepy = _nepy.NEPY(
142 self.model_filename,
143 self.natoms,
144 cell,
145 symbols,
146 positions,
147 masses,
148 )
149 self._nepy_atoms = self.atoms.copy()
151 def _check_species(self, atoms: Atoms):
152 """Checks that a structure only contains species covered by the model.
154 Parameters
155 ----------
156 atoms : Atoms
157 Structure to check
158 """
159 species_in_atoms_object = set(np.unique(atoms.get_chemical_symbols()))
160 if not species_in_atoms_object.issubset(self.supported_species):
161 raise ValueError('Structure contains species that are not supported by the NEP model.')
163 def _validate_atoms(self):
164 """Checks that the attached structure can be handled by the model."""
165 if self.atoms is None:
166 raise ValueError('Atoms must be defined when calculating properties.')
167 if self.atoms.cell.rank == 0:
168 raise ValueError('Atoms must have a defined cell.')
169 self._check_species(self.atoms)
171 def _sync_nepy(self):
172 """Brings the NEPY object in sync with the attached structure.
174 The structure that was last handed to NEPY is kept in
175 :attr:`_nepy_atoms`. The difference between that snapshot and
176 :attr:`atoms` determines what needs to be updated. Deriving the changes
177 here rather than from the :attr:`system_changes` argument of
178 :func:`calculate` also covers the gradient methods, which take no
179 structure argument, as well as structures that are modified in place.
181 A NEPY object is bound to a fixed number of atoms, since its setters
182 write into buffers that they cannot resize. It is therefore rebuilt
183 whenever the number of atoms changes.
184 """
185 self._validate_atoms()
187 if (
188 self.nepy is None
189 or self._nepy_atoms is None
190 or len(self._nepy_atoms) != len(self.atoms)
191 ):
192 self._setup_nepy()
193 self.results = {}
194 return
196 # `nep_cpu` always treats the cell as periodic and does not consume the
197 # initial charges or magnetic moments.
198 changes = compare_atoms(
199 self._nepy_atoms,
200 self.atoms,
201 tol=0,
202 excluded_properties={'pbc', 'initial_charges', 'initial_magmoms'},
203 )
204 # ASE does not track the masses, which NEPY uses for the center of mass
205 # correction of the dipole gradient.
206 masses_changed = not np.array_equal(
207 self._nepy_atoms.get_masses(), self.atoms.get_masses()
208 )
209 if not changes and not masses_changed:
210 return
212 if 'numbers' in changes:
213 self._update_symbols()
214 if 'numbers' in changes or masses_changed:
215 self._update_masses()
216 if 'positions' in changes:
217 self._update_positions()
218 if 'cell' in changes:
219 self._update_cell()
221 self._nepy_atoms = self.atoms.copy()
222 # The cached results belong to the structure that was replaced.
223 self.results = {}
225 def set_atoms(self, atoms: Atoms):
226 """Updates the Atoms object.
228 Parameters
229 ----------
230 atoms : Atoms
231 Atoms to attach the calculator to
232 """
233 self._check_species(atoms)
234 self.atoms = atoms
235 self.results = {}
236 self.nepy = None
237 self._nepy_atoms = None
239 def _update_symbols(self):
240 """Update atom symbols in NEPY."""
241 symbols = self.atoms.get_chemical_symbols()
242 self.nepy.set_symbols(symbols)
244 def _update_masses(self):
245 """Update atom masses in NEPY"""
246 masses = self.atoms.get_masses()
247 self.nepy.set_masses(masses)
249 def _update_cell(self):
250 """Update cell parameters in NEPY."""
251 c = self.atoms.get_cell(complete=True).flatten()
252 cell = [c[0], c[3], c[6], c[1], c[4], c[7], c[2], c[5], c[8]]
253 self.nepy.set_cell(cell)
255 def _update_positions(self):
256 """Update atom positions in NEPY."""
257 positions = list(
258 self.atoms.get_positions().T.flatten()
259 ) # [x1, ..., xN, y1, ... yN,...]
260 self.nepy.set_positions(positions)
262 def calculate(
263 self,
264 atoms: Atoms = None,
265 properties: List[str] = None,
266 system_changes: List[str] = all_changes,
267 ):
268 """Calculate energy, per atom energies, forces, stress and dipole.
270 Parameters
271 ----------
272 atoms : Atoms, optional
273 System for which to calculate properties, by default None
274 properties : List[str], optional
275 Properties to calculate, by default None
276 system_changes : List[str], optional
277 Changes to the system since last call, by default all_changes.
278 Accepted for compatibility with the ASE calculator interface and
279 not used, since :func:`_sync_nepy` derives the changes itself.
280 """
281 if properties is None:
282 properties = self.implemented_properties
284 Calculator.calculate(self, atoms, properties, system_changes)
286 self._sync_nepy()
287 natoms = len(self.atoms)
289 if 'dipole' in properties:
290 dipole = np.array(self.nepy.get_dipole())
291 self.results['dipole'] = dipole
292 elif 'polarizability' in properties:
293 # NEPY.get_polarizability() returns components in NEP_REDUCED6_ORDER
294 # (xx,yy,zz,xy,yz,zx); see `find_polarizability` in src/nepy/nep.cpp.
295 pol = np.array(self.nepy.get_polarizability())
296 polarizability = reduced6_to_full_3x3(pol)
297 self.results['polarizability'] = polarizability
298 elif 'descriptors' in properties:
299 descriptors = np.array(self.nepy.get_descriptors())
300 descriptors_per_atom = descriptors.reshape(-1, natoms).T
301 self.results['descriptors'] = descriptors_per_atom
302 else:
303 if 'charge' in self.model_type:
304 energies, forces, virials, charges, becs = \
305 self.nepy.get_potential_forces_virials_and_charges()
306 else:
307 energies, forces, virials = self.nepy.get_potential_forces_and_virials()
309 energies_per_atom = np.array(energies)
310 energy = energies_per_atom.sum()
311 forces_per_atom = np.array(forces).reshape(-1, natoms).T
312 # NEPY's per-atom virial is a raw, row-major full-3x3 tensor
313 # [xx,xy,xz,yx,yy,yz,zx,zy,zz]; documented in src/nepy/nep.h and
314 # matching the accumulation pattern in src/nepy/nep.cpp. Given
315 # that layout, `.reshape((3,3))` below is already correctly
316 # oriented, and ASE's own `full_3x3_to_voigt_6_stress` produces
317 # a correct ASE-Voigt-6 `stress` -- independently confirmed
318 # empirically via isolated shear strains, each producing a
319 # stress response concentrated at the expected component.
320 virials_per_atom = np.array(virials).reshape(-1, natoms).T
321 stresses_per_atom = virials_per_atom / self.atoms.get_volume()
322 stress = -(np.sum(virials_per_atom, axis=0) / self.atoms.get_volume()).reshape((3, 3))
323 stress = full_3x3_to_voigt_6_stress(stress)
325 self.results['energy'] = energy
326 self.results['energies'] = energies_per_atom
327 self.results['forces'] = forces_per_atom
328 self.results['stress'] = stress
329 self.results['stresses'] = stresses_per_atom
331 if 'charge' in self.model_type:
332 charges_per_atom = np.array(charges)
333 # Row-major full-3x3 per atom, same convention as the virial
334 # above (BEC is not symmetric, so no reduced-6 form applies);
335 # see src/nepy/nep.cpp. Matches GPUMD's own `dump_xyz`/
336 # `bec_*.out` convention (src/force/nep_charge.cu,
337 # src/measure/dump_xyz.cu -- no reindexing applied to BEC).
338 becs_per_atom = np.array(becs).reshape(-1, natoms).T
340 self.results['charges'] = charges_per_atom
341 self.results['born_effective_charges'] = becs_per_atom
343 def get_dipole_gradient(
344 self,
345 displacement: float = 0.01,
346 method: str = 'central difference',
347 charge: float = 1.0,
348 atoms: Atoms = None,
349 ):
350 """Calculates the dipole gradient using finite differences.
352 Parameters
353 ----------
354 displacement
355 Displacement in Å to use for finite differences. Defaults to 0.01 Å.
356 method
357 Method for computing gradient with finite differences.
358 One of 'forward difference' and 'central difference'.
359 Defaults to 'central difference'
360 charge
361 System charge in units of the elemental charge.
362 Used for correcting the dipoles before computing the gradient.
363 Defaults to 1.0.
364 atoms
365 System for which to compute the gradient. A copy is attached to the
366 calculator. Defaults to ``None``, in which case the structure that
367 is already attached is used.
369 Returns
370 -------
371 dipole gradient with shape `(N, 3, 3)` where ``N`` are the number of atoms.
372 """
373 if 'dipole' not in self.implemented_properties:
374 raise ValueError('Dipole gradients are only defined for dipole NEP models.')
376 if displacement <= 0:
377 raise ValueError('displacement must be > 0 Å')
379 implemented_methods = {
380 'forward difference': 0,
381 'central difference': 1,
382 'second order central difference': 2,
383 }
385 if method not in implemented_methods.keys():
386 raise ValueError(f'Invalid method {method} for calculating gradient')
388 if atoms is not None:
389 self.atoms = atoms.copy()
390 self._sync_nepy()
392 dipole_gradient = np.array(
393 self.nepy.get_dipole_gradient(
394 displacement, implemented_methods[method], charge
395 )
396 ).reshape(len(self.atoms), 3, 3)
397 return dipole_gradient
399 def get_polarizability(
400 self,
401 atoms: Atoms = None,
402 properties: List[str] = None,
403 system_changes: List[str] = all_changes,
404 ) -> np.ndarray:
405 """Calculates the polarizability tensor for the current structure.
406 The model must have been trained to predict the polarizability.
407 This is a wrapper function for :func:`calculate`.
409 Parameters
410 ----------
411 atoms : Atoms, optional
412 System for which to calculate properties, by default None
413 properties : List[str], optional
414 Properties to calculate, by default None
415 system_changes : List[str], optional
416 Changes to the system since last call, by default all_changes
418 Returns
419 -------
420 polarizability with shape ``(3, 3)``
421 """
422 if properties is None:
423 properties = self.implemented_properties
425 if 'polarizability' not in properties:
426 raise ValueError('Polarizability is only defined for polarizability NEP models.')
427 self.calculate(atoms, properties, system_changes)
428 return self.results['polarizability']
430 def get_polarizability_gradient(
431 self,
432 displacement: float = 0.01,
433 component: Union[str, List[str]] = 'full',
434 atoms: Atoms = None,
435 ) -> np.ndarray:
436 """Calculates the dipole gradient for a given structure using finite differences.
437 This function computes the derivatives using the second-order central difference
438 method with a C++ backend.
440 Parameters
441 ----------
442 displacement
443 Displacement in Å to use for finite differences. Defaults to ``0.01``.
444 component
445 Component or components of the polarizability tensor that the gradient
446 should be computed for.
447 The following components are available: `x`, `y`, `z`, `full`
448 Option ``full`` computes the derivative whilst moving the atoms in each Cartesian
449 direction, which yields a tensor of shape ``(N, 3, 3, 3)``,
450 where ``N`` is the number of atoms.
451 Multiple components may be specified.
452 Defaults to ``full``.
453 atoms
454 System for which to compute the gradient. A copy is attached to the
455 calculator. Defaults to ``None``, in which case the structure that
456 is already attached is used.
458 Returns
459 -------
460 polarizability gradient with shape ``(N, C, 3, 3)`` with ``C`` components chosen.
461 """
462 if 'polarizability' not in self.implemented_properties:
463 raise ValueError('Polarizability gradients are only defined'
464 ' for polarizability NEP models.')
466 if displacement <= 0:
467 raise ValueError('displacement must be > 0 Å')
469 if atoms is not None:
470 self.atoms = atoms.copy()
471 self._sync_nepy()
473 component_array = _check_components_polarizability_gradient(component)
475 pg = np.array(
476 self.nepy.get_polarizability_gradient(
477 displacement, component_array
478 )
479 ).reshape(len(self.atoms), 3, 6)
480 polarizability_gradient = _polarizability_gradient_to_3x3(pg)
481 return polarizability_gradient[:, component_array, :, :]
483 def get_descriptors(
484 self,
485 atoms: Atoms = None,
486 properties: List[str] = None,
487 system_changes: List[str] = all_changes,
488 ) -> np.ndarray:
489 """Calculates the descriptor tensor for the current structure.
490 This is a wrapper function for :func:`calculate`.
492 Parameters
493 ----------
494 atoms : Atoms, optional
495 System for which to calculate properties, by default None
496 properties : List[str], optional
497 Properties to calculate, by default None
498 system_changes : List[str], optional
499 Changes to the system since last call, by default all_changes
501 Returns
502 -------
503 descriptors with shape ``(number_of_atoms, descriptor_components)``
504 """
505 self.calculate(atoms, ['descriptors'], system_changes)
506 return self.results['descriptors']
508 def get_born_effective_charges(
509 self,
510 atoms: Atoms = None,
511 properties: List[str] = None,
512 system_changes: List[str] = all_changes,
513 ) -> np.ndarray:
514 """Calculates (if needed) and returns the Born effective charges.
515 Note that this requires a qNEP model.
517 Parameters
518 ----------
519 atoms
520 System for which to calculate properties, by default `None`.
521 properties
522 Properties to calculate, by default `None`.
523 system_changes
524 Changes to the system since last call, by default all_changes.
525 """
526 if 'born_effective_charges' not in self.implemented_properties:
527 raise ValueError(
528 'This model does not support the calculation of Born effective charges.')
529 self.calculate(atoms, properties, system_changes)
530 return self.results['born_effective_charges']