from __future__ import annotations
import contextlib
import os
from tempfile import TemporaryFile
from typing import List, Union
import numpy as np
from ase import Atoms
from ase.calculators.calculator import Calculator, all_changes
from ase.stress import full_3x3_to_voigt_6_stress
import _nepy
from calorine.nep.model import _get_nep_contents
from calorine.nep.nep import _check_components_polarizability_gradient, \
_polarizability_gradient_to_3x3
[docs]
class CPUNEP(Calculator):
"""This class provides an ASE calculator for `nep_cpu`,
the in-memory CPU implementation of GPUMD.
Parameters
----------
model_filename : str
Path to file in ``nep.txt`` format with model parameters
atoms : Atoms
Atoms to attach the calculator to
label : str
Label for this calclator
debug : bool, optional
Flag to toggle debug mode. Prints GPUMD output. Defaults to False.
Raises
------
FileNotFoundError
Raises :class:`FileNotFoundError` if :attr:`model_filename` does not point to a valid file.
ValueError
Raises :class:`ValueError` atoms are not defined when trying to get energies and forces.
Example
-------
>>> calc = CPUNEP('nep.txt')
>>> atoms.calc = calc
>>> atoms.get_potential_energy()
"""
implemented_properties = [
'energy',
'energies',
'forces',
'stress',
]
debug = False
nepy = None
def __init__(
self,
model_filename: str,
atoms: Atoms | None = None,
label: str | None = None,
debug: bool = False,
):
self.debug = debug
if not os.path.exists(model_filename):
raise FileNotFoundError(f'{model_filename} does not exist.')
self.model_filename = str(model_filename)
# Get model type from first row in nep.txt
header, _ = _get_nep_contents(self.model_filename)
self.model_type = header['model_type']
self.supported_species = set(header['types'])
self.nep_version = header['version']
if self.model_type == 'dipole':
# Only available for dipole models
self.implemented_properties = ['dipole']
elif self.model_type == 'polarizability':
# Only available for polarizability models
self.implemented_properties = ['polarizability']
# Initialize atoms, results and nepy - note that this is also done in Calculator.__init__()
if atoms is not None:
self.set_atoms(atoms)
parameters = {'model_filename': model_filename}
Calculator.__init__(self, label=label, atoms=atoms, **parameters)
if atoms is not None:
self._setup_nepy()
def __str__(self) -> str:
def indent(s: str, i: int) -> str:
s = '\n'.join([i * ' ' + line for line in s.split('\n')])
return s
parameters = '\n'.join(
[f'{key}: {value}' for key, value in self.parameters.items()]
)
parameters = indent(parameters, 4)
using_debug = '\nIn debug mode' if self.debug else ''
s = f'{self.__class__.__name__}\n{parameters}{using_debug}'
return s
def _setup_nepy(self):
"""
Creates an instance of the NEPY class and attaches it to the calculator object.
The output from `nep.cpp` is only written to STDOUT if debug == True
"""
if self.atoms is None:
raise ValueError('Atoms must be defined to get energies and forces.')
if self.atoms.cell.rank == 0:
raise ValueError('Atoms must have a defined cell.')
natoms = len(self.atoms)
self.natoms = natoms
c = self.atoms.get_cell(complete=True).flatten()
cell = [c[0], c[3], c[6], c[1], c[4], c[7], c[2], c[5], c[8]]
symbols = self.atoms.get_chemical_symbols()
positions = list(
self.atoms.get_positions().T.flatten()
) # [x1, ..., xN, y1, ... yN,...]
masses = self.atoms.get_masses()
# Disable output from C++ code by default
if self.debug:
self.nepy = _nepy.NEPY(
self.model_filename, self.natoms, cell, symbols, positions, masses
)
else:
with TemporaryFile('w') as f:
with contextlib.redirect_stdout(f):
self.nepy = _nepy.NEPY(
self.model_filename,
self.natoms,
cell,
symbols,
positions,
masses,
)
[docs]
def set_atoms(self, atoms: Atoms):
"""Updates the Atoms object.
Parameters
----------
atoms : Atoms
Atoms to attach the calculator to
"""
species_in_atoms_object = set(np.unique(atoms.get_chemical_symbols()))
if not species_in_atoms_object.issubset(self.supported_species):
raise ValueError('Structure contains species that are not supported by the NEP model.')
self.atoms = atoms
self.results = {}
self.nepy = None
def _update_symbols(self):
"""Update atom symbols in NEPY."""
symbols = self.atoms.get_chemical_symbols()
self.nepy.set_symbols(symbols)
def _update_masses(self):
"""Update atom masses in NEPY"""
masses = self.atoms.get_masses()
self.nepy.set_masses(masses)
def _update_cell(self):
"""Update cell parameters in NEPY."""
c = self.atoms.get_cell(complete=True).flatten()
cell = [c[0], c[3], c[6], c[1], c[4], c[7], c[2], c[5], c[8]]
self.nepy.set_cell(cell)
def _update_positions(self):
"""Update atom positions in NEPY."""
positions = list(
self.atoms.get_positions().T.flatten()
) # [x1, ..., xN, y1, ... yN,...]
self.nepy.set_positions(positions)
def calculate(
self,
atoms: Atoms = None,
properties: List[str] = None,
system_changes: List[str] = all_changes,
):
"""Calculate energy, per atom energies, forces, stress and dipole.
Parameters
----------
atoms : Atoms, optional
System for which to calculate properties, by default None
properties : List[str], optional
Properties to calculate, by default None
system_changes : List[str], optional
Changes to the system since last call, by default all_changes
"""
if properties is None:
properties = self.implemented_properties
Calculator.calculate(self, atoms, properties, system_changes)
if self.nepy is None:
# Create new NEPY interface
self._setup_nepy()
# Update existing NEPY interface
for change in system_changes:
if change == 'positions':
self._update_positions()
elif change == 'numbers':
self._update_symbols()
self._update_masses()
elif change == 'cell':
self._update_cell()
if 'dipole' in self.implemented_properties:
dipole = np.array(self.nepy.get_dipole())
self.results['dipole'] = dipole
elif 'polarizability' in self.implemented_properties:
pol = np.array(self.nepy.get_polarizability())
polarizability = np.array([
[pol[0], pol[3], pol[5]],
[pol[3], pol[1], pol[4]],
[pol[5], pol[4], pol[2]]
])
self.results['polarizability'] = polarizability
else:
energies, forces, virials = self.nepy.get_potential_forces_and_virials()
energies_per_atom = np.array(energies)
energy = energies_per_atom.sum()
forces_per_atom = np.array(forces).reshape(-1, self.natoms).T
virials_per_atom = np.array(virials).reshape(-1, self.natoms).T
stress = -(np.sum(virials_per_atom, axis=0) / self.atoms.get_volume()).reshape((3, 3))
stress = full_3x3_to_voigt_6_stress(stress)
self.results['energy'] = energy
self.results['forces'] = forces_per_atom
self.results['stress'] = stress
def get_dipole_gradient(
self,
displacement: float = 0.01,
method: str = 'central difference',
charge: float = 1.0,
):
"""Calculates the dipole gradient using finite differences.
Parameters
----------
displacement
Displacement in Å to use for finite differences. Defaults to 0.01 Å.
method
Method for computing gradient with finite differences.
One of 'forward difference' and 'central difference'.
Defaults to 'central difference'
charge
System charge in units of the elemental charge.
Used for correcting the dipoles before computing the gradient.
Defaults to 1.0.
Returns
------- dipole gradient with shape `(N, 3, 3)`
"""
if 'dipole' not in self.implemented_properties:
raise ValueError('Dipole gradients are only defined for dipole NEP models.')
if displacement <= 0:
raise ValueError('displacement must be > 0 Å')
implemented_methods = {
'forward difference': 0,
'central difference': 1,
'second order central difference': 2,
}
if method not in implemented_methods.keys():
raise ValueError(f'Invalid method {method} for calculating gradient')
if self.nepy is None:
# Create new NEPY interface
self._setup_nepy()
dipole_gradient = np.array(
self.nepy.get_dipole_gradient(
displacement, implemented_methods[method], charge
)
).reshape(self.natoms, 3, 3)
return dipole_gradient
def get_polarizability(
self,
atoms: Atoms = None,
properties: List[str] = None,
system_changes: List[str] = all_changes,
) -> np.ndarray:
"""Calculates the polarizability tensor for the current structure.
The model must have been trained to predict the polarizability.
This is a wrapper function for :func:`calculate`.
Parameters
----------
atoms : Atoms, optional
System for which to calculate properties, by default None
properties : List[str], optional
Properties to calculate, by default None
system_changes : List[str], optional
Changes to the system since last call, by default all_changes
Returns
-------
polarizability with shape ``(3, 3)``
"""
if properties is None:
properties = self.implemented_properties
if 'polarizability' not in properties:
raise ValueError('Polarizability is only defined for polarizability NEP models.')
self.calculate(atoms, properties, system_changes)
return self.results['polarizability']
def get_polarizability_gradient(
self,
displacement: float = 0.01,
component: Union[str, List[str]] = 'full',
) -> np.ndarray:
"""Calculates the dipole gradient for a given structure using finite differences.
This function computes the derivatives using the second-order central difference
method with a C++ backend.
Parameters
----------
displacement
Displacement in Å to use for finite differences. Defaults to ``0.01``.
component
Component or components of the polarizability tensor that the gradient
should be computed for. The following components are available:
``x`, ``y``, ``z``, ``full``
Option ``full`` computes the derivative whilst moving the atoms in each Cartesian
direction, which yields a tensor of shape (N, 3, 3, 3).
Multiple components may be specified.
Defaults to ``full``.
Returns
-------
polarizability gradient with shape ``(N, C, 3, 3)`` where ``C``
are the number of components chosen.
"""
if 'polarizability' not in self.implemented_properties:
raise ValueError('Polarizability gradients are only defined'
' for polarizability NEP models.')
if displacement <= 0:
raise ValueError('displacement must be > 0 Å')
if self.nepy is None:
# Create new NEPY interface
self._setup_nepy()
component_array = _check_components_polarizability_gradient(component)
pg = np.array(
self.nepy.get_polarizability_gradient(
displacement, component_array
)
).reshape(self.natoms, 3, 6)
polarizability_gradient = _polarizability_gradient_to_3x3(self.natoms, pg)
return polarizability_gradient[:, component_array, :, :]