import os
import shutil
import subprocess
import tempfile
import warnings
from typing import List, Optional, Union
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
-------
>>> _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: Union[str, Model],
command: Optional[str] = None,
directory: Optional[str] = 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']
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(parameters, directory)
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}')
with warnings.catch_warnings():
# prediction mode never writes a test.xyz; the resulting warning is expected noise
warnings.filterwarnings(
'ignore', message=r'File .*test\.xyz not found\.', category=UserWarning)
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
elif model_type_int == 2:
# 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
