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