Coverage for calorine/tools/prediction.py: 99%
91 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-20 12:52 +0000
1import os
2import shutil
3import subprocess
4import tempfile
5import warnings
6from typing import List, Optional, Union
8import numpy as np
9from ase import Atoms
10from ase.calculators.singlepoint import SinglePointCalculator
11from ase.units import GPa
13from calorine.env import calorine_getenv
14from calorine.nep.io import read_structures, write_nepfile, write_structures
15from calorine.nep.model import Model, read_model
16from calorine.nep.nep import set_default_cell
17from calorine.nep.tensor_conventions import ASE_VOIGT6_ORDER, reduced6_to_full_3x3
20# Files that an earlier prediction run in the same directory leaves behind: the
21# input files written here and the `*_train.out` files the `nep` executable
22# produces, the latter covered by the `.out` suffix. Mirrors
23# `_PREVIOUS_RUN_FILES` in calorine/calculators/gpunep.py, but lists the files of
24# the `nep` executable rather than those of `gpumd`.
25_PREVIOUS_RUN_FILES = ('nep.in', 'nep.txt', 'train.xyz')
28def _find_previous_run_files(directory: str) -> List[str]:
29 """Return the names of the files in :attr:`directory` that indicate an
30 earlier prediction run took place there.
32 The presence of such files, rather than a non-empty directory, is what
33 signals that results may be read back from an earlier run by mistake. A
34 directory holding only unrelated files is not reported.
36 Parameters
37 ----------
38 directory
39 Directory to inspect.
41 Returns
42 -------
43 Sorted names of the files that indicate an earlier run, empty if there
44 are none.
46 Example
47 -------
48 >>> _find_previous_run_files('some_directory_with_an_energy_train_out_file')
49 ['energy_train.out']
50 """
51 return sorted(filename for filename in os.listdir(directory)
52 if filename in _PREVIOUS_RUN_FILES or filename.endswith('.out'))
55def batch_predict_properties(
56 structures: List[Atoms],
57 model: Union[str, Model],
58 command: Optional[str] = None,
59 directory: Optional[str] = None,
60) -> List[Atoms]:
61 """Evaluates NEP model properties for a list of structures in a single pass,
62 using the ``prediction`` mode of the ``nep`` executable (see `here
63 <https://gpumd.org/nep/input_parameters/prediction.html>`__). This is
64 substantially faster than evaluating structures one at a time with
65 :class:`CPUNEP <calorine.calculators.CPUNEP>` or
66 :class:`GPUNEP <calorine.calculators.GPUNEP>`, since all structures are
67 transferred to the GPU in a single pass.
69 Parameters
70 ----------
71 structures
72 Structures for which to evaluate properties.
73 model
74 Either a path to a NEP model in ``nep.txt`` format, or a
75 :class:`Model <calorine.nep.model.Model>` object.
76 command
77 Command used to invoke the ``nep`` executable.
78 Default: ``nep``, or the value of the ``CALORINE_NEP_COMMAND``
79 environment variable if set.
80 directory
81 Directory in which to run ``nep``. If ``None``, a temporary directory
82 is created and removed once the calculation is finished. If
83 specified, the directory is created if needed and is *not* deleted
84 afterward, which is useful for debugging or for further analysis of
85 the raw ``nep`` output files (e.g. via :func:`read_structures
86 <calorine.nep.read_structures>`).
88 Returns
89 -------
90 list of Atoms
91 A new list of :class:`Atoms <ase.Atoms>` objects, in the same order
92 as :attr:`structures`, each with a :class:`SinglePointCalculator
93 <ase.calculators.singlepoint.SinglePointCalculator>` attached
94 exposing the predicted properties in the standard way (``energy``,
95 ``forces``, and ``stress``, plus ``charges`` and
96 ``born_effective_charges`` for qNEP models, or
97 ``dipole``/``polarizability`` for TNEP models). The input
98 :attr:`structures` are not modified.
99 """
100 if len(structures) == 0:
101 return []
103 if isinstance(model, Model):
104 model_obj = model
105 else:
106 if not os.path.exists(model):
107 raise FileNotFoundError(f'{model} does not exist.')
108 model_obj = read_model(model)
110 # `training_parameters` carries every nep.in keyword that describes the model, including
111 # `model_type` and, for a charge-aware model, `charge_mode` and the muNEP head layout
112 parameters = dict(model_obj.training_parameters)
113 parameters['prediction'] = 1
114 model_type_int = parameters['model_type']
115 charge_mode = parameters.get('charge_mode', 0)
117 # If no directory is given, run in a temporary one that is cleaned up
118 # afterward; otherwise keep the nep.in/nep.txt/train.xyz/*_train.out files
119 # around for debugging or further analysis, mirroring GPUNEP.
120 use_temporary_directory = directory is None
121 if use_temporary_directory:
122 directory = tempfile.mkdtemp()
123 else:
124 os.makedirs(directory, exist_ok=True)
125 previous_run_files = _find_previous_run_files(directory)
126 if previous_run_files:
127 warnings.warn(f'{directory} already contains files from an earlier run: '
128 f'{", ".join(previous_run_files)}. Results may be read from '
129 'those rather than from the run about to take place.')
131 try:
132 if isinstance(model, Model):
133 model_obj.write(os.path.join(directory, 'nep.txt'))
134 else:
135 shutil.copy2(model, os.path.join(directory, 'nep.txt'))
136 write_nepfile(parameters, directory)
138 prepared_structures = []
139 for structure in structures:
140 prepared = structure.copy()
141 if prepared.cell.rank == 0:
142 warnings.warn('Using default unit cell (cubic with side 100 Å).')
143 set_default_cell(prepared)
144 forces = np.zeros((len(prepared), 3))
145 prepared.calc = SinglePointCalculator(prepared, energy=0.0, forces=forces)
146 if charge_mode:
147 prepared.arrays['bec'] = np.zeros((len(prepared), 9))
148 prepared_structures.append(prepared)
149 with warnings.catch_warnings():
150 # prepared_structures carry a placeholder zero energy/forces (stress is not
151 # part of that placeholder), so the resulting warning is expected noise
152 warnings.filterwarnings(
153 'ignore', message='Failed to retrieve stresses for structure',
154 category=UserWarning)
155 write_structures(os.path.join(directory, 'train.xyz'), prepared_structures)
157 run_command = command or calorine_getenv('NEP_COMMAND')
158 try:
159 completed = subprocess.run(
160 [run_command], cwd=directory, capture_output=True, text=True)
161 except OSError as e:
162 raise RuntimeError(f'Failed to run `{run_command}` in prediction mode: {e}') from e
163 if completed.returncode != 0:
164 raise RuntimeError(
165 f'Failed to run `{run_command}` in prediction mode:\n{completed.stderr}')
167 with warnings.catch_warnings():
168 # prediction mode never writes a test.xyz; the resulting warning is expected noise
169 warnings.filterwarnings(
170 'ignore', message=r'File .*test\.xyz not found\.', category=UserWarning)
171 predicted_structures, _ = read_structures(directory)
172 finally:
173 if use_temporary_directory:
174 shutil.rmtree(directory)
176 results_structures = []
177 for original, predicted in zip(structures, predicted_structures):
178 natoms = len(original)
179 results = {}
180 if model_type_int == 0:
181 # `energy_predicted` is a per-atom average, not a per-structure total;
182 # multiply by natoms to match the ASE convention.
183 results['energy'] = float(predicted.info['energy_predicted'][0]) * natoms
184 results['forces'] = predicted.arrays['force_predicted']
185 # read_structures() already converts `stress_predicted` to
186 # ASE-Voigt order; it's also already normalized by the
187 # (GPUMD-internal) cell volume and given in GPa, matching the
188 # convention used for GPUNEP's thermo.out parsing.
189 results['stress'] = -np.array(predicted.info['stress_predicted']) * GPa
190 if charge_mode:
191 results['charges'] = predicted.arrays['charge_predicted'][:, 0]
192 results['born_effective_charges'] = predicted.arrays['bec_predicted']
193 elif model_type_int == 1:
194 results['dipole'] = np.array(predicted.info['dipole_predicted']) * natoms
195 elif model_type_int == 2: 195 ↛ 202line 195 didn't jump to line 202 because the condition on line 195 was always true
196 # Also already converted to ASE-Voigt order by read_structures().
197 p = np.array(predicted.info['polarizability_predicted']) * natoms
198 results['polarizability'] = reduced6_to_full_3x3(p, order=ASE_VOIGT6_ORDER)
200 # `polarizability` is not among ASE's SinglePointCalculator-recognized
201 # properties, so it has to be added to `results` after construction.
202 polarizability = results.pop('polarizability', None)
203 new_atoms = original.copy()
204 new_atoms.calc = SinglePointCalculator(new_atoms, **results)
205 if polarizability is not None:
206 new_atoms.calc.results['polarizability'] = polarizability
207 results_structures.append(new_atoms)
209 return results_structures