Coverage for calorine/calculators/gpunep.py: 100%
118 statements
« prev ^ index » next coverage.py v7.2.5, created at 2024-07-26 09:34 +0000
1import os
2import tempfile
3import shutil
4import subprocess
5import warnings
6from collections.abc import Iterable
7from typing import Any, List, Tuple
9import numpy as np
10from ase import Atoms
11from ase.calculators.calculator import FileIOCalculator
12from ase.io import read as ase_read
13from ase.units import GPa
14from ..gpumd import write_xyz
17class GPUNEP(FileIOCalculator):
18 """This class provides an ASE calculator for NEP calculations with
19 GPUMD.
21 This calculator writes files that are input to the `gpumd`
22 executable. It is thus likely to be slow if many calculations
23 are to be performed.
25 Parameters
26 ----------
27 model_filename : str
28 Path to file in ``nep.txt`` format with model parameters.
29 directory : str
30 Directory to run GPUMD in. If None, a temporary directory
31 will be created and removed once the calculations are finished.
32 If specified, the directory will not be deleted. In the latter
33 case, it is advisable to do no more than one calculation with
34 this calculator (unless you know exactly what you are doing).
35 label : str
36 Label for this calculator.
37 atoms : Atoms
38 Atoms to attach to this calculator.
39 command : str
40 Command to run GPUMD with.
41 Default: ``gpumd``
42 prefix : str
43 Filename (excluding file ending) for run script.
44 Default: ``to_run``
46 Example
47 -------
49 >>> calc = GPUNEP('nep.txt')
50 >>> atoms.calc = calc
51 >>> atoms.get_potential_energy()
52 """
54 command = 'gpumd'
55 implemented_properties = ['energy', 'forces', 'stress']
56 discard_results_on_any_change = True
58 # We use list of tuples to define parameters for
59 # MD simulations. Looks like a dictionary, but sometimes
60 # we want to repeat the same keyword.
61 single_point_parameters = [('dump_thermo', 1),
62 ('dump_force', 1),
63 ('dump_position', 1),
64 ('velocity', 1e-24),
65 ('time_step', 1e-6), # 1 zeptosecond
66 ('ensemble', 'nve'),
67 ('run', 1)]
69 def __init__(self,
70 model_filename: str,
71 directory: str = None,
72 label: str = 'GPUNEP',
73 atoms: Atoms = None,
74 command: str = command):
76 # Determine run command
77 # Determine whether to save stdout or not
78 if directory is None and '>' not in command:
79 # No need to save stdout if we run in temporary directory
80 command += ' > /dev/null'
81 elif '>' not in command:
82 command += ' > stdout'
83 self.command = command
85 self.model_filename = model_filename
87 # Determine directory to run in
88 self._use_temporary_directory = directory is None
89 self._directory = directory
90 if self._use_temporary_directory:
91 self._make_new_tmp_directory()
92 else:
93 self._potential_path = os.path.relpath(
94 os.path.abspath(self.model_filename), self._directory)
96 FileIOCalculator.__init__(self, directory=self._directory, label=label,
97 atoms=atoms)
99 # If the model file is missing we should abort immediately
100 # such that we can provide a more clear error message
101 # (otherwise the code would fail without telling what is wrong).
102 if not os.path.exists(model_filename):
103 raise FileNotFoundError(f'{model_filename} does not exist.')
105 # Read species from nep.txt
106 with open(model_filename, 'r') as f:
107 for line in f:
108 if 'nep' in line:
109 self.species = line.split()[2:]
111 def run_custom_md(
112 self,
113 parameters: List[Tuple[str, Any]],
114 return_last_atoms: bool = False,
115 only_prepare: bool = False,
116 ):
117 """
118 Run a custom MD simulation.
120 Parameters
121 ----------
122 parameters
123 Parameters to be specified in the run.in file.
124 The potential keyword is set automatically, all other
125 keywords need to be set via this argument.
126 Example::
128 [('dump_thermo', 100),
129 ('dump_position', 1000),
130 ('velocity', 300),
131 ('time_step', 1),
132 ('ensemble', ['nvt_ber', 300, 300, 100]),
133 ('run', 10000)]
135 return_last_atoms
136 If ``True`` the last saved snapshot will be returned.
137 only_prepare
138 If ``True`` the necessary input files will be written
139 but theMD run will not be executed.
141 Returns
142 -------
143 The last snapshot if :attr:`return_last_atoms` is ``True``.
144 """
145 if self._use_temporary_directory:
146 self._make_new_tmp_directory()
148 if self._use_temporary_directory and not return_last_atoms:
149 raise ValueError('Refusing to run in temporary directory '
150 'and not returning atoms; all results will be gone.')
152 if self._use_temporary_directory and only_prepare:
153 raise ValueError('Refusing to only prepare in temporary directory, '
154 'all files will be removed.')
156 # Write files and run
157 FileIOCalculator.write_input(self, self.atoms)
158 self._write_runfile(parameters)
159 write_xyz(filename=os.path.join(self._directory, 'model.xyz'),
160 structure=self.atoms)
162 if only_prepare:
163 return None
165 # Execute the calculation.
166 # Once a new release of ASE is made
167 # (>3.22.1), the below lines can be replaced with
168 # self.execute()
169 command = self.command
170 proc = subprocess.Popen(command, shell=True, cwd=self.directory)
172 errorcode = proc.wait()
173 if errorcode:
174 path = os.path.abspath(self.directory)
175 msg = ('Calculator "{}" failed with command "{}" failed in '
176 '{} with error code {}'.format(self.name, command,
177 path, errorcode))
178 raise RuntimeError(msg)
180 # Extract last snapshot if needed
181 if return_last_atoms:
182 last_atoms = ase_read(os.path.join(self._directory, 'movie.xyz'),
183 format='extxyz', index=-1)
185 if self._use_temporary_directory:
186 self._clean()
188 if return_last_atoms:
189 return last_atoms
190 else:
191 return None
193 def write_input(self, atoms, properties=None, system_changes=None):
194 """
195 Write the input files necessary for a single-point calculation.
196 """
197 if self._use_temporary_directory:
198 self._make_new_tmp_directory()
199 FileIOCalculator.write_input(self, atoms, properties, system_changes)
200 self._write_runfile(parameters=self.single_point_parameters)
201 write_xyz(filename=os.path.join(self._directory, 'model.xyz'),
202 structure=atoms)
204 def _write_runfile(self, parameters):
205 """Write run.in file to define input parameters for MD simulation.
207 Parameters
208 ----------
209 parameters : dict
210 Defines all key-value pairs used in run.in file
211 (see GPUMD documentation for a complete list).
212 Values can be either floats, integers, or lists/tuples.
213 """
214 if len(os.listdir(self._directory)) > 0:
215 warnings.warn(f'{self._directory} is not empty.')
217 with open(os.path.join(self._directory, 'run.in'), 'w') as f:
218 # Custom potential is allowed but normally it can be deduced
219 if 'potential' not in [keyval[0] for keyval in parameters]:
220 f.write(f'potential {self._potential_path} \n')
221 # Write all keywords with parameter(s)
222 for key, val in parameters:
223 f.write(f'{key} ')
224 if isinstance(val, Iterable) and not isinstance(val, str):
225 for v in val:
226 f.write(f'{v} ')
227 else:
228 f.write(f'{val}')
229 f.write('\n')
231 def get_potential_energy_and_stresses_from_file(self):
232 """
233 Extract potential energy (third column of last line in thermo.out) and stresses
234 from thermo.out
235 """
236 data = np.loadtxt(os.path.join(self._directory, 'thermo.out'))
237 if len(data.shape) == 1:
238 line = data
239 else:
240 line = data[-1, :]
242 # Energy
243 energy = line[2]
245 # Stress
246 stress = [v for v in line[3:9]]
247 stress = -GPa * np.array(stress) # to eV/A^3
249 if np.any(np.isnan(stress)) or np.isnan(energy):
250 raise ValueError(f'Failed to extract energy and/or stresses:\n {line}')
251 return energy, stress
253 def _read_potential_energy_and_stresses(self):
254 """Reads potential energy and stresses."""
255 self.results['energy'], self.results['stress'] = \
256 self.get_potential_energy_and_stresses_from_file()
258 def get_forces_from_file(self):
259 """
260 Extract forces (in eV/A) from last snapshot in force.out
261 """
262 data = np.loadtxt(os.path.join(self._directory, 'force.out'))
263 return data[-len(self.atoms):, :]
265 def _read_forces(self):
266 """Reads forces (the last snapshot in force.out) in eV/A"""
267 self.results['forces'] = self.get_forces_from_file()
269 def read_results(self):
270 """
271 Read results from last step of MD calculation.
272 """
273 self._read_potential_energy_and_stresses()
274 self._read_forces()
275 if self._use_temporary_directory:
276 self._clean()
278 def _clean(self):
279 """
280 Remove directory with calculations.
281 """
282 shutil.rmtree(self._directory)
284 def _make_new_tmp_directory(self):
285 """
286 Create a new temporary directory.
287 """
288 # We do not need to create a new temporary directory
289 # if the current one is empty
290 if self._directory is None or \
291 (os.path.isdir(self._directory) and len(os.listdir(self._directory)) > 0):
292 self._directory = tempfile.mkdtemp()
293 self._potential_path = os.path.relpath(os.path.abspath(self.model_filename),
294 self._directory)
296 def set_atoms(self, atoms):
297 """
298 Set Atoms object.
299 Used also when attaching calculator to Atoms object.
300 """
301 self.atoms = atoms
302 self.results = {}
304 def set_directory(self, directory):
305 """
306 Set path to a new directory. This makes it possible to run
307 several calculations with the same calculator while saving
308 all results
309 """
310 self._directory = directory
311 self._use_temporary_directory = False
312 self._potential_path = os.path.relpath(os.path.abspath(self.model_filename),
313 self._directory)