Coverage for calorine/calculators/gpunep.py: 100%
146 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-23 12:50 +0000
1import os
2import shutil
3import warnings
4import tempfile
5from collections.abc import Iterable
6from typing import Any, List, Tuple, Union
8import numpy as np
9from ase import Atoms
10from ase.calculators.calculator import FileIOCalculator, OldShellProfile, all_changes
11from ase.io import read as ase_read
12from ase.units import GPa
14from calorine.nep.model import _get_nep_contents
15from ..env import calorine_getenv
16from ..gpumd import write_xyz
19class GPUMDShellProfile(OldShellProfile):
20 """This class provides an ASE calculator for NEP calculations with
21 GPUMD.
23 Parameters
24 ----------
25 command : str
26 Command to run GPUMD with.
27 Default: ``gpumd``, or the value of the ``CALORINE_GPUMD_COMMAND``
28 environment variable if set.
29 gpu_identifier_index : int, None
30 Index that identifies the GPU that GPUNEP should be run with.
31 Typically, NVIDIA GPUs are enumerated with integer indices.
32 See https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#env-vars.
33 Set to None in order to use all available GPUs. Note that GPUMD exit with an error
34 when running with more than one GPU if your system is not large enough.
35 Default: None
36 """
37 def __init__(self, command : str, gpu_identifier_index: Union[int, None]):
38 if gpu_identifier_index is not None:
39 # Do not set a specific device to use = use all available GPUs
40 self.cuda_environment_variables = f'CUDA_VISIBLE_DEVICES={gpu_identifier_index}'
41 command_with_gpus = f'export {self.cuda_environment_variables} && ' + command
42 else:
43 command_with_gpus = command
44 super().__init__(command_with_gpus)
47class GPUNEP(FileIOCalculator):
48 """This class provides an ASE calculator for NEP calculations with
49 GPUMD.
51 This calculator writes files that are input to the `gpumd`
52 executable. It is thus likely to be slow if many calculations
53 are to be performed.
55 Parameters
56 ----------
57 model_filename : str
58 Path to file in ``nep.txt`` format with model parameters.
59 directory : str
60 Directory to run GPUMD in. If None, a temporary directory
61 will be created and removed once the calculations are finished.
62 If specified, the directory will not be deleted. In the latter
63 case, it is advisable to do no more than one calculation with
64 this calculator (unless you know exactly what you are doing).
65 label : str
66 Label for this calculator.
67 atoms : Atoms
68 Atoms to attach to this calculator.
69 command : str
70 Command to run GPUMD with.
71 Default: ``gpumd``, or the value of the ``CALORINE_GPUMD_COMMAND``
72 environment variable if set.
73 gpu_identifier_index : int
74 Index that identifies the GPU that GPUNEP should be run with.
75 Typically, NVIDIA GPUs are enumerated with integer indices.
76 See https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#env-vars.
77 Set to None in order to use all available GPUs. Note that GPUMD exit with an error
78 when running with more than one GPU if your system is not large enough.
79 Default: 0
82 Example
83 -------
85 >>> calc = GPUNEP('nep.txt')
86 >>> atoms.calc = calc
87 >>> atoms.get_potential_energy()
88 """
90 # Shadows the `command` property inherited from FileIOCalculator, which
91 # otherwise routes `self.command = ...` through `self.profile` before
92 # `self.profile` has been set up (see __init__ below).
93 command = 'gpumd'
94 base_implemented_properties = ['energy', 'forces', 'stress']
95 discard_results_on_any_change = True
97 # We use list of tuples to define parameters for
98 # MD simulations. Looks like a dictionary, but sometimes
99 # we want to repeat the same keyword.
100 base_single_point_parameters = [('dump_thermo', 1),
101 ('dump_force', 1),
102 ('dump_position', 1),
103 ('velocity', 1e-24),
104 ('time_step', 1e-6), # 1 zeptosecond
105 ('ensemble', 'nve'),
106 ('run', 1)]
108 def __init__(self,
109 model_filename: str,
110 directory: str = None,
111 label: str = 'GPUNEP',
112 atoms: Atoms = None,
113 command: str = None,
114 gpu_identifier_index: Union[int, None] = 0
115 ):
116 if command is None:
117 command = calorine_getenv('GPUMD_COMMAND')
118 if not os.path.exists(model_filename):
119 raise FileNotFoundError(f'{model_filename} does not exist.')
120 self.model_filename = str(model_filename)
122 # Get model type from first row in nep.txt
123 header, _ = _get_nep_contents(self.model_filename)
124 self.model_type = header['model_type']
125 self.supported_species = set(header['types'])
126 self.nep_version = header['version']
127 self.model_filename = model_filename
129 self.implemented_properties = list(self.base_implemented_properties)
130 self.single_point_parameters = self.base_single_point_parameters
131 if 'charge' in self.model_type:
132 # Only available for charge models
133 self.implemented_properties.extend(
134 ['charges', 'born_effective_charges'])
135 qnep_parameters = [('dump_xyz', (-1, 1, 1, 'charges_and_bec.xyz', 'charge', 'bec'))]
136 self.single_point_parameters = qnep_parameters + self.base_single_point_parameters
138 # Determine run command
139 # Determine whether to save stdout or not
140 if directory is None and '>' not in command:
141 # No need to save stdout if we run in temporary directory
142 command += ' > /dev/null'
143 elif '>' not in command:
144 command += ' > stdout'
145 self.command = command
147 # Determine directory to run in
148 self._use_temporary_directory = directory is None
149 self._directory = directory
150 if self._use_temporary_directory:
151 self._make_new_tmp_directory()
152 else:
153 self._potential_path = os.path.relpath(
154 os.path.abspath(self.model_filename), self._directory)
156 # Override the profile in ~/.config/ase/config.ini.
157 # See https://docs.ase-lib.org/ase/calculators/calculators.html#calculator-configuration
158 profile = GPUMDShellProfile(command, gpu_identifier_index)
159 FileIOCalculator.__init__(self,
160 directory=self._directory,
161 label=label,
162 atoms=atoms,
163 profile=profile)
165 def run_custom_md(
166 self,
167 parameters: List[Tuple[str, Any]],
168 return_last_atoms: bool = False,
169 only_prepare: bool = False,
170 ):
171 """
172 Run a custom MD simulation.
174 Parameters
175 ----------
176 parameters
177 Parameters to be specified in the run.in file.
178 The potential keyword is set automatically, all other
179 keywords need to be set via this argument.
180 Example::
182 [('dump_thermo', 100),
183 ('dump_position', 1000),
184 ('velocity', 300),
185 ('time_step', 1),
186 ('ensemble', ['nvt_ber', 300, 300, 100]),
187 ('run', 10000)]
189 return_last_atoms
190 If ``True`` the last saved snapshot will be returned.
191 only_prepare
192 If ``True`` the necessary input files will be written
193 but the MD run will not be executed.
195 Returns
196 -------
197 The last snapshot if :attr:`return_last_atoms` is ``True``.
198 """
199 if self._use_temporary_directory:
200 self._make_new_tmp_directory()
202 if self._use_temporary_directory and not return_last_atoms:
203 raise ValueError('Refusing to run in temporary directory '
204 'and not returning atoms; all results will be gone.')
206 if self._use_temporary_directory and only_prepare:
207 raise ValueError('Refusing to only prepare in temporary directory, '
208 'all files will be removed.')
210 # Write files and run
211 FileIOCalculator.write_input(self, self.atoms)
212 self._write_runfile(parameters)
213 write_xyz(filename=os.path.join(self._directory, 'model.xyz'),
214 structure=self.atoms)
216 if only_prepare:
217 return None
219 # Execute the calculation.
220 self.execute()
222 # Extract last snapshot if needed
223 if return_last_atoms:
224 last_atoms = ase_read(os.path.join(self._directory, 'movie.xyz'),
225 format='extxyz', index=-1)
227 if self._use_temporary_directory:
228 self._clean()
230 if return_last_atoms:
231 return last_atoms
232 else:
233 return None
235 def write_input(self, atoms, properties=None, system_changes=None):
236 """
237 Write the input files necessary for a single-point calculation.
238 """
239 if self._use_temporary_directory:
240 self._make_new_tmp_directory()
241 FileIOCalculator.write_input(self, atoms, properties, system_changes)
242 self._write_runfile(parameters=self.single_point_parameters)
243 write_xyz(filename=os.path.join(self._directory, 'model.xyz'),
244 structure=atoms)
246 def _write_runfile(self, parameters):
247 """Write run.in file to define input parameters for MD simulation.
249 Parameters
250 ----------
251 parameters : dict
252 Defines all key-value pairs used in run.in file
253 (see GPUMD documentation for a complete list).
254 Values can be either floats, integers, or lists/tuples.
255 """
256 if len(os.listdir(self._directory)) > 0:
257 warnings.warn(f'{self._directory} is not empty.')
259 with open(os.path.join(self._directory, 'run.in'), 'w') as f:
260 # Custom potential is allowed but normally it can be deduced
261 if 'potential' not in [keyval[0] for keyval in parameters]:
262 f.write(f'potential {self._potential_path} \n')
263 # Write all keywords with parameter(s)
264 for key, val in parameters:
265 f.write(f'{key} ')
266 if isinstance(val, Iterable) and not isinstance(val, str):
267 for v in val:
268 f.write(f'{v} ')
269 else:
270 f.write(f'{val}')
271 f.write('\n')
273 def get_potential_energy_and_stresses_from_file(self):
274 """
275 Extract potential energy (third column of last line in thermo.out) and stresses
276 from thermo.out
277 """
278 data = np.loadtxt(os.path.join(self._directory, 'thermo.out'))
279 if len(data.shape) == 1:
280 line = data
281 else:
282 line = data[-1, :]
284 # Energy
285 energy = line[2]
287 # Stress. GPUMD's src/measure/dump_thermo.cu (and dump_observer.cu)
288 # explicitly re-permutes its internal thermo array (xx,yy,zz,xy,xz,yz)
289 # into true ASE-Voigt order (xx,yy,zz,yz,xz,xy) before writing
290 # thermo.out, so these columns need no further permutation here --
291 # unlike the `nep` executable's virial_*.out/stress_*.out training
292 # output, which uses a different native order. Independently
293 # confirmed empirically (ad hoc, not a standing test) via isolated
294 # shear strains applied to a GPUNEP-attached structure, each
295 # producing a stress response concentrated at the expected
296 # component.
297 stress = [v for v in line[3:9]]
298 stress = -GPa * np.array(stress) # to eV/A^3
300 if np.any(np.isnan(stress)) or np.isnan(energy):
301 raise ValueError(f'Failed to extract energy and/or stresses:\n {line}')
302 return energy, stress
304 def _read_potential_energy_and_stresses(self):
305 """Reads potential energy and stresses."""
306 self.results['energy'], self.results['stress'] = \
307 self.get_potential_energy_and_stresses_from_file()
309 def get_forces_from_file(self):
310 """
311 Extract forces (in eV/A) from last snapshot in force.out
312 """
313 data = np.loadtxt(os.path.join(self._directory, 'force.out'))
314 return data[-len(self.atoms):, :]
316 def _read_forces(self):
317 """Reads forces (the last snapshot in force.out) in eV/A"""
318 self.results['forces'] = self.get_forces_from_file()
320 def get_charges_and_becs_from_file(self):
321 """Extract charges and Born Effective Charges from last
322 snapshot in `charges_and_bec.xyz`"""
323 structure = ase_read(os.path.join(self._directory, 'charges_and_bec.xyz'), '-1')
324 charges = structure.get_charges()
325 # Raw, row-major full-3x3 per atom (xx,xy,xz,yx,yy,yz,zx,zy,zz); BEC
326 # is not symmetric, so no reduced-6 form applies. GPUMD's
327 # src/measure/dump_xyz.cu writes `bec` with no reindexing, straight
328 # from src/force/nep_charge.cu's row-major per-atom buffer -- the
329 # same convention CPUNEP uses for its own BEC.
330 becs = structure.get_array('bec')
331 return charges, becs
333 def _read_charges_and_becs(self):
334 """Reads charges and Born Effective Charges from file."""
335 charges, becs = self.get_charges_and_becs_from_file()
336 self.results['charges'] = charges
337 self.results['born_effective_charges'] = becs
339 def read_results(self):
340 """
341 Read results from last step of MD calculation.
342 """
343 self._read_potential_energy_and_stresses()
344 self._read_forces()
346 if 'charge' in self.model_type:
347 self._read_charges_and_becs()
348 if self._use_temporary_directory:
349 self._clean()
351 def _clean(self):
352 """
353 Remove directory with calculations.
354 """
355 shutil.rmtree(self._directory)
357 def _make_new_tmp_directory(self):
358 """
359 Create a new temporary directory.
360 """
361 # We do not need to create a new temporary directory
362 # if the current one is empty
363 if self._directory is None or \
364 (os.path.isdir(self._directory) and len(os.listdir(self._directory)) > 0):
365 self._directory = tempfile.mkdtemp()
366 self._potential_path = os.path.relpath(os.path.abspath(self.model_filename),
367 self._directory)
369 def set_atoms(self, atoms):
370 """
371 Set Atoms object.
372 Used also when attaching calculator to Atoms object.
373 """
374 self.atoms = atoms
375 self.results = {}
377 def set_directory(self, directory):
378 """
379 Set path to a new directory. This makes it possible to run
380 several calculations with the same calculator while saving
381 all results
382 """
383 self._directory = directory
384 self._use_temporary_directory = False
385 self._potential_path = os.path.relpath(os.path.abspath(self.model_filename),
386 self._directory)
388 def get_born_effective_charges(
389 self,
390 atoms: Atoms = None,
391 properties: List[str] = None,
392 system_changes: List[str] = all_changes,
393 ) -> np.ndarray:
394 """Calculates (if needed) and returns the Born effective charges.
395 Note that this requires a qNEP model.
397 Parameters
398 ----------
399 atoms
400 System for which to calculate properties, by default `None`.
401 properties
402 Properties to calculate, by default `None`.
403 system_changes
404 Changes to the system since last call, by default all_changes.
405 """
406 if 'born_effective_charges' not in self.implemented_properties:
407 raise ValueError(
408 'This model does not support the calculation of Born effective charges.')
409 self.calculate(atoms, properties, system_changes)
410 return self.results['born_effective_charges']