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
[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)
if model_obj.model_type in ('potential', 'potential_with_charges'):
model_type_int = 0
elif model_obj.model_type == 'dipole':
model_type_int = 1
elif model_obj.model_type == 'polarizability':
model_type_int = 2
else:
raise ValueError(f'Unknown model_type: {model_obj.model_type}')
charge_mode = model_obj.charge_mode
parameters = dict(model_obj.training_parameters)
parameters['prediction'] = 1
parameters['model_type'] = model_type_int
if charge_mode:
parameters['charge_mode'] = charge_mode
# 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)
if os.listdir(directory):
warnings.warn(f'{directory} is not empty.')
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)
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
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
