Coverage for calorine/calculators/cpunep.py: 100%
146 statements
« prev ^ index » next coverage.py v7.6.4, created at 2024-12-10 08:26 +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
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.
34 Raises
35 ------
36 FileNotFoundError
37 Raises :class:`FileNotFoundError` if :attr:`model_filename` does not point to a valid file.
38 ValueError
39 Raises :class:`ValueError` atoms are not defined when trying to get energies and forces.
40 Example
41 -------
43 >>> calc = CPUNEP('nep.txt')
44 >>> atoms.calc = calc
45 >>> atoms.get_potential_energy()
46 """
48 implemented_properties = [
49 'energy',
50 'energies',
51 'forces',
52 'stress',
53 ]
54 debug = False
55 nepy = None
57 def __init__(
58 self,
59 model_filename: str,
60 atoms: Atoms | None = None,
61 label: str | None = None,
62 debug: bool = False,
63 ):
64 self.debug = debug
66 if not os.path.exists(model_filename):
67 raise FileNotFoundError(f'{model_filename} does not exist.')
68 self.model_filename = str(model_filename)
70 # Get model type from first row in nep.txt
71 header, _ = _get_nep_contents(self.model_filename)
72 self.model_type = header['model_type']
73 self.supported_species = set(header['types'])
74 self.nep_version = header['version']
76 if self.model_type == 'dipole':
77 # Only available for dipole models
78 self.implemented_properties = ['dipole']
79 elif self.model_type == 'polarizability':
80 # Only available for polarizability models
81 self.implemented_properties = ['polarizability']
83 # Initialize atoms, results and nepy - note that this is also done in Calculator.__init__()
84 if atoms is not None:
85 self.set_atoms(atoms)
86 parameters = {'model_filename': model_filename}
87 Calculator.__init__(self, label=label, atoms=atoms, **parameters)
88 if atoms is not None:
89 self._setup_nepy()
91 def __str__(self) -> str:
92 def indent(s: str, i: int) -> str:
93 s = '\n'.join([i * ' ' + line for line in s.split('\n')])
94 return s
96 parameters = '\n'.join(
97 [f'{key}: {value}' for key, value in self.parameters.items()]
98 )
99 parameters = indent(parameters, 4)
100 using_debug = '\nIn debug mode' if self.debug else ''
102 s = f'{self.__class__.__name__}\n{parameters}{using_debug}'
103 return s
105 def _setup_nepy(self):
106 """
107 Creates an instance of the NEPY class and attaches it to the calculator object.
108 The output from `nep.cpp` is only written to STDOUT if debug == True
109 """
110 if self.atoms is None:
111 raise ValueError('Atoms must be defined when calculating properties.')
112 if self.atoms.cell.rank == 0:
113 raise ValueError('Atoms must have a defined cell.')
115 natoms = len(self.atoms)
116 self.natoms = natoms
117 c = self.atoms.get_cell(complete=True).flatten()
118 cell = [c[0], c[3], c[6], c[1], c[4], c[7], c[2], c[5], c[8]]
119 symbols = self.atoms.get_chemical_symbols()
120 positions = list(
121 self.atoms.get_positions().T.flatten()
122 ) # [x1, ..., xN, y1, ... yN,...]
123 masses = self.atoms.get_masses()
125 # Disable output from C++ code by default
126 if self.debug:
127 self.nepy = _nepy.NEPY(
128 self.model_filename, self.natoms, cell, symbols, positions, masses
129 )
130 else:
131 with TemporaryFile('w') as f:
132 with contextlib.redirect_stdout(f):
133 self.nepy = _nepy.NEPY(
134 self.model_filename,
135 self.natoms,
136 cell,
137 symbols,
138 positions,
139 masses,
140 )
142 def set_atoms(self, atoms: Atoms):
143 """Updates the Atoms object.
145 Parameters
146 ----------
147 atoms : Atoms
148 Atoms to attach the calculator to
149 """
150 species_in_atoms_object = set(np.unique(atoms.get_chemical_symbols()))
151 if not species_in_atoms_object.issubset(self.supported_species):
152 raise ValueError('Structure contains species that are not supported by the NEP model.')
153 self.atoms = atoms
154 self.results = {}
155 self.nepy = None
157 def _update_symbols(self):
158 """Update atom symbols in NEPY."""
159 symbols = self.atoms.get_chemical_symbols()
160 self.nepy.set_symbols(symbols)
162 def _update_masses(self):
163 """Update atom masses in NEPY"""
164 masses = self.atoms.get_masses()
165 self.nepy.set_masses(masses)
167 def _update_cell(self):
168 """Update cell parameters in NEPY."""
169 c = self.atoms.get_cell(complete=True).flatten()
170 cell = [c[0], c[3], c[6], c[1], c[4], c[7], c[2], c[5], c[8]]
171 self.nepy.set_cell(cell)
173 def _update_positions(self):
174 """Update atom positions in NEPY."""
175 positions = list(
176 self.atoms.get_positions().T.flatten()
177 ) # [x1, ..., xN, y1, ... yN,...]
178 self.nepy.set_positions(positions)
180 def calculate(
181 self,
182 atoms: Atoms = None,
183 properties: List[str] = None,
184 system_changes: List[str] = all_changes,
185 ):
186 """Calculate energy, per atom energies, forces, stress and dipole.
188 Parameters
189 ----------
190 atoms : Atoms, optional
191 System for which to calculate properties, by default None
192 properties : List[str], optional
193 Properties to calculate, by default None
194 system_changes : List[str], optional
195 Changes to the system since last call, by default all_changes
196 """
197 if properties is None:
198 properties = self.implemented_properties
200 Calculator.calculate(self, atoms, properties, system_changes)
202 if self.nepy is None:
203 # Create new NEPY interface
204 self._setup_nepy()
205 # Update existing NEPY interface
206 for change in system_changes:
207 if change == 'positions':
208 self._update_positions()
209 elif change == 'numbers':
210 self._update_symbols()
211 self._update_masses()
212 elif change == 'cell':
213 self._update_cell()
215 if 'dipole' in properties:
216 dipole = np.array(self.nepy.get_dipole())
217 self.results['dipole'] = dipole
218 elif 'polarizability' in properties:
219 pol = np.array(self.nepy.get_polarizability())
220 polarizability = np.array([
221 [pol[0], pol[3], pol[5]],
222 [pol[3], pol[1], pol[4]],
223 [pol[5], pol[4], pol[2]]
224 ])
225 self.results['polarizability'] = polarizability
226 elif 'descriptors' in properties:
227 descriptors = np.array(self.nepy.get_descriptors())
228 descriptors_per_atom = descriptors.reshape(-1, self.natoms).T
229 self.results['descriptors'] = descriptors_per_atom
230 else:
231 energies, forces, virials = self.nepy.get_potential_forces_and_virials()
232 energies_per_atom = np.array(energies)
233 energy = energies_per_atom.sum()
234 forces_per_atom = np.array(forces).reshape(-1, self.natoms).T
235 virials_per_atom = np.array(virials).reshape(-1, self.natoms).T
236 stress = -(np.sum(virials_per_atom, axis=0) / self.atoms.get_volume()).reshape((3, 3))
237 stress = full_3x3_to_voigt_6_stress(stress)
239 self.results['energy'] = energy
240 self.results['forces'] = forces_per_atom
241 self.results['stress'] = stress
243 def get_dipole_gradient(
244 self,
245 displacement: float = 0.01,
246 method: str = 'central difference',
247 charge: float = 1.0,
248 ):
249 """Calculates the dipole gradient using finite differences.
251 Parameters
252 ----------
253 displacement
254 Displacement in Å to use for finite differences. Defaults to 0.01 Å.
255 method
256 Method for computing gradient with finite differences.
257 One of 'forward difference' and 'central difference'.
258 Defaults to 'central difference'
259 charge
260 System charge in units of the elemental charge.
261 Used for correcting the dipoles before computing the gradient.
262 Defaults to 1.0.
264 Returns
265 -------
266 dipole gradient with shape `(N, 3, 3)` where ``N`` are the number of atoms.
267 """
268 if 'dipole' not in self.implemented_properties:
269 raise ValueError('Dipole gradients are only defined for dipole NEP models.')
271 if displacement <= 0:
272 raise ValueError('displacement must be > 0 Å')
274 implemented_methods = {
275 'forward difference': 0,
276 'central difference': 1,
277 'second order central difference': 2,
278 }
280 if method not in implemented_methods.keys():
281 raise ValueError(f'Invalid method {method} for calculating gradient')
283 if self.nepy is None:
284 # Create new NEPY interface
285 self._setup_nepy()
287 dipole_gradient = np.array(
288 self.nepy.get_dipole_gradient(
289 displacement, implemented_methods[method], charge
290 )
291 ).reshape(self.natoms, 3, 3)
292 return dipole_gradient
294 def get_polarizability(
295 self,
296 atoms: Atoms = None,
297 properties: List[str] = None,
298 system_changes: List[str] = all_changes,
299 ) -> np.ndarray:
300 """Calculates the polarizability tensor for the current structure.
301 The model must have been trained to predict the polarizability.
302 This is a wrapper function for :func:`calculate`.
304 Parameters
305 ----------
306 atoms : Atoms, optional
307 System for which to calculate properties, by default None
308 properties : List[str], optional
309 Properties to calculate, by default None
310 system_changes : List[str], optional
311 Changes to the system since last call, by default all_changes
313 Returns
314 -------
315 polarizability with shape ``(3, 3)``
316 """
317 if properties is None:
318 properties = self.implemented_properties
320 if 'polarizability' not in properties:
321 raise ValueError('Polarizability is only defined for polarizability NEP models.')
322 self.calculate(atoms, properties, system_changes)
323 return self.results['polarizability']
325 def get_polarizability_gradient(
326 self,
327 displacement: float = 0.01,
328 component: Union[str, List[str]] = 'full',
329 ) -> np.ndarray:
330 """Calculates the dipole gradient for a given structure using finite differences.
331 This function computes the derivatives using the second-order central difference
332 method with a C++ backend.
334 Parameters
335 ----------
336 displacement
337 Displacement in Å to use for finite differences. Defaults to ``0.01``.
338 component
339 Component or components of the polarizability tensor that the gradient
340 should be computed for.
341 The following components are available: `x`, `y`, `z`, `full`
342 Option ``full`` computes the derivative whilst moving the atoms in each Cartesian
343 direction, which yields a tensor of shape ``(N, 3, 3, 3)``,
344 where ``N`` is the number of atoms.
345 Multiple components may be specified.
346 Defaults to ``full``.
348 Returns
349 -------
350 polarizability gradient with shape ``(N, C, 3, 3)`` with ``C`` components chosen.
351 """
352 if 'polarizability' not in self.implemented_properties:
353 raise ValueError('Polarizability gradients are only defined'
354 ' for polarizability NEP models.')
356 if displacement <= 0:
357 raise ValueError('displacement must be > 0 Å')
359 if self.nepy is None:
360 # Create new NEPY interface
361 self._setup_nepy()
363 component_array = _check_components_polarizability_gradient(component)
365 pg = np.array(
366 self.nepy.get_polarizability_gradient(
367 displacement, component_array
368 )
369 ).reshape(self.natoms, 3, 6)
370 polarizability_gradient = _polarizability_gradient_to_3x3(self.natoms, pg)
371 return polarizability_gradient[:, component_array, :, :]
373 def get_descriptors(
374 self,
375 atoms: Atoms = None,
376 properties: List[str] = None,
377 system_changes: List[str] = all_changes,
378 ) -> np.ndarray:
379 """Calculates the descriptor tensor for the current structure.
380 This is a wrapper function for :func:`calculate`.
382 Parameters
383 ----------
384 atoms : Atoms, optional
385 System for which to calculate properties, by default None
386 properties : List[str], optional
387 Properties to calculate, by default None
388 system_changes : List[str], optional
389 Changes to the system since last call, by default all_changes
391 Returns
392 -------
393 descriptors with shape ``(number_of_atoms, descriptor_components)``
394 """
395 self.calculate(atoms, ['descriptors'], system_changes)
396 return self.results['descriptors']