Coverage for calorine/calculators/cpunep.py: 100%
164 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-23 12:50 +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
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
59 def __init__(
60 self,
61 model_filename: str,
62 atoms: Atoms | None = None,
63 label: str | None = None,
64 debug: bool = False,
65 ):
66 self.debug = debug
68 if not os.path.exists(model_filename):
69 raise FileNotFoundError(f'{model_filename} does not exist.')
70 self.model_filename = str(model_filename)
72 # Get model type from first row in nep.txt
73 header, _ = _get_nep_contents(self.model_filename)
74 self.model_type = header['model_type']
75 self.supported_species = set(header['types'])
76 self.nep_version = header['version']
78 # Set implemented properties -- not use class-level property
79 # to avoid leaking state between calculator instances.
80 self.implemented_properties = list(self.base_implemented_properties)
81 if 'charge' in self.model_type:
82 # Only available for charge models
83 self.implemented_properties.extend(
84 ['charges', 'born_effective_charges'])
85 elif self.model_type == 'dipole':
86 # Only available for dipole models
87 self.implemented_properties = ['dipole']
88 elif self.model_type == 'polarizability':
89 # Only available for polarizability models
90 self.implemented_properties = ['polarizability']
92 # Initialize atoms, results and nepy - note that this is also done in Calculator.__init__()
93 if atoms is not None:
94 self.set_atoms(atoms)
95 parameters = {'model_filename': model_filename}
96 Calculator.__init__(self, label=label, atoms=atoms, **parameters)
97 if atoms is not None:
98 self._setup_nepy()
100 def __str__(self) -> str:
101 def indent(s: str, i: int) -> str:
102 s = '\n'.join([i * ' ' + line for line in s.split('\n')])
103 return s
105 parameters = '\n'.join(
106 [f'{key}: {value}' for key, value in self.parameters.items()]
107 )
108 parameters = indent(parameters, 4)
109 using_debug = '\nIn debug mode' if self.debug else ''
111 s = f'{self.__class__.__name__}\n{parameters}{using_debug}'
112 return s
114 def _setup_nepy(self):
115 """
116 Creates an instance of the NEPY class and attaches it to the calculator object.
117 The output from `nep.cpp` is only written to STDOUT if debug == True
118 """
119 if self.atoms is None:
120 raise ValueError('Atoms must be defined when calculating properties.')
121 if self.atoms.cell.rank == 0:
122 raise ValueError('Atoms must have a defined cell.')
124 natoms = len(self.atoms)
125 self.natoms = natoms
126 c = self.atoms.get_cell(complete=True).flatten()
127 cell = [c[0], c[3], c[6], c[1], c[4], c[7], c[2], c[5], c[8]]
128 symbols = self.atoms.get_chemical_symbols()
129 positions = list(
130 self.atoms.get_positions().T.flatten()
131 ) # [x1, ..., xN, y1, ... yN,...]
132 masses = self.atoms.get_masses()
134 # Disable output from C++ code by default
135 if self.debug:
136 self.nepy = _nepy.NEPY(
137 self.model_filename, self.natoms, cell, symbols, positions, masses
138 )
139 else:
140 with TemporaryFile('w') as f:
141 with contextlib.redirect_stdout(f):
142 self.nepy = _nepy.NEPY(
143 self.model_filename,
144 self.natoms,
145 cell,
146 symbols,
147 positions,
148 masses,
149 )
151 def set_atoms(self, atoms: Atoms):
152 """Updates the Atoms object.
154 Parameters
155 ----------
156 atoms : Atoms
157 Atoms to attach the calculator to
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.')
162 self.atoms = atoms
163 self.results = {}
164 self.nepy = None
166 def _update_symbols(self):
167 """Update atom symbols in NEPY."""
168 symbols = self.atoms.get_chemical_symbols()
169 self.nepy.set_symbols(symbols)
171 def _update_masses(self):
172 """Update atom masses in NEPY"""
173 masses = self.atoms.get_masses()
174 self.nepy.set_masses(masses)
176 def _update_cell(self):
177 """Update cell parameters in NEPY."""
178 c = self.atoms.get_cell(complete=True).flatten()
179 cell = [c[0], c[3], c[6], c[1], c[4], c[7], c[2], c[5], c[8]]
180 self.nepy.set_cell(cell)
182 def _update_positions(self):
183 """Update atom positions in NEPY."""
184 positions = list(
185 self.atoms.get_positions().T.flatten()
186 ) # [x1, ..., xN, y1, ... yN,...]
187 self.nepy.set_positions(positions)
189 def calculate(
190 self,
191 atoms: Atoms = None,
192 properties: List[str] = None,
193 system_changes: List[str] = all_changes,
194 ):
195 """Calculate energy, per atom energies, forces, stress and dipole.
197 Parameters
198 ----------
199 atoms : Atoms, optional
200 System for which to calculate properties, by default None
201 properties : List[str], optional
202 Properties to calculate, by default None
203 system_changes : List[str], optional
204 Changes to the system since last call, by default all_changes
205 """
206 if properties is None:
207 properties = self.implemented_properties
209 Calculator.calculate(self, atoms, properties, system_changes)
211 if self.nepy is None:
212 # Create new NEPY interface
213 self._setup_nepy()
214 # Update existing NEPY interface
215 for change in system_changes:
216 if change == 'positions':
217 self._update_positions()
218 elif change == 'numbers':
219 self._update_symbols()
220 self._update_masses()
221 elif change == 'cell':
222 self._update_cell()
224 if 'dipole' in properties:
225 dipole = np.array(self.nepy.get_dipole())
226 self.results['dipole'] = dipole
227 elif 'polarizability' in properties:
228 # NEPY.get_polarizability() returns components in NEP_REDUCED6_ORDER
229 # (xx,yy,zz,xy,yz,zx); see `find_polarizability` in src/nepy/nep.cpp.
230 pol = np.array(self.nepy.get_polarizability())
231 polarizability = reduced6_to_full_3x3(pol)
232 self.results['polarizability'] = polarizability
233 elif 'descriptors' in properties:
234 descriptors = np.array(self.nepy.get_descriptors())
235 descriptors_per_atom = descriptors.reshape(-1, self.natoms).T
236 self.results['descriptors'] = descriptors_per_atom
237 else:
238 if 'charge' in self.model_type:
239 energies, forces, virials, charges, becs = \
240 self.nepy.get_potential_forces_virials_and_charges()
241 else:
242 energies, forces, virials = self.nepy.get_potential_forces_and_virials()
244 energies_per_atom = np.array(energies)
245 energy = energies_per_atom.sum()
246 forces_per_atom = np.array(forces).reshape(-1, self.natoms).T
247 # NEPY's per-atom virial is a raw, row-major full-3x3 tensor
248 # [xx,xy,xz,yx,yy,yz,zx,zy,zz]; documented in src/nepy/nep.h and
249 # matching the accumulation pattern in src/nepy/nep.cpp. Given
250 # that layout, `.reshape((3,3))` below is already correctly
251 # oriented, and ASE's own `full_3x3_to_voigt_6_stress` produces
252 # a correct ASE-Voigt-6 `stress` -- independently confirmed
253 # empirically via isolated shear strains, each producing a
254 # stress response concentrated at the expected component.
255 virials_per_atom = np.array(virials).reshape(-1, self.natoms).T
256 stresses_per_atom = virials_per_atom / self.atoms.get_volume()
257 stress = -(np.sum(virials_per_atom, axis=0) / self.atoms.get_volume()).reshape((3, 3))
258 stress = full_3x3_to_voigt_6_stress(stress)
260 self.results['energy'] = energy
261 self.results['forces'] = forces_per_atom
262 self.results['stress'] = stress
263 self.results['stresses'] = stresses_per_atom
265 if 'charge' in self.model_type:
266 charges_per_atom = np.array(charges)
267 # Row-major full-3x3 per atom, same convention as the virial
268 # above (BEC is not symmetric, so no reduced-6 form applies);
269 # see src/nepy/nep.cpp. Matches GPUMD's own `dump_xyz`/
270 # `bec_*.out` convention (src/force/nep_charge.cu,
271 # src/measure/dump_xyz.cu -- no reindexing applied to BEC).
272 becs_per_atom = np.array(becs).reshape(-1, self.natoms).T
274 self.results['charges'] = charges_per_atom
275 self.results['born_effective_charges'] = becs_per_atom
277 def get_dipole_gradient(
278 self,
279 displacement: float = 0.01,
280 method: str = 'central difference',
281 charge: float = 1.0,
282 ):
283 """Calculates the dipole gradient using finite differences.
285 Parameters
286 ----------
287 displacement
288 Displacement in Å to use for finite differences. Defaults to 0.01 Å.
289 method
290 Method for computing gradient with finite differences.
291 One of 'forward difference' and 'central difference'.
292 Defaults to 'central difference'
293 charge
294 System charge in units of the elemental charge.
295 Used for correcting the dipoles before computing the gradient.
296 Defaults to 1.0.
298 Returns
299 -------
300 dipole gradient with shape `(N, 3, 3)` where ``N`` are the number of atoms.
301 """
302 if 'dipole' not in self.implemented_properties:
303 raise ValueError('Dipole gradients are only defined for dipole NEP models.')
305 if displacement <= 0:
306 raise ValueError('displacement must be > 0 Å')
308 implemented_methods = {
309 'forward difference': 0,
310 'central difference': 1,
311 'second order central difference': 2,
312 }
314 if method not in implemented_methods.keys():
315 raise ValueError(f'Invalid method {method} for calculating gradient')
317 if self.nepy is None:
318 # Create new NEPY interface
319 self._setup_nepy()
321 dipole_gradient = np.array(
322 self.nepy.get_dipole_gradient(
323 displacement, implemented_methods[method], charge
324 )
325 ).reshape(self.natoms, 3, 3)
326 return dipole_gradient
328 def get_polarizability(
329 self,
330 atoms: Atoms = None,
331 properties: List[str] = None,
332 system_changes: List[str] = all_changes,
333 ) -> np.ndarray:
334 """Calculates the polarizability tensor for the current structure.
335 The model must have been trained to predict the polarizability.
336 This is a wrapper function for :func:`calculate`.
338 Parameters
339 ----------
340 atoms : Atoms, optional
341 System for which to calculate properties, by default None
342 properties : List[str], optional
343 Properties to calculate, by default None
344 system_changes : List[str], optional
345 Changes to the system since last call, by default all_changes
347 Returns
348 -------
349 polarizability with shape ``(3, 3)``
350 """
351 if properties is None:
352 properties = self.implemented_properties
354 if 'polarizability' not in properties:
355 raise ValueError('Polarizability is only defined for polarizability NEP models.')
356 self.calculate(atoms, properties, system_changes)
357 return self.results['polarizability']
359 def get_polarizability_gradient(
360 self,
361 displacement: float = 0.01,
362 component: Union[str, List[str]] = 'full',
363 ) -> np.ndarray:
364 """Calculates the dipole gradient for a given structure using finite differences.
365 This function computes the derivatives using the second-order central difference
366 method with a C++ backend.
368 Parameters
369 ----------
370 displacement
371 Displacement in Å to use for finite differences. Defaults to ``0.01``.
372 component
373 Component or components of the polarizability tensor that the gradient
374 should be computed for.
375 The following components are available: `x`, `y`, `z`, `full`
376 Option ``full`` computes the derivative whilst moving the atoms in each Cartesian
377 direction, which yields a tensor of shape ``(N, 3, 3, 3)``,
378 where ``N`` is the number of atoms.
379 Multiple components may be specified.
380 Defaults to ``full``.
382 Returns
383 -------
384 polarizability gradient with shape ``(N, C, 3, 3)`` with ``C`` components chosen.
385 """
386 if 'polarizability' not in self.implemented_properties:
387 raise ValueError('Polarizability gradients are only defined'
388 ' for polarizability NEP models.')
390 if displacement <= 0:
391 raise ValueError('displacement must be > 0 Å')
393 if self.nepy is None:
394 # Create new NEPY interface
395 self._setup_nepy()
397 component_array = _check_components_polarizability_gradient(component)
399 pg = np.array(
400 self.nepy.get_polarizability_gradient(
401 displacement, component_array
402 )
403 ).reshape(self.natoms, 3, 6)
404 polarizability_gradient = _polarizability_gradient_to_3x3(pg)
405 return polarizability_gradient[:, component_array, :, :]
407 def get_descriptors(
408 self,
409 atoms: Atoms = None,
410 properties: List[str] = None,
411 system_changes: List[str] = all_changes,
412 ) -> np.ndarray:
413 """Calculates the descriptor tensor for the current structure.
414 This is a wrapper function for :func:`calculate`.
416 Parameters
417 ----------
418 atoms : Atoms, optional
419 System for which to calculate properties, by default None
420 properties : List[str], optional
421 Properties to calculate, by default None
422 system_changes : List[str], optional
423 Changes to the system since last call, by default all_changes
425 Returns
426 -------
427 descriptors with shape ``(number_of_atoms, descriptor_components)``
428 """
429 self.calculate(atoms, ['descriptors'], system_changes)
430 return self.results['descriptors']
432 def get_born_effective_charges(
433 self,
434 atoms: Atoms = None,
435 properties: List[str] = None,
436 system_changes: List[str] = all_changes,
437 ) -> np.ndarray:
438 """Calculates (if needed) and returns the Born effective charges.
439 Note that this requires a qNEP model.
441 Parameters
442 ----------
443 atoms
444 System for which to calculate properties, by default `None`.
445 properties
446 Properties to calculate, by default `None`.
447 system_changes
448 Changes to the system since last call, by default all_changes.
449 """
450 if 'born_effective_charges' not in self.implemented_properties:
451 raise ValueError(
452 'This model does not support the calculation of Born effective charges.')
453 self.calculate(atoms, properties, system_changes)
454 return self.results['born_effective_charges']