Coverage for calorine/calculators/gpunep.py: 100%
196 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-10-04 20:18 +0000
1import os
2import shutil
3import warnings
4import tempfile
5from collections.abc import Iterable
6from typing import Any
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 read_thermo, write_xyz
18from ._dftd3 import DFTD3_KEYS, parse_dftd3
21# Where the standard output of GPUMD is kept when a directory is given.
22_STDOUT_FILE = 'stdout'
24# Files that a previous run in the same directory leaves behind: the input files
25# GPUNEP writes itself, the result files it reads back, and the redirected stdout.
26# Any file ending in `.out` is counted as well, since that covers the output of
27# every GPUMD keyword without having to enumerate them.
28_PREVIOUS_RUN_FILES = ('run.in', 'model.xyz', 'movie.xyz', 'charges_and_bec.xyz', _STDOUT_FILE)
30# Where the standard error of GPUMD is kept. It is deliberately not part of _PREVIOUS_RUN_FILES,
31# since the shell creates it on every run and an empty one says nothing about an earlier
32# calculation.
33_STDERR_FILE = 'stderr'
35# Below this a cutoff is shorter than any interatomic distance, and GPUMD does not return.
36_MINIMUM_CUTOFF = 1.0
39def _find_previous_run_files(directory: str) -> list[str]:
40 """Return the names of the files in :attr:`directory` that indicate an
41 earlier calculation was run there.
43 The presence of such files, rather than a non-empty directory, is what
44 signals that results may be read back from an earlier calculation by
45 mistake. A directory holding only unrelated files, such as the ``nep.txt``
46 the model is loaded from, is not reported.
48 Parameters
49 ----------
50 directory
51 Directory to inspect.
53 Returns
54 -------
55 Sorted names of the files that indicate an earlier calculation,
56 empty if there are none.
58 Example
59 -------
60 >>> # xdoctest: +SKIP
61 >>> _find_previous_run_files('some_directory_with_a_thermo_out_file')
62 ['thermo.out']
63 """
64 return sorted(filename for filename in os.listdir(directory)
65 if filename in _PREVIOUS_RUN_FILES or filename.endswith('.out'))
68class GPUMDShellProfile(OldShellProfile):
69 """This class provides an ASE calculator for NEP calculations with
70 GPUMD.
72 Parameters
73 ----------
74 command : str
75 Command to run GPUMD with.
76 Default: ``gpumd``, or the value of the ``CALORINE_GPUMD_COMMAND``
77 environment variable if set.
78 gpu_identifier_index : int, None
79 Index that identifies the GPU that GPUNEP should be run with.
80 Typically, NVIDIA GPUs are enumerated with integer indices.
81 See https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#env-vars.
82 Set to None in order to use all available GPUs. Note that GPUMD exit with an error
83 when running with more than one GPU if your system is not large enough.
84 Default: None
85 """
86 def __init__(self, command : str, gpu_identifier_index: int | None):
87 if gpu_identifier_index is not None:
88 # Do not set a specific device to use = use all available GPUs
89 self.cuda_environment_variables = f'CUDA_VISIBLE_DEVICES={gpu_identifier_index}'
90 command_with_gpus = f'export {self.cuda_environment_variables} && ' + command
91 else:
92 command_with_gpus = command
93 super().__init__(command_with_gpus)
96class GPUNEP(FileIOCalculator):
97 """This class provides an ASE calculator for NEP calculations with
98 GPUMD.
100 This calculator writes files that are input to the `gpumd`
101 executable. It is thus likely to be slow if many calculations
102 are to be performed.
104 Parameters
105 ----------
106 model_filename : str
107 Path to file in ``nep.txt`` format with model parameters.
108 directory : str
109 Directory to run GPUMD in. If None, a temporary directory
110 will be created and removed once the calculations are finished.
111 If specified, the directory will not be deleted. In the latter
112 case, it is advisable to do no more than one calculation with
113 this calculator (unless you know exactly what you are doing).
114 label : str
115 Label for this calculator.
116 atoms : Atoms
117 Atoms to attach to this calculator.
118 command : str
119 Command to run GPUMD with.
120 Default: ``gpumd``, or the value of the ``CALORINE_GPUMD_COMMAND``
121 environment variable if set.
122 gpu_identifier_index : int
123 Index that identifies the GPU that GPUNEP should be run with.
124 Typically, NVIDIA GPUs are enumerated with integer indices.
125 See https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#env-vars.
126 Set to None in order to use all available GPUs. Note that GPUMD exit with an error
127 when running with more than one GPU if your system is not large enough.
128 Default: 0
129 dftd3 : dict
130 Settings for the DFT-D3 dispersion correction, which is applied to every run of this
131 calculator, meaning single-point calculations as well as :func:`run_custom_md`.
132 The keys are ``functional``, ``potential_cutoff`` and ``coordination_number_cutoff``,
133 the two cutoffs in units of Angstrom.
134 Example: ``dict(functional='pbe', potential_cutoff=12, coordination_number_cutoff=6)``.
135 GPUMD carries the list of functionals, documented under the ``dftd3`` keyword at
136 https://gpumd.org, and implements the correction for a single GPU only.
137 The dispersion energy is truncated at the potential cutoff without a switching
138 function, so a cutoff that coincides with a neighbor shell of the structure makes the
139 energy and the pressure step as an atom crosses it. Choose a cutoff that falls between
140 shells.
141 Default: None
144 Example
145 -------
147 >>> # xdoctest: +SKIP
148 >>> calc = GPUNEP('nep.txt')
149 >>> atoms.calc = calc
150 >>> atoms.get_potential_energy()
151 """
153 # Shadows the `command` property inherited from FileIOCalculator, which
154 # otherwise routes `self.command = ...` through `self.profile` before
155 # `self.profile` has been set up (see __init__ below).
156 command = 'gpumd'
157 base_implemented_properties = ['energy', 'forces', 'stress']
158 discard_results_on_any_change = True
160 # We use list of tuples to define parameters for
161 # MD simulations. Looks like a dictionary, but sometimes
162 # we want to repeat the same keyword.
163 # A single dump_xyz writes both the positions and the forces to movie.xyz.
164 # `precision double` is needed because dump_xyz writes nine significant digits by default,
165 # which is fewer than the forces carry.
166 base_single_point_parameters = [('dump_thermo', 1),
167 ('dump_xyz', (1, 'movie.xyz', 'precision', 'double', 'force')),
168 ('velocity', 1e-24),
169 ('time_step', 1e-6), # 1 zeptosecond
170 ('ensemble', 'nve'),
171 ('run', 1)]
173 def __init__(self,
174 model_filename: str,
175 directory: str = None,
176 label: str = 'GPUNEP',
177 atoms: Atoms = None,
178 command: str = None,
179 gpu_identifier_index: int | None = 0,
180 dftd3: dict[str, Any] | None = None,
181 ):
182 self._dftd3 = parse_dftd3(dftd3)
183 if self._dftd3 is not None:
184 for key, cutoff in zip(DFTD3_KEYS[1:], self._dftd3[1:]):
185 if cutoff < _MINIMUM_CUTOFF:
186 raise ValueError(f'{key} must be at least {_MINIMUM_CUTOFF} Angstrom, not '
187 f'{cutoff}. GPUMD divides the cell into bins of the cutoff '
188 'and does not return for a cutoff far below an interatomic '
189 'distance.')
190 if command is None:
191 command = calorine_getenv('GPUMD_COMMAND')
192 if not os.path.exists(model_filename):
193 raise FileNotFoundError(f'{model_filename} does not exist.')
194 self.model_filename = str(model_filename)
196 # Get model type from first row in nep.txt
197 header, _ = _get_nep_contents(self.model_filename)
198 self.model_type = header['model_type']
199 self.supported_species = set(header['types'])
200 self.nep_version = header['version']
201 self.model_filename = model_filename
203 self.implemented_properties = list(self.base_implemented_properties)
204 self.single_point_parameters = self.base_single_point_parameters
205 if 'charge' in self.model_type:
206 # Only available for charge models
207 self.implemented_properties.extend(
208 ['charges', 'born_effective_charges'])
209 qnep_parameters = [('dump_xyz', (1, 'charges_and_bec.xyz', 'precision', 'double',
210 'charge', 'bec'))]
211 self.single_point_parameters = qnep_parameters + self.base_single_point_parameters
213 # Determine run command
214 # Determine whether to save stdout or not.
215 # A stderr redirection carries a `>` as well, so it is removed before the command is
216 # checked for a stdout redirection.
217 # The stdout redirection is inserted ahead of a stderr one, since `2>&1` duplicates
218 # whichever file descriptor 1 points at when the shell reads it.
219 # Appending would leave `2>&1 > stdout`, which sends stdout to the file and stderr to
220 # the terminal.
221 stdout_target = '/dev/null' if directory is None else _STDOUT_FILE
222 if '>' not in command.replace('2>', ''):
223 head, separator, tail = command.partition('2>')
224 command = f'{head.rstrip()} > {stdout_target}'
225 if separator:
226 command += f' {separator}{tail}'
227 # GPUMD reports input and CUDA errors on stderr (see PRINT_INPUT_ERROR in
228 # src/utilities/error.cuh), while stdout only carries the banner. Keeping stderr in its own
229 # file is what lets a failure be reported with the reason GPUMD gave for it.
230 if '2>' not in command:
231 command += f' 2> {_STDERR_FILE}'
232 self.command = command
234 # Determine directory to run in
235 self._use_temporary_directory = directory is None
236 self._directory = directory
237 if self._use_temporary_directory:
238 self._make_new_tmp_directory()
239 else:
240 self._potential_path = os.path.relpath(
241 os.path.abspath(self.model_filename), self._directory)
243 # Override the profile in ~/.config/ase/config.ini.
244 # See https://docs.ase-lib.org/ase/calculators/calculators.html#calculator-configuration
245 profile = GPUMDShellProfile(command, gpu_identifier_index)
246 FileIOCalculator.__init__(self,
247 directory=self._directory,
248 label=label,
249 atoms=atoms,
250 profile=profile)
251 if self._dftd3 is not None:
252 # Recorded among the parameters, so that a printed calculator and anything written
253 # from todict() name the correction the energy carries. It is written after the
254 # base class has run, since Calculator.__init__ resets the parameters to the
255 # defaults.
256 self.parameters['dftd3'] = self.dftd3
258 def set(self, **kwargs):
259 """Set parameters as defined by the ASE calculator interface.
261 Raises ``ValueError`` for ``dftd3`` settings other than those in use, since they are
262 fixed when the calculator is created.
263 """
264 if 'dftd3' in kwargs and parse_dftd3(kwargs.pop('dftd3')) != self._dftd3:
265 raise ValueError('The DFT-D3 settings are fixed when the calculator is created.')
266 return FileIOCalculator.set(self, **kwargs)
268 @property
269 def dftd3(self) -> dict[str, Any] | None:
270 """Settings for the DFT-D3 dispersion correction, None if the calculator applies none.
272 The settings are fixed when the calculator is created, since the results already
273 obtained were computed with them.
274 """
275 if self._dftd3 is None:
276 return None
277 return dict(zip(DFTD3_KEYS, self._dftd3))
279 def run_custom_md(
280 self,
281 parameters: list[tuple[str, Any]],
282 return_last_atoms: bool = False,
283 only_prepare: bool = False,
284 ):
285 """
286 Run a custom MD simulation.
288 Parameters
289 ----------
290 parameters
291 Parameters to be specified in the run.in file.
292 The potential keyword is set automatically, all other
293 keywords need to be set via this argument.
294 A ``dftd3`` keyword raises here if the calculator carries a dispersion
295 correction of its own.
296 Example::
298 [('dump_thermo', 100),
299 ('dump_xyz', (1000, 'movie.xyz')),
300 ('velocity', 300),
301 ('time_step', 1),
302 ('ensemble', ['nvt_ber', 300, 300, 100]),
303 ('run', 10000)]
305 return_last_atoms
306 If ``True`` the last saved snapshot will be returned.
307 This requires :attr:`parameters` to contain a ``dump_xyz`` keyword writing to
308 ``movie.xyz``, since that is the file read back.
309 only_prepare
310 If ``True`` the necessary input files will be written
311 but the MD run will not be executed.
313 Returns
314 -------
315 The last snapshot if :attr:`return_last_atoms` is ``True``.
316 """
317 if self._use_temporary_directory:
318 self._make_new_tmp_directory()
320 if self._use_temporary_directory and not return_last_atoms:
321 raise ValueError('Refusing to run in temporary directory '
322 'and not returning atoms; all results will be gone.')
324 if self._use_temporary_directory and only_prepare:
325 raise ValueError('Refusing to only prepare in temporary directory, '
326 'all files will be removed.')
328 # Write files and run
329 FileIOCalculator.write_input(self, self.atoms)
330 self._write_runfile(parameters)
331 write_xyz(filename=os.path.join(self._directory, 'model.xyz'),
332 structure=self.atoms)
334 if only_prepare:
335 return None
337 # Execute the calculation.
338 self.execute()
340 # Extract last snapshot if needed
341 if return_last_atoms:
342 last_atoms = ase_read(os.path.join(self._directory, 'movie.xyz'),
343 format='extxyz', index=-1)
345 if self._use_temporary_directory:
346 self._clean()
348 if return_last_atoms:
349 return last_atoms
350 else:
351 return None
353 def execute(self):
354 """
355 Run GPUMD, reporting what it wrote to standard error if it fails.
357 :program:`ase` raises :class:`CalculationFailed` carrying only the exit code, which is not
358 enough to tell an input error apart from a missing GPU or an unusable model. GPUMD explains
359 itself on standard error, so that text is attached to the exception.
361 An unsupported ``dftd3`` functional and a cutoff GPUMD cannot parse are reported on
362 standard output instead, which is kept only when :attr:`directory` is given, so the
363 tail of that file is used when standard error is empty.
364 """
365 try:
366 FileIOCalculator.execute(self)
367 except CalculationFailed as exception:
368 message = self._read_error_output()
369 if message is None:
370 raise
371 raise CalculationFailed(f'{exception}\nGPUMD reported:\n{message}') from exception
373 def _read_error_output(self):
374 """Return what GPUMD wrote about the last run.
376 Standard error is returned if it is not empty.
377 Otherwise the last 20 lines of standard output are returned, since GPUMD reports an
378 unusable ``dftd3`` line there.
380 Returns
381 -------
382 The captured text with surrounding whitespace removed, or ``None`` where neither
383 file is present, readable or non-empty.
384 """
385 for filename, lines_kept in ((_STDERR_FILE, None), (_STDOUT_FILE, 20)):
386 try:
387 with open(os.path.join(self._directory, filename)) as handle:
388 message = handle.read().strip()
389 except OSError:
390 continue
391 if message:
392 if lines_kept is not None:
393 message = '\n'.join(message.splitlines()[-lines_kept:])
394 return message
395 return None
397 def write_input(self, atoms, properties=None, system_changes=None):
398 """
399 Write the input files necessary for a single-point calculation.
400 """
401 if self._use_temporary_directory:
402 self._make_new_tmp_directory()
403 FileIOCalculator.write_input(self, atoms, properties, system_changes)
404 self._write_runfile(parameters=self.single_point_parameters)
405 write_xyz(filename=os.path.join(self._directory, 'model.xyz'),
406 structure=atoms)
408 def _write_runfile(self, parameters):
409 """Write run.in file to define input parameters for MD simulation.
411 Parameters
412 ----------
413 parameters : dict
414 Defines all key-value pairs used in run.in file
415 (see GPUMD documentation for a complete list).
416 Values can be either floats, integers, or lists/tuples.
417 """
418 if self._dftd3 is not None:
419 if 'dftd3' in [keyval[0] for keyval in parameters]:
420 raise ValueError('The calculator carries a dispersion correction, so the '
421 'parameters must not contain a dftd3 keyword.')
422 # A new list, since self.single_point_parameters is bound to the class attribute.
423 parameters = [('dftd3', self._dftd3)] + list(parameters)
425 previous_run_files = _find_previous_run_files(self._directory)
426 if previous_run_files:
427 warnings.warn(f'{self._directory} already contains files from an earlier '
428 f'calculation: {", ".join(previous_run_files)}. Results may be '
429 'read from those rather than from the calculation about to run.')
431 with open(os.path.join(self._directory, 'run.in'), 'w') as f:
432 # Custom potential is allowed but normally it can be deduced
433 if 'potential' not in [keyval[0] for keyval in parameters]:
434 f.write(f'potential {self._potential_path} \n')
435 # Write all keywords with parameter(s)
436 for key, val in parameters:
437 f.write(f'{key} ')
438 if isinstance(val, Iterable) and not isinstance(val, str):
439 for v in val:
440 f.write(f'{v} ')
441 else:
442 f.write(f'{val}')
443 f.write('\n')
445 def get_potential_energy_and_stresses_from_file(self):
446 """Reads the potential energy and the stresses from the last row of
447 ``thermo.out``, taking them by column name.
449 Returns
450 -------
451 Potential energy in eV and the stress tensor in eV/Å^3, the latter in
452 ASE Voigt order (xx, yy, zz, yz, xz, xy).
454 Raises
455 ------
456 ValueError
457 If the last row of ``thermo.out`` holds no energy or no stress.
458 """
459 thermo = read_thermo(os.path.join(self._directory, 'thermo.out'),
460 normalize=False)
461 line = thermo.iloc[-1]
463 # Energy
464 energy = line['potential_energy']
466 # Stress. GPUMD's src/measure/dump_thermo.cu (and dump_observer.cu)
467 # explicitly re-permutes its internal thermo array (xx,yy,zz,xy,xz,yz)
468 # into true ASE-Voigt order (xx,yy,zz,yz,xz,xy) before writing
469 # thermo.out, so these columns need no further permutation here --
470 # unlike the `nep` executable's virial_*.out/stress_*.out training
471 # output, which uses a different native order. Independently
472 # confirmed empirically (ad hoc, not a standing test) via isolated
473 # shear strains applied to a GPUNEP-attached structure, each
474 # producing a stress response concentrated at the expected
475 # component.
476 stress = line[['stress_xx', 'stress_yy', 'stress_zz',
477 'stress_yz', 'stress_xz', 'stress_xy']].to_numpy(dtype=float)
478 stress = -GPa * stress # to eV/A^3
480 if np.any(np.isnan(stress)) or np.isnan(energy):
481 raise ValueError(f'Failed to extract energy and/or stresses:\n {line}')
482 return energy, stress
484 def _read_potential_energy_and_stresses(self):
485 """Reads potential energy and stresses."""
486 self.results['energy'], self.results['stress'] = \
487 self.get_potential_energy_and_stresses_from_file()
489 def get_forces_from_file(self):
490 """Reads the forces from the last frame of ``movie.xyz``.
492 Returns
493 -------
494 Forces in eV/Å, of shape ``(len(atoms), 3)``.
495 """
496 # GPUMD writes the energy and the stress on the comment line, so ase attaches a
497 # SinglePointCalculator and the forces are reached through get_forces rather than arrays.
498 structure = ase_read(os.path.join(self._directory, 'movie.xyz'),
499 format='extxyz', index=-1)
500 return structure.get_forces()
502 def _read_forces(self):
503 """Reads forces (the last snapshot in movie.xyz) in eV/A"""
504 self.results['forces'] = self.get_forces_from_file()
506 def get_charges_and_becs_from_file(self):
507 """Reads the charges and the Born effective charges from the last frame of
508 ``charges_and_bec.xyz``, which a charge model writes.
510 Returns
511 -------
512 Charges in e, of shape ``(len(atoms),)``, and Born effective charges of
513 shape ``(len(atoms), 9)`` in row-major full-3x3 order.
514 """
515 structure = ase_read(os.path.join(self._directory, 'charges_and_bec.xyz'), '-1')
516 charges = structure.get_charges()
517 # Raw, row-major full-3x3 per atom (xx,xy,xz,yx,yy,yz,zx,zy,zz); BEC
518 # is not symmetric, so no reduced-6 form applies. GPUMD's
519 # src/measure/dump_xyz.cu writes `bec` with no reindexing, straight
520 # from src/force/nep_charge.cu's row-major per-atom buffer -- the
521 # same convention CPUNEP uses for its own BEC.
522 becs = structure.get_array('bec')
523 return charges, becs
525 def _read_charges_and_becs(self):
526 """Reads charges and Born Effective Charges from file."""
527 charges, becs = self.get_charges_and_becs_from_file()
528 self.results['charges'] = charges
529 self.results['born_effective_charges'] = becs
531 def read_results(self):
532 """
533 Read results from last step of MD calculation.
534 """
535 self._read_potential_energy_and_stresses()
536 self._read_forces()
538 if 'charge' in self.model_type:
539 self._read_charges_and_becs()
540 if self._use_temporary_directory:
541 self._clean()
543 def _clean(self):
544 """
545 Remove directory with calculations.
546 """
547 shutil.rmtree(self._directory)
549 def _make_new_tmp_directory(self):
550 """
551 Create a new temporary directory.
552 """
553 # We do not need to create a new temporary directory
554 # if the current one is empty
555 if self._directory is None or \
556 (os.path.isdir(self._directory) and len(os.listdir(self._directory)) > 0):
557 self._directory = tempfile.mkdtemp()
558 self._potential_path = os.path.relpath(os.path.abspath(self.model_filename),
559 self._directory)
561 def set_atoms(self, atoms):
562 """
563 Set Atoms object.
564 Used also when attaching calculator to Atoms object.
565 """
566 self.atoms = atoms
567 self.results = {}
569 def set_directory(self, directory):
570 """
571 Set path to a new directory. This makes it possible to run
572 several calculations with the same calculator while saving
573 all results
574 """
575 self._directory = directory
576 self._use_temporary_directory = False
577 self._potential_path = os.path.relpath(os.path.abspath(self.model_filename),
578 self._directory)
580 def get_born_effective_charges(
581 self,
582 atoms: Atoms = None,
583 properties: list[str] = None,
584 system_changes: list[str] = all_changes,
585 ) -> np.ndarray:
586 """Calculates (if needed) and returns the Born effective charges.
587 Note that this requires a qNEP model.
589 Parameters
590 ----------
591 atoms
592 System for which to calculate properties, by default `None`.
593 properties
594 Properties to calculate, by default `None`.
595 system_changes
596 Changes to the system since last call, by default all_changes.
597 """
598 if 'born_effective_charges' not in self.implemented_properties:
599 raise ValueError(
600 'This model does not support the calculation of Born effective charges.')
601 self.calculate(atoms, properties, system_changes)
602 return self.results['born_effective_charges']