import os
import shutil
import subprocess
import tempfile
import warnings
import numpy as np
from ase import Atoms
from ase.calculators.singlepoint import SinglePointCalculator
from ase.units import GPa
from calorine.env import calorine_getenv
from calorine.nep.io import read_structures, write_nepfile, write_structures
from calorine.nep.model import Model, read_model
from calorine.nep.nep import set_default_cell
from calorine.nep.tensor_conventions import ASE_VOIGT6_ORDER, reduced6_to_full_3x3
# Files that an earlier prediction run in the same directory leaves behind: the
# input files written here and the `*_train.out` files the `nep` executable
# produces, the latter covered by the `.out` suffix. Mirrors
# `_PREVIOUS_RUN_FILES` in calorine/calculators/gpunep.py, but lists the files of
# the `nep` executable rather than those of `gpumd`.
_PREVIOUS_RUN_FILES = ('nep.in', 'nep.txt', 'train.xyz')
def _find_previous_run_files(directory: str) -> list[str]:
"""Return the names of the files in :attr:`directory` that indicate an
earlier prediction run took place there.
The presence of such files, rather than a non-empty directory, is what
signals that results may be read back from an earlier run by mistake. A
directory holding only unrelated files is not reported.
Parameters
----------
directory
Directory to inspect.
Returns
-------
Sorted names of the files that indicate an earlier run, empty if there
are none.
Example
-------
>>> # xdoctest: +SKIP
>>> _find_previous_run_files('some_directory_with_an_energy_train_out_file')
['energy_train.out']
"""
return sorted(filename for filename in os.listdir(directory)
if filename in _PREVIOUS_RUN_FILES or filename.endswith('.out'))
[docs]
def batch_predict_properties(
structures: list[Atoms],
model: str | Model,
command: str | None = None,
directory: str | None = None,
) -> list[Atoms]:
"""Evaluates NEP model properties for a list of structures in a single pass,
using the ``prediction`` mode of the ``nep`` executable (see `here
<https://gpumd.org/nep/input_parameters/prediction.html>`__). This is
substantially faster than evaluating structures one at a time with
:class:`CPUNEP <calorine.calculators.CPUNEP>` or
:class:`GPUNEP <calorine.calculators.GPUNEP>`, since all structures are
transferred to the GPU in a single pass.
Parameters
----------
structures
Structures for which to evaluate properties.
model
Either a path to a NEP model in ``nep.txt`` format, or a
:class:`Model <calorine.nep.model.Model>` object.
command
Command used to invoke the ``nep`` executable.
Default: ``nep``, or the value of the ``CALORINE_NEP_COMMAND``
environment variable if set.
directory
Directory in which to run ``nep``. If ``None``, a temporary directory
is created and removed once the calculation is finished. If
specified, the directory is created if needed and is *not* deleted
afterward, which is useful for debugging or for further analysis of
the raw ``nep`` output files (e.g. via :func:`read_structures
<calorine.nep.read_structures>`).
Returns
-------
list of Atoms
A new list of :class:`Atoms <ase.Atoms>` objects, in the same order
as :attr:`structures`, each with a :class:`SinglePointCalculator
<ase.calculators.singlepoint.SinglePointCalculator>` attached
exposing the predicted properties in the standard way (``energy``,
``forces``, and ``stress``, plus ``charges`` and
``born_effective_charges`` for qNEP models, or
``dipole``/``polarizability`` for TNEP models). The input
:attr:`structures` are not modified.
"""
if len(structures) == 0:
return []
if isinstance(model, Model):
model_obj = model
else:
if not os.path.exists(model):
raise FileNotFoundError(f'{model} does not exist.')
model_obj = read_model(model)
# `training_parameters` carries every nep.in keyword that describes the model, including
# `model_type` and, for a charge-aware model, `charge_mode` and the muNEP head layout
parameters = dict(model_obj.training_parameters)
parameters['prediction'] = 1
model_type_int = parameters['model_type']
# This is the set the dispatch below handles, which is why it is written here rather
# than read from `_MODEL_TYPE_TO_INT`: a model type added to that mapping has to reach
# this guard as a rejection until an arm is written for it.
if model_type_int not in (0, 1, 2):
raise ValueError(f'batch_predict_properties does not support model_type {model_type_int}')
charge_mode = parameters.get('charge_mode', 0)
# If no directory is given, run in a temporary one that is cleaned up
# afterward; otherwise keep the nep.in/nep.txt/train.xyz/*_train.out files
# around for debugging or further analysis, mirroring GPUNEP.
use_temporary_directory = directory is None
if use_temporary_directory:
directory = tempfile.mkdtemp()
else:
os.makedirs(directory, exist_ok=True)
previous_run_files = _find_previous_run_files(directory)
if previous_run_files:
warnings.warn(f'{directory} already contains files from an earlier run: '
f'{", ".join(previous_run_files)}. Results may be read from '
'those rather than from the run about to take place.')
try:
if isinstance(model, Model):
model_obj.write(os.path.join(directory, 'nep.txt'))
else:
shutil.copy2(model, os.path.join(directory, 'nep.txt'))
write_nepfile(os.path.join(directory, 'nep.in'), parameters)
prepared_structures = []
for structure in structures:
prepared = structure.copy()
if prepared.cell.rank == 0:
warnings.warn('Using default unit cell (cubic with side 100 Å).')
set_default_cell(prepared)
forces = np.zeros((len(prepared), 3))
prepared.calc = SinglePointCalculator(prepared, energy=0.0, forces=forces)
if charge_mode:
prepared.arrays['bec'] = np.zeros((len(prepared), 9))
prepared_structures.append(prepared)
with warnings.catch_warnings():
# prepared_structures carry a placeholder zero energy/forces (stress is not
# part of that placeholder), so the resulting warning is expected noise
warnings.filterwarnings(
'ignore', message='Failed to retrieve stresses for structure',
category=UserWarning)
write_structures(os.path.join(directory, 'train.xyz'), prepared_structures)
run_command = command or calorine_getenv('NEP_COMMAND')
try:
completed = subprocess.run(
[run_command], cwd=directory, capture_output=True, text=True)
except OSError as e:
raise RuntimeError(f'Failed to run `{run_command}` in prediction mode: {e}') from e
if completed.returncode != 0:
raise RuntimeError(
f'Failed to run `{run_command}` in prediction mode:\n{completed.stderr}')
predicted_structures, _ = read_structures(directory)
finally:
if use_temporary_directory:
shutil.rmtree(directory)
results_structures = []
for original, predicted in zip(structures, predicted_structures):
natoms = len(original)
results = {}
if model_type_int == 0:
# `energy_predicted` is a per-atom average, not a per-structure total;
# multiply by natoms to match the ASE convention.
results['energy'] = float(predicted.info['energy_predicted'][0]) * natoms
results['forces'] = predicted.arrays['force_predicted']
# read_structures() already converts `stress_predicted` to
# ASE-Voigt order; it's also already normalized by the
# (GPUMD-internal) cell volume and given in GPa, matching the
# convention used for GPUNEP's thermo.out parsing.
results['stress'] = -np.array(predicted.info['stress_predicted']) * GPa
if charge_mode:
results['charges'] = predicted.arrays['charge_predicted'][:, 0]
results['born_effective_charges'] = predicted.arrays['bec_predicted']
elif model_type_int == 1:
results['dipole'] = np.array(predicted.info['dipole_predicted']) * natoms
else:
# Also already converted to ASE-Voigt order by read_structures().
p = np.array(predicted.info['polarizability_predicted']) * natoms
results['polarizability'] = reduced6_to_full_3x3(p, order=ASE_VOIGT6_ORDER)
# `polarizability` is not among ASE's SinglePointCalculator-recognized
# properties, so it has to be added to `results` after construction.
polarizability = results.pop('polarizability', None)
new_atoms = original.copy()
new_atoms.calc = SinglePointCalculator(new_atoms, **results)
if polarizability is not None:
new_atoms.calc.results['polarizability'] = polarizability
results_structures.append(new_atoms)
return results_structures
