Coverage for calorine/tools/prediction.py: 98%
96 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-23 12:50 +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
20def batch_predict_properties(
21 structures: List[Atoms],
22 model: Union[str, Model],
23 command: Optional[str] = None,
24 directory: Optional[str] = None,
25) -> List[Atoms]:
26 """Evaluates NEP model properties for a list of structures in a single pass,
27 using the ``prediction`` mode of the ``nep`` executable (see `here
28 <https://gpumd.org/nep/input_parameters/prediction.html>`__). This is
29 substantially faster than evaluating structures one at a time with
30 :class:`CPUNEP <calorine.calculators.CPUNEP>` or
31 :class:`GPUNEP <calorine.calculators.GPUNEP>`, since all structures are
32 transferred to the GPU in a single pass.
34 Parameters
35 ----------
36 structures
37 Structures for which to evaluate properties.
38 model
39 Either a path to a NEP model in ``nep.txt`` format, or a
40 :class:`Model <calorine.nep.model.Model>` object.
41 command
42 Command used to invoke the ``nep`` executable.
43 Default: ``nep``, or the value of the ``CALORINE_NEP_COMMAND``
44 environment variable if set.
45 directory
46 Directory in which to run ``nep``. If ``None``, a temporary directory
47 is created and removed once the calculation is finished. If
48 specified, the directory is created if needed and is *not* deleted
49 afterward, which is useful for debugging or for further analysis of
50 the raw ``nep`` output files (e.g. via :func:`read_structures
51 <calorine.nep.read_structures>`).
53 Returns
54 -------
55 list of Atoms
56 A new list of :class:`Atoms <ase.Atoms>` objects, in the same order
57 as :attr:`structures`, each with a :class:`SinglePointCalculator
58 <ase.calculators.singlepoint.SinglePointCalculator>` attached
59 exposing the predicted properties in the standard way (``energy``,
60 ``forces``, and ``stress``, plus ``charges`` and
61 ``born_effective_charges`` for qNEP models, or
62 ``dipole``/``polarizability`` for TNEP models). The input
63 :attr:`structures` are not modified.
64 """
65 if len(structures) == 0:
66 return []
68 if isinstance(model, Model):
69 model_obj = model
70 else:
71 if not os.path.exists(model):
72 raise FileNotFoundError(f'{model} does not exist.')
73 model_obj = read_model(model)
75 if model_obj.model_type in ('potential', 'potential_with_charges'):
76 model_type_int = 0
77 elif model_obj.model_type == 'dipole':
78 model_type_int = 1
79 elif model_obj.model_type == 'polarizability': 79 ↛ 82line 79 didn't jump to line 82 because the condition on line 79 was always true
80 model_type_int = 2
81 else:
82 raise ValueError(f'Unknown model_type: {model_obj.model_type}')
83 charge_mode = model_obj.charge_mode
85 parameters = dict(model_obj.training_parameters)
86 parameters['prediction'] = 1
87 parameters['model_type'] = model_type_int
88 if charge_mode:
89 parameters['charge_mode'] = charge_mode
91 # If no directory is given, run in a temporary one that is cleaned up
92 # afterward; otherwise keep the nep.in/nep.txt/train.xyz/*_train.out files
93 # around for debugging or further analysis, mirroring GPUNEP.
94 use_temporary_directory = directory is None
95 if use_temporary_directory:
96 directory = tempfile.mkdtemp()
97 else:
98 os.makedirs(directory, exist_ok=True)
99 if os.listdir(directory):
100 warnings.warn(f'{directory} is not empty.')
102 try:
103 if isinstance(model, Model):
104 model_obj.write(os.path.join(directory, 'nep.txt'))
105 else:
106 shutil.copy2(model, os.path.join(directory, 'nep.txt'))
107 write_nepfile(parameters, directory)
109 prepared_structures = []
110 for structure in structures:
111 prepared = structure.copy()
112 if prepared.cell.rank == 0:
113 warnings.warn('Using default unit cell (cubic with side 100 Å).')
114 set_default_cell(prepared)
115 forces = np.zeros((len(prepared), 3))
116 prepared.calc = SinglePointCalculator(prepared, energy=0.0, forces=forces)
117 if charge_mode:
118 prepared.arrays['bec'] = np.zeros((len(prepared), 9))
119 prepared_structures.append(prepared)
120 with warnings.catch_warnings():
121 # prepared_structures carry a placeholder zero energy/forces (stress is not
122 # part of that placeholder), so the resulting warning is expected noise
123 warnings.filterwarnings(
124 'ignore', message='Failed to retrieve stresses for structure',
125 category=UserWarning)
126 write_structures(os.path.join(directory, 'train.xyz'), prepared_structures)
128 run_command = command or calorine_getenv('NEP_COMMAND')
129 try:
130 completed = subprocess.run(
131 [run_command], cwd=directory, capture_output=True, text=True)
132 except OSError as e:
133 raise RuntimeError(f'Failed to run `{run_command}` in prediction mode: {e}') from e
134 if completed.returncode != 0:
135 raise RuntimeError(
136 f'Failed to run `{run_command}` in prediction mode:\n{completed.stderr}')
138 with warnings.catch_warnings():
139 # prediction mode never writes a test.xyz; the resulting warning is expected noise
140 warnings.filterwarnings(
141 'ignore', message=r'File .*test\.xyz not found\.', category=UserWarning)
142 predicted_structures, _ = read_structures(directory)
143 finally:
144 if use_temporary_directory:
145 shutil.rmtree(directory)
147 results_structures = []
148 for original, predicted in zip(structures, predicted_structures):
149 natoms = len(original)
150 results = {}
151 if model_type_int == 0:
152 # `energy_predicted` is a per-atom average, not a per-structure total;
153 # multiply by natoms to match the ASE convention.
154 results['energy'] = float(predicted.info['energy_predicted'][0]) * natoms
155 results['forces'] = predicted.arrays['force_predicted']
156 # read_structures() already converts `stress_predicted` to
157 # ASE-Voigt order; it's also already normalized by the
158 # (GPUMD-internal) cell volume and given in GPa, matching the
159 # convention used for GPUNEP's thermo.out parsing.
160 results['stress'] = -np.array(predicted.info['stress_predicted']) * GPa
161 if charge_mode:
162 results['charges'] = predicted.arrays['charge_predicted'][:, 0]
163 results['born_effective_charges'] = predicted.arrays['bec_predicted']
164 elif model_type_int == 1:
165 results['dipole'] = np.array(predicted.info['dipole_predicted']) * natoms
166 elif model_type_int == 2: 166 ↛ 173line 166 didn't jump to line 173 because the condition on line 166 was always true
167 # Also already converted to ASE-Voigt order by read_structures().
168 p = np.array(predicted.info['polarizability_predicted']) * natoms
169 results['polarizability'] = reduced6_to_full_3x3(p, order=ASE_VOIGT6_ORDER)
171 # `polarizability` is not among ASE's SinglePointCalculator-recognized
172 # properties, so it has to be added to `results` after construction.
173 polarizability = results.pop('polarizability', None)
174 new_atoms = original.copy()
175 new_atoms.calc = SinglePointCalculator(new_atoms, **results)
176 if polarizability is not None:
177 new_atoms.calc.results['polarizability'] = polarizability
178 results_structures.append(new_atoms)
180 return results_structures