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