Coverage for calorine/calculators/cpunep.py: 100%
114 statements
« prev ^ index » next coverage.py v7.2.5, created at 2024-07-26 09:34 +0000
1from __future__ import annotations
3import contextlib
4import os
5from tempfile import TemporaryFile
6from typing import List
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
17class CPUNEP(Calculator):
18 """This class provides an ASE calculator for `nep_cpu`,
19 the in-memory CPU implementation of GPUMD.
21 Parameters
22 ----------
23 model_filename : str
24 Path to file in ``nep.txt`` format with model parameters
25 atoms : Atoms
26 Atoms to attach the calculator to
27 label : str
28 Label for this calclator
29 debug : bool, optional
30 Flag to toggle debug mode. Prints GPUMD output. Defaults to False.
32 Raises
33 ------
34 FileNotFoundError
35 Raises :class:`FileNotFoundError` if :attr:`model_filename` does not point to a valid file.
36 ValueError
37 Raises :class:`ValueError` atoms are not defined when trying to get energies and forces.
38 Example
39 -------
41 >>> calc = CPUNEP('nep.txt')
42 >>> atoms.calc = calc
43 >>> atoms.get_potential_energy()
44 """
46 implemented_properties = [
47 'energy',
48 'energies',
49 'forces',
50 'stress',
51 ]
52 debug = False
53 nepy = None
55 def __init__(
56 self,
57 model_filename: str,
58 atoms: Atoms | None = None,
59 label: str | None = None,
60 debug: bool = False,
61 ):
62 self.debug = debug
64 if not os.path.exists(model_filename):
65 raise FileNotFoundError(f'{model_filename} does not exist.')
66 self.model_filename = str(model_filename)
68 # Get model type from first row in nep.txt
69 header, _ = _get_nep_contents(self.model_filename)
70 self.model_type = header['model_type']
71 self.supported_species = set(header['types'])
72 self.nep_version = header['version']
74 if self.model_type == 'dipole':
75 # Only available for dipole models
76 self.implemented_properties = ['dipole']
78 # Initialize atoms, results and nepy - note that this is also done in Calculator.__init__()
79 if atoms is not None:
80 self.set_atoms(atoms)
81 parameters = {'model_filename': model_filename}
82 Calculator.__init__(self, label=label, atoms=atoms, **parameters)
83 if atoms is not None:
84 self._setup_nepy()
86 def __str__(self) -> str:
87 def indent(s: str, i: int) -> str:
88 s = '\n'.join([i * ' ' + line for line in s.split('\n')])
89 return s
91 parameters = '\n'.join(
92 [f'{key}: {value}' for key, value in self.parameters.items()]
93 )
94 parameters = indent(parameters, 4)
95 using_debug = '\nIn debug mode' if self.debug else ''
97 s = f'{self.__class__.__name__}\n{parameters}{using_debug}'
98 return s
100 def _setup_nepy(self):
101 """
102 Creates an instance of the NEPY class and attaches it to the calculator object.
103 The output from `nep.cpp` is only written to STDOUT if debug == True
104 """
105 if self.atoms is None:
106 raise ValueError('Atoms must be defined to get energies and forces.')
107 if self.atoms.cell.rank == 0:
108 raise ValueError('Atoms must have a defined cell.')
110 natoms = len(self.atoms)
111 self.natoms = natoms
112 c = self.atoms.get_cell(complete=True).flatten()
113 cell = [c[0], c[3], c[6], c[1], c[4], c[7], c[2], c[5], c[8]]
114 symbols = self.atoms.get_chemical_symbols()
115 positions = list(
116 self.atoms.get_positions().T.flatten()
117 ) # [x1, ..., xN, y1, ... yN,...]
118 masses = self.atoms.get_masses()
120 # Disable output from C++ code by default
121 if self.debug:
122 self.nepy = _nepy.NEPY(
123 self.model_filename, self.natoms, cell, symbols, positions, masses
124 )
125 else:
126 with TemporaryFile('w') as f:
127 with contextlib.redirect_stdout(f):
128 self.nepy = _nepy.NEPY(
129 self.model_filename,
130 self.natoms,
131 cell,
132 symbols,
133 positions,
134 masses,
135 )
137 def set_atoms(self, atoms: Atoms):
138 """Updates the Atoms object.
140 Parameters
141 ----------
142 atoms : Atoms
143 Atoms to attach the calculator to
144 """
145 species_in_atoms_object = set(np.unique(atoms.get_chemical_symbols()))
146 if not species_in_atoms_object.issubset(self.supported_species):
147 raise ValueError('Structure contains species that are not supported by the NEP model.')
148 self.atoms = atoms
149 self.results = {}
150 self.nepy = None
152 def _update_symbols(self):
153 """Update atom symbols in NEPY."""
154 symbols = self.atoms.get_chemical_symbols()
155 self.nepy.set_symbols(symbols)
157 def _update_masses(self):
158 """Update atom masses in NEPY"""
159 masses = self.atoms.get_masses()
160 self.nepy.set_masses(masses)
162 def _update_cell(self):
163 """Update cell parameters in NEPY."""
164 c = self.atoms.get_cell(complete=True).flatten()
165 cell = [c[0], c[3], c[6], c[1], c[4], c[7], c[2], c[5], c[8]]
166 self.nepy.set_cell(cell)
168 def _update_positions(self):
169 """Update atom positions in NEPY."""
170 positions = list(
171 self.atoms.get_positions().T.flatten()
172 ) # [x1, ..., xN, y1, ... yN,...]
173 self.nepy.set_positions(positions)
175 def calculate(
176 self,
177 atoms: Atoms = None,
178 properties: List[str] = None,
179 system_changes: List[str] = all_changes,
180 ):
181 """Calculate energy, per atom energies, forces, stress and dipole.
183 Parameters
184 ----------
185 atoms : Atoms, optional
186 System for which to calculate properties, by default None
187 properties : List[str], optional
188 Properties to calculate, by default None
189 system_changes : List[str], optional
190 Changes to the system since last call, by default all_changes
191 """
192 if properties is None:
193 properties = self.implemented_properties
195 Calculator.calculate(self, atoms, properties, system_changes)
197 if self.nepy is None:
198 # Create new NEPY interface
199 self._setup_nepy()
200 # Update existing NEPY interface
201 for change in system_changes:
202 if change == 'positions':
203 self._update_positions()
204 elif change == 'numbers':
205 self._update_symbols()
206 self._update_masses()
207 elif change == 'cell':
208 self._update_cell()
210 if 'dipole' in self.implemented_properties:
211 dipole = np.array(self.nepy.get_dipole())
212 self.results['dipole'] = dipole
213 else:
214 energies, forces, virials = self.nepy.get_potential_forces_and_virials()
215 energies_per_atom = np.array(energies)
216 energy = energies_per_atom.sum()
217 forces_per_atom = np.array(forces).reshape(-1, self.natoms).T
218 virials_per_atom = np.array(virials).reshape(-1, self.natoms).T
219 stress = -(np.sum(virials_per_atom, axis=0) / self.atoms.get_volume()).reshape((3, 3))
220 stress = full_3x3_to_voigt_6_stress(stress)
222 self.results['energy'] = energy
223 self.results['forces'] = forces_per_atom
224 self.results['stress'] = stress
226 def get_dipole_gradient(
227 self,
228 displacement: float = 0.01,
229 method: str = 'central difference',
230 charge: float = 1.0,
231 ):
232 """Calculates the dipole gradient using finite differences.
234 Parameters
235 ----------
236 structure
237 Input structure
238 model_filename
239 Path to NEP model. Defaults to ``None``.
240 method
241 Method for computing gradient with finite differences.
242 One of 'forward difference' and 'central difference'.
243 Defaults to 'central difference'
244 displacement
245 Displacement in Å to use for finite differences. Defaults to 0.01 Å.
246 charge
247 System charge in units of the elemental charge.
248 Used for correcting the dipoles before computing the gradient.
249 Defaults to 1.0.
251 Returns
252 ------- dipole gradient with shape `(N, 3, 3)`
253 """
254 if 'dipole' not in self.implemented_properties:
255 raise ValueError('Dipole gradients are only defined for dipole NEP models.')
257 if displacement <= 0:
258 raise ValueError('displacement must be > 0 Å')
260 implemented_methods = {
261 'forward difference': 0,
262 'central difference': 1,
263 'second order central difference': 2,
264 }
266 if method not in implemented_methods.keys():
267 raise ValueError(f'Invalid method {method} for calculating gradient')
269 if self.nepy is None:
270 # Create new NEPY interface
271 self._setup_nepy()
273 dipole_gradient = np.array(
274 self.nepy.get_dipole_gradient(
275 displacement, implemented_methods[method], charge
276 )
277 ).reshape(self.natoms, 3, 3)
278 return dipole_gradient