from __future__ import annotations
import contextlib
import os
from tempfile import TemporaryFile
from typing import List
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
[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']
# 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
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
----------
structure
Input structure
model_filename
Path to NEP model. Defaults to ``None``.
method
Method for computing gradient with finite differences.
One of 'forward difference' and 'central difference'.
Defaults to 'central difference'
displacement
Displacement in Å to use for finite differences. Defaults to 0.01 Å.
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