Coverage for calorine/calculators/gpunep.py: 100%
119 statements
« prev ^ index » next coverage.py v7.6.4, created at 2024-12-10 08:26 +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
11from ase.io import read as ase_read
12from ase.units import GPa
14from ..gpumd import write_xyz
17class GPUMDShellProfile(OldShellProfile):
18 """This class provides an ASE calculator for NEP calculations with
19 GPUMD.
21 Parameters
22 ----------
23 command : str
24 Command to run GPUMD with.
25 Default: ``gpumd``
26 gpu_identifier_index : int, None
27 Index that identifies the GPU that GPUNEP should be run with.
28 Typically, NVIDIA GPUs are enumerated with integer indices.
29 See https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#env-vars.
30 Set to None in order to use all available GPUs. Note that GPUMD exit with an error
31 when running with more than one GPU if your system is not large enough.
32 Default: 1
33 """
34 def __init__(self, command : str, gpu_identifier_index: Union[int, None]):
35 if gpu_identifier_index is not None:
36 # Do not set a specific device to use = use all available GPUs
37 self.cuda_environment_variables = f'CUDA_VISIBLE_DEVICES={gpu_identifier_index}'
38 command_with_gpus = f'export {self.cuda_environment_variables} && ' + command
39 else:
40 command_with_gpus = command
41 super().__init__(command_with_gpus)
44class GPUNEP(FileIOCalculator):
45 """This class provides an ASE calculator for NEP calculations with
46 GPUMD.
48 This calculator writes files that are input to the `gpumd`
49 executable. It is thus likely to be slow if many calculations
50 are to be performed.
52 Parameters
53 ----------
54 model_filename : str
55 Path to file in ``nep.txt`` format with model parameters.
56 directory : str
57 Directory to run GPUMD in. If None, a temporary directory
58 will be created and removed once the calculations are finished.
59 If specified, the directory will not be deleted. In the latter
60 case, it is advisable to do no more than one calculation with
61 this calculator (unless you know exactly what you are doing).
62 label : str
63 Label for this calculator.
64 atoms : Atoms
65 Atoms to attach to this calculator.
66 command : str
67 Command to run GPUMD with.
68 Default: ``gpumd``
69 gpu_identifier_index : int
70 Index that identifies the GPU that GPUNEP should be run with.
71 Typically, NVIDIA GPUs are enumerated with integer indices.
72 See https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#env-vars.
73 Set to None in order to use all available GPUs. Note that GPUMD exit with an error
74 when running with more than one GPU if your system is not large enough.
75 Default: 1
78 Example
79 -------
81 >>> calc = GPUNEP('nep.txt')
82 >>> atoms.calc = calc
83 >>> atoms.get_potential_energy()
84 """
86 command = 'gpumd'
87 implemented_properties = ['energy', 'forces', 'stress']
88 discard_results_on_any_change = True
90 # We use list of tuples to define parameters for
91 # MD simulations. Looks like a dictionary, but sometimes
92 # we want to repeat the same keyword.
93 single_point_parameters = [('dump_thermo', 1),
94 ('dump_force', 1),
95 ('dump_position', 1),
96 ('velocity', 1e-24),
97 ('time_step', 1e-6), # 1 zeptosecond
98 ('ensemble', 'nve'),
99 ('run', 1)]
101 def __init__(self,
102 model_filename: str,
103 directory: str = None,
104 label: str = 'GPUNEP',
105 atoms: Atoms = None,
106 command: str = command,
107 gpu_identifier_index: Union[int, None] = 0
108 ):
110 # Determine run command
111 # Determine whether to save stdout or not
112 if directory is None and '>' not in command:
113 # No need to save stdout if we run in temporary directory
114 command += ' > /dev/null'
115 elif '>' not in command:
116 command += ' > stdout'
117 self.command = command
118 self.model_filename = model_filename
120 # Determine directory to run in
121 self._use_temporary_directory = directory is None
122 self._directory = directory
123 if self._use_temporary_directory:
124 self._make_new_tmp_directory()
125 else:
126 self._potential_path = os.path.relpath(
127 os.path.abspath(self.model_filename), self._directory)
129 # Override the profile in ~/.config/ase/config.ini.
130 # See https://wiki.fysik.dtu.dk/ase/ase/
131 # calculators/calculators.html#calculator-configuration
132 profile = GPUMDShellProfile(command, gpu_identifier_index)
133 FileIOCalculator.__init__(self,
134 directory=self._directory,
135 label=label,
136 atoms=atoms,
137 profile=profile)
139 # If the model file is missing we should abort immediately
140 # such that we can provide a more clear error message
141 # (otherwise the code would fail without telling what is wrong).
142 if not os.path.exists(model_filename):
143 raise FileNotFoundError(f'{model_filename} does not exist.')
145 # Read species from nep.txt
146 with open(model_filename, 'r') as f:
147 for line in f:
148 if 'nep' in line:
149 self.species = line.split()[2:]
151 def run_custom_md(
152 self,
153 parameters: List[Tuple[str, Any]],
154 return_last_atoms: bool = False,
155 only_prepare: bool = False,
156 ):
157 """
158 Run a custom MD simulation.
160 Parameters
161 ----------
162 parameters
163 Parameters to be specified in the run.in file.
164 The potential keyword is set automatically, all other
165 keywords need to be set via this argument.
166 Example::
168 [('dump_thermo', 100),
169 ('dump_position', 1000),
170 ('velocity', 300),
171 ('time_step', 1),
172 ('ensemble', ['nvt_ber', 300, 300, 100]),
173 ('run', 10000)]
175 return_last_atoms
176 If ``True`` the last saved snapshot will be returned.
177 only_prepare
178 If ``True`` the necessary input files will be written
179 but theMD run will not be executed.
181 Returns
182 -------
183 The last snapshot if :attr:`return_last_atoms` is ``True``.
184 """
185 if self._use_temporary_directory:
186 self._make_new_tmp_directory()
188 if self._use_temporary_directory and not return_last_atoms:
189 raise ValueError('Refusing to run in temporary directory '
190 'and not returning atoms; all results will be gone.')
192 if self._use_temporary_directory and only_prepare:
193 raise ValueError('Refusing to only prepare in temporary directory, '
194 'all files will be removed.')
196 # Write files and run
197 FileIOCalculator.write_input(self, self.atoms)
198 self._write_runfile(parameters)
199 write_xyz(filename=os.path.join(self._directory, 'model.xyz'),
200 structure=self.atoms)
202 if only_prepare:
203 return None
205 # Execute the calculation.
206 self.execute()
208 # Extract last snapshot if needed
209 if return_last_atoms:
210 last_atoms = ase_read(os.path.join(self._directory, 'movie.xyz'),
211 format='extxyz', index=-1)
213 if self._use_temporary_directory:
214 self._clean()
216 if return_last_atoms:
217 return last_atoms
218 else:
219 return None
221 def write_input(self, atoms, properties=None, system_changes=None):
222 """
223 Write the input files necessary for a single-point calculation.
224 """
225 if self._use_temporary_directory:
226 self._make_new_tmp_directory()
227 FileIOCalculator.write_input(self, atoms, properties, system_changes)
228 self._write_runfile(parameters=self.single_point_parameters)
229 write_xyz(filename=os.path.join(self._directory, 'model.xyz'),
230 structure=atoms)
232 def _write_runfile(self, parameters):
233 """Write run.in file to define input parameters for MD simulation.
235 Parameters
236 ----------
237 parameters : dict
238 Defines all key-value pairs used in run.in file
239 (see GPUMD documentation for a complete list).
240 Values can be either floats, integers, or lists/tuples.
241 """
242 if len(os.listdir(self._directory)) > 0:
243 warnings.warn(f'{self._directory} is not empty.')
245 with open(os.path.join(self._directory, 'run.in'), 'w') as f:
246 # Custom potential is allowed but normally it can be deduced
247 if 'potential' not in [keyval[0] for keyval in parameters]:
248 f.write(f'potential {self._potential_path} \n')
249 # Write all keywords with parameter(s)
250 for key, val in parameters:
251 f.write(f'{key} ')
252 if isinstance(val, Iterable) and not isinstance(val, str):
253 for v in val:
254 f.write(f'{v} ')
255 else:
256 f.write(f'{val}')
257 f.write('\n')
259 def get_potential_energy_and_stresses_from_file(self):
260 """
261 Extract potential energy (third column of last line in thermo.out) and stresses
262 from thermo.out
263 """
264 data = np.loadtxt(os.path.join(self._directory, 'thermo.out'))
265 if len(data.shape) == 1:
266 line = data
267 else:
268 line = data[-1, :]
270 # Energy
271 energy = line[2]
273 # Stress
274 stress = [v for v in line[3:9]]
275 stress = -GPa * np.array(stress) # to eV/A^3
277 if np.any(np.isnan(stress)) or np.isnan(energy):
278 raise ValueError(f'Failed to extract energy and/or stresses:\n {line}')
279 return energy, stress
281 def _read_potential_energy_and_stresses(self):
282 """Reads potential energy and stresses."""
283 self.results['energy'], self.results['stress'] = \
284 self.get_potential_energy_and_stresses_from_file()
286 def get_forces_from_file(self):
287 """
288 Extract forces (in eV/A) from last snapshot in force.out
289 """
290 data = np.loadtxt(os.path.join(self._directory, 'force.out'))
291 return data[-len(self.atoms):, :]
293 def _read_forces(self):
294 """Reads forces (the last snapshot in force.out) in eV/A"""
295 self.results['forces'] = self.get_forces_from_file()
297 def read_results(self):
298 """
299 Read results from last step of MD calculation.
300 """
301 self._read_potential_energy_and_stresses()
302 self._read_forces()
303 if self._use_temporary_directory:
304 self._clean()
306 def _clean(self):
307 """
308 Remove directory with calculations.
309 """
310 shutil.rmtree(self._directory)
312 def _make_new_tmp_directory(self):
313 """
314 Create a new temporary directory.
315 """
316 # We do not need to create a new temporary directory
317 # if the current one is empty
318 if self._directory is None or \
319 (os.path.isdir(self._directory) and len(os.listdir(self._directory)) > 0):
320 self._directory = tempfile.mkdtemp()
321 self._potential_path = os.path.relpath(os.path.abspath(self.model_filename),
322 self._directory)
324 def set_atoms(self, atoms):
325 """
326 Set Atoms object.
327 Used also when attaching calculator to Atoms object.
328 """
329 self.atoms = atoms
330 self.results = {}
332 def set_directory(self, directory):
333 """
334 Set path to a new directory. This makes it possible to run
335 several calculations with the same calculator while saving
336 all results
337 """
338 self._directory = directory
339 self._use_temporary_directory = False
340 self._potential_path = os.path.relpath(os.path.abspath(self.model_filename),
341 self._directory)