Coverage for calorine/calculators/gpunep.py: 99%
168 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-14 16:36 +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 (CalculationFailed, FileIOCalculator, OldShellProfile,
11 all_changes)
12from ase.io import read as ase_read
13from ase.units import GPa
15from calorine.nep.model import _get_nep_contents
16from ..env import calorine_getenv
17from ..gpumd import write_xyz
20# Files that a previous run in the same directory leaves behind: the input files
21# GPUNEP writes itself, the result files it reads back, and the redirected stdout.
22# Any file ending in `.out` is counted as well, since that covers the output of
23# every GPUMD keyword without having to enumerate them.
24_PREVIOUS_RUN_FILES = ('run.in', 'model.xyz', 'movie.xyz', 'charges_and_bec.xyz', 'stdout')
26# Where the standard error of GPUMD is kept. It is deliberately not part of _PREVIOUS_RUN_FILES,
27# since the shell creates it on every run and an empty one says nothing about an earlier
28# calculation.
29_STDERR_FILE = 'stderr'
32def _find_previous_run_files(directory: str) -> List[str]:
33 """Return the names of the files in :attr:`directory` that indicate an
34 earlier calculation was run there.
36 The presence of such files, rather than a non-empty directory, is what
37 signals that results may be read back from an earlier calculation by
38 mistake. A directory holding only unrelated files, such as the ``nep.txt``
39 the model is loaded from, is not reported.
41 Parameters
42 ----------
43 directory
44 Directory to inspect.
46 Returns
47 -------
48 Sorted names of the files that indicate an earlier calculation,
49 empty if there are none.
51 Example
52 -------
53 >>> _find_previous_run_files('some_directory_with_a_thermo_out_file')
54 ['thermo.out']
55 """
56 return sorted(filename for filename in os.listdir(directory)
57 if filename in _PREVIOUS_RUN_FILES or filename.endswith('.out'))
60class GPUMDShellProfile(OldShellProfile):
61 """This class provides an ASE calculator for NEP calculations with
62 GPUMD.
64 Parameters
65 ----------
66 command : str
67 Command to run GPUMD with.
68 Default: ``gpumd``, or the value of the ``CALORINE_GPUMD_COMMAND``
69 environment variable if set.
70 gpu_identifier_index : int, None
71 Index that identifies the GPU that GPUNEP should be run with.
72 Typically, NVIDIA GPUs are enumerated with integer indices.
73 See https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#env-vars.
74 Set to None in order to use all available GPUs. Note that GPUMD exit with an error
75 when running with more than one GPU if your system is not large enough.
76 Default: None
77 """
78 def __init__(self, command : str, gpu_identifier_index: Union[int, None]):
79 if gpu_identifier_index is not None:
80 # Do not set a specific device to use = use all available GPUs
81 self.cuda_environment_variables = f'CUDA_VISIBLE_DEVICES={gpu_identifier_index}'
82 command_with_gpus = f'export {self.cuda_environment_variables} && ' + command
83 else:
84 command_with_gpus = command
85 super().__init__(command_with_gpus)
88class GPUNEP(FileIOCalculator):
89 """This class provides an ASE calculator for NEP calculations with
90 GPUMD.
92 This calculator writes files that are input to the `gpumd`
93 executable. It is thus likely to be slow if many calculations
94 are to be performed.
96 Parameters
97 ----------
98 model_filename : str
99 Path to file in ``nep.txt`` format with model parameters.
100 directory : str
101 Directory to run GPUMD in. If None, a temporary directory
102 will be created and removed once the calculations are finished.
103 If specified, the directory will not be deleted. In the latter
104 case, it is advisable to do no more than one calculation with
105 this calculator (unless you know exactly what you are doing).
106 label : str
107 Label for this calculator.
108 atoms : Atoms
109 Atoms to attach to this calculator.
110 command : str
111 Command to run GPUMD with.
112 Default: ``gpumd``, or the value of the ``CALORINE_GPUMD_COMMAND``
113 environment variable if set.
114 gpu_identifier_index : int
115 Index that identifies the GPU that GPUNEP should be run with.
116 Typically, NVIDIA GPUs are enumerated with integer indices.
117 See https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#env-vars.
118 Set to None in order to use all available GPUs. Note that GPUMD exit with an error
119 when running with more than one GPU if your system is not large enough.
120 Default: 0
123 Example
124 -------
126 >>> calc = GPUNEP('nep.txt')
127 >>> atoms.calc = calc
128 >>> atoms.get_potential_energy()
129 """
131 # Shadows the `command` property inherited from FileIOCalculator, which
132 # otherwise routes `self.command = ...` through `self.profile` before
133 # `self.profile` has been set up (see __init__ below).
134 command = 'gpumd'
135 base_implemented_properties = ['energy', 'forces', 'stress']
136 discard_results_on_any_change = True
138 # We use list of tuples to define parameters for
139 # MD simulations. Looks like a dictionary, but sometimes
140 # we want to repeat the same keyword.
141 # A single dump_xyz writes both the positions and the forces to movie.xyz.
142 # `precision double` is needed because dump_xyz writes nine significant digits by default,
143 # which is fewer than the forces carry.
144 base_single_point_parameters = [('dump_thermo', 1),
145 ('dump_xyz', (1, 'movie.xyz', 'precision', 'double', 'force')),
146 ('velocity', 1e-24),
147 ('time_step', 1e-6), # 1 zeptosecond
148 ('ensemble', 'nve'),
149 ('run', 1)]
151 def __init__(self,
152 model_filename: str,
153 directory: str = None,
154 label: str = 'GPUNEP',
155 atoms: Atoms = None,
156 command: str = None,
157 gpu_identifier_index: Union[int, None] = 0
158 ):
159 if command is None:
160 command = calorine_getenv('GPUMD_COMMAND')
161 if not os.path.exists(model_filename):
162 raise FileNotFoundError(f'{model_filename} does not exist.')
163 self.model_filename = str(model_filename)
165 # Get model type from first row in nep.txt
166 header, _ = _get_nep_contents(self.model_filename)
167 self.model_type = header['model_type']
168 self.supported_species = set(header['types'])
169 self.nep_version = header['version']
170 self.model_filename = model_filename
172 self.implemented_properties = list(self.base_implemented_properties)
173 self.single_point_parameters = self.base_single_point_parameters
174 if 'charge' in self.model_type:
175 # Only available for charge models
176 self.implemented_properties.extend(
177 ['charges', 'born_effective_charges'])
178 qnep_parameters = [('dump_xyz', (1, 'charges_and_bec.xyz', 'precision', 'double',
179 'charge', 'bec'))]
180 self.single_point_parameters = qnep_parameters + self.base_single_point_parameters
182 # Determine run command
183 # Determine whether to save stdout or not
184 if directory is None and '>' not in command:
185 # No need to save stdout if we run in temporary directory
186 command += ' > /dev/null'
187 elif '>' not in command:
188 command += ' > stdout'
189 # GPUMD reports input and CUDA errors on stderr (see PRINT_INPUT_ERROR in
190 # src/utilities/error.cuh), while stdout only carries the banner. Keeping stderr in its own
191 # file is what lets a failure be reported with the reason GPUMD gave for it.
192 if '2>' not in command: 192 ↛ 194line 192 didn't jump to line 194 because the condition on line 192 was always true
193 command += f' 2> {_STDERR_FILE}'
194 self.command = command
196 # Determine directory to run in
197 self._use_temporary_directory = directory is None
198 self._directory = directory
199 if self._use_temporary_directory:
200 self._make_new_tmp_directory()
201 else:
202 self._potential_path = os.path.relpath(
203 os.path.abspath(self.model_filename), self._directory)
205 # Override the profile in ~/.config/ase/config.ini.
206 # See https://docs.ase-lib.org/ase/calculators/calculators.html#calculator-configuration
207 profile = GPUMDShellProfile(command, gpu_identifier_index)
208 FileIOCalculator.__init__(self,
209 directory=self._directory,
210 label=label,
211 atoms=atoms,
212 profile=profile)
214 def run_custom_md(
215 self,
216 parameters: List[Tuple[str, Any]],
217 return_last_atoms: bool = False,
218 only_prepare: bool = False,
219 ):
220 """
221 Run a custom MD simulation.
223 Parameters
224 ----------
225 parameters
226 Parameters to be specified in the run.in file.
227 The potential keyword is set automatically, all other
228 keywords need to be set via this argument.
229 Example::
231 [('dump_thermo', 100),
232 ('dump_xyz', (1000, 'movie.xyz')),
233 ('velocity', 300),
234 ('time_step', 1),
235 ('ensemble', ['nvt_ber', 300, 300, 100]),
236 ('run', 10000)]
238 return_last_atoms
239 If ``True`` the last saved snapshot will be returned.
240 This requires :attr:`parameters` to contain a ``dump_xyz`` keyword writing to
241 ``movie.xyz``, since that is the file read back.
242 only_prepare
243 If ``True`` the necessary input files will be written
244 but the MD run will not be executed.
246 Returns
247 -------
248 The last snapshot if :attr:`return_last_atoms` is ``True``.
249 """
250 if self._use_temporary_directory:
251 self._make_new_tmp_directory()
253 if self._use_temporary_directory and not return_last_atoms:
254 raise ValueError('Refusing to run in temporary directory '
255 'and not returning atoms; all results will be gone.')
257 if self._use_temporary_directory and only_prepare:
258 raise ValueError('Refusing to only prepare in temporary directory, '
259 'all files will be removed.')
261 # Write files and run
262 FileIOCalculator.write_input(self, self.atoms)
263 self._write_runfile(parameters)
264 write_xyz(filename=os.path.join(self._directory, 'model.xyz'),
265 structure=self.atoms)
267 if only_prepare:
268 return None
270 # Execute the calculation.
271 self.execute()
273 # Extract last snapshot if needed
274 if return_last_atoms:
275 last_atoms = ase_read(os.path.join(self._directory, 'movie.xyz'),
276 format='extxyz', index=-1)
278 if self._use_temporary_directory:
279 self._clean()
281 if return_last_atoms:
282 return last_atoms
283 else:
284 return None
286 def execute(self):
287 """
288 Run GPUMD, reporting what it wrote to standard error if it fails.
290 :program:`ase` raises :class:`CalculationFailed` carrying only the exit code, which is not
291 enough to tell an input error apart from a missing GPU or an unusable model. GPUMD explains
292 itself on standard error, so that text is attached to the exception.
293 """
294 try:
295 FileIOCalculator.execute(self)
296 except CalculationFailed as exception:
297 message = self._read_stderr()
298 if message is None:
299 raise
300 raise CalculationFailed(f'{exception}\nGPUMD reported:\n{message}') from exception
302 def _read_stderr(self):
303 """Return what GPUMD wrote to standard error during the last run.
305 Returns
306 -------
307 The captured text with surrounding whitespace removed, or ``None`` if the file is
308 absent, unreadable, or empty.
309 """
310 try:
311 with open(os.path.join(self._directory, _STDERR_FILE)) as handle:
312 message = handle.read().strip()
313 except OSError:
314 return None
315 return message if message else None
317 def write_input(self, atoms, properties=None, system_changes=None):
318 """
319 Write the input files necessary for a single-point calculation.
320 """
321 if self._use_temporary_directory:
322 self._make_new_tmp_directory()
323 FileIOCalculator.write_input(self, atoms, properties, system_changes)
324 self._write_runfile(parameters=self.single_point_parameters)
325 write_xyz(filename=os.path.join(self._directory, 'model.xyz'),
326 structure=atoms)
328 def _write_runfile(self, parameters):
329 """Write run.in file to define input parameters for MD simulation.
331 Parameters
332 ----------
333 parameters : dict
334 Defines all key-value pairs used in run.in file
335 (see GPUMD documentation for a complete list).
336 Values can be either floats, integers, or lists/tuples.
337 """
338 previous_run_files = _find_previous_run_files(self._directory)
339 if previous_run_files:
340 warnings.warn(f'{self._directory} already contains files from an earlier '
341 f'calculation: {", ".join(previous_run_files)}. Results may be '
342 'read from those rather than from the calculation about to run.')
344 with open(os.path.join(self._directory, 'run.in'), 'w') as f:
345 # Custom potential is allowed but normally it can be deduced
346 if 'potential' not in [keyval[0] for keyval in parameters]:
347 f.write(f'potential {self._potential_path} \n')
348 # Write all keywords with parameter(s)
349 for key, val in parameters:
350 f.write(f'{key} ')
351 if isinstance(val, Iterable) and not isinstance(val, str):
352 for v in val:
353 f.write(f'{v} ')
354 else:
355 f.write(f'{val}')
356 f.write('\n')
358 def get_potential_energy_and_stresses_from_file(self):
359 """
360 Extract potential energy (third column of last line in thermo.out) and stresses
361 from thermo.out
362 """
363 data = np.loadtxt(os.path.join(self._directory, 'thermo.out'))
364 if len(data.shape) == 1:
365 line = data
366 else:
367 line = data[-1, :]
369 # Energy
370 energy = line[2]
372 # Stress. GPUMD's src/measure/dump_thermo.cu (and dump_observer.cu)
373 # explicitly re-permutes its internal thermo array (xx,yy,zz,xy,xz,yz)
374 # into true ASE-Voigt order (xx,yy,zz,yz,xz,xy) before writing
375 # thermo.out, so these columns need no further permutation here --
376 # unlike the `nep` executable's virial_*.out/stress_*.out training
377 # output, which uses a different native order. Independently
378 # confirmed empirically (ad hoc, not a standing test) via isolated
379 # shear strains applied to a GPUNEP-attached structure, each
380 # producing a stress response concentrated at the expected
381 # component.
382 stress = [v for v in line[3:9]]
383 stress = -GPa * np.array(stress) # to eV/A^3
385 if np.any(np.isnan(stress)) or np.isnan(energy):
386 raise ValueError(f'Failed to extract energy and/or stresses:\n {line}')
387 return energy, stress
389 def _read_potential_energy_and_stresses(self):
390 """Reads potential energy and stresses."""
391 self.results['energy'], self.results['stress'] = \
392 self.get_potential_energy_and_stresses_from_file()
394 def get_forces_from_file(self):
395 """
396 Extract forces (in eV/A) from last snapshot in movie.xyz
397 """
398 # GPUMD writes the energy and the stress on the comment line, so ase attaches a
399 # SinglePointCalculator and the forces are reached through get_forces rather than arrays.
400 structure = ase_read(os.path.join(self._directory, 'movie.xyz'),
401 format='extxyz', index=-1)
402 return structure.get_forces()
404 def _read_forces(self):
405 """Reads forces (the last snapshot in movie.xyz) in eV/A"""
406 self.results['forces'] = self.get_forces_from_file()
408 def get_charges_and_becs_from_file(self):
409 """Extract charges and Born Effective Charges from last
410 snapshot in `charges_and_bec.xyz`"""
411 structure = ase_read(os.path.join(self._directory, 'charges_and_bec.xyz'), '-1')
412 charges = structure.get_charges()
413 # Raw, row-major full-3x3 per atom (xx,xy,xz,yx,yy,yz,zx,zy,zz); BEC
414 # is not symmetric, so no reduced-6 form applies. GPUMD's
415 # src/measure/dump_xyz.cu writes `bec` with no reindexing, straight
416 # from src/force/nep_charge.cu's row-major per-atom buffer -- the
417 # same convention CPUNEP uses for its own BEC.
418 becs = structure.get_array('bec')
419 return charges, becs
421 def _read_charges_and_becs(self):
422 """Reads charges and Born Effective Charges from file."""
423 charges, becs = self.get_charges_and_becs_from_file()
424 self.results['charges'] = charges
425 self.results['born_effective_charges'] = becs
427 def read_results(self):
428 """
429 Read results from last step of MD calculation.
430 """
431 self._read_potential_energy_and_stresses()
432 self._read_forces()
434 if 'charge' in self.model_type:
435 self._read_charges_and_becs()
436 if self._use_temporary_directory:
437 self._clean()
439 def _clean(self):
440 """
441 Remove directory with calculations.
442 """
443 shutil.rmtree(self._directory)
445 def _make_new_tmp_directory(self):
446 """
447 Create a new temporary directory.
448 """
449 # We do not need to create a new temporary directory
450 # if the current one is empty
451 if self._directory is None or \
452 (os.path.isdir(self._directory) and len(os.listdir(self._directory)) > 0):
453 self._directory = tempfile.mkdtemp()
454 self._potential_path = os.path.relpath(os.path.abspath(self.model_filename),
455 self._directory)
457 def set_atoms(self, atoms):
458 """
459 Set Atoms object.
460 Used also when attaching calculator to Atoms object.
461 """
462 self.atoms = atoms
463 self.results = {}
465 def set_directory(self, directory):
466 """
467 Set path to a new directory. This makes it possible to run
468 several calculations with the same calculator while saving
469 all results
470 """
471 self._directory = directory
472 self._use_temporary_directory = False
473 self._potential_path = os.path.relpath(os.path.abspath(self.model_filename),
474 self._directory)
476 def get_born_effective_charges(
477 self,
478 atoms: Atoms = None,
479 properties: List[str] = None,
480 system_changes: List[str] = all_changes,
481 ) -> np.ndarray:
482 """Calculates (if needed) and returns the Born effective charges.
483 Note that this requires a qNEP model.
485 Parameters
486 ----------
487 atoms
488 System for which to calculate properties, by default `None`.
489 properties
490 Properties to calculate, by default `None`.
491 system_changes
492 Changes to the system since last call, by default all_changes.
493 """
494 if 'born_effective_charges' not in self.implemented_properties:
495 raise ValueError(
496 'This model does not support the calculation of Born effective charges.')
497 self.calculate(atoms, properties, system_changes)
498 return self.results['born_effective_charges']