import numpy as np
from pathlib import Path
from deepmd.infer import DeepPot
from pharmaforge.interfaces.abstractio import AbstractIO
[docs]
class DeepMDInterface(AbstractIO):
""" This class is used to interface with DeepMD-kit for running machine learning potentials, and will
be used to calculate the energies and forces of the system.
.. note::
This interface only takes a model as the option, and thus does not have default options.
Parameters
----------
model : str
The path to the DeepMD model file. This should be a .pb file.
Attributes
----------
calculator : object
The DeepMD calculator object.
model : str
The path to the DeepMD model file.
allow_parallel : bool
Whether to allow parallel calculations. Default is False, ans currently cannot pass the threads to DeepMD.
See Also
--------
pharmaforge.interfaces.abstractio.AbstractIO : The abstract interface class for the Psi4Interface.
deepmd.infer.DeepPot : The deepmd-kit inference calculator object.
"""
def __init__(self, model=None):
self.model = Path(model)
if not self.model.exists():
raise FileNotFoundError(f"The model file {self.model} does not exist.")
self.calculator = DeepPot(self.model)
self.allow_parallel = False
self.level_of_theory = model
self.options_dict = {
"model": self.model,
}
def _calculate(self, structure, charge=0):
"""Calculate the system data.
Parameters
----------
structure : ase.Atoms
The ASE atoms object containing the system data.
Returns
-------
data : DeepMDData
The energy and forces calculated by DeepMD.
"""
xyz = structure.get_positions()
cell = structure.get_cell()
if 0.0 in cell:
cell = np.diag(10 * np.ones(3))
atype_raw = structure.get_chemical_symbols()
atype = self._remap_atom_types(atype_raw)
e, f, v = self.calculator.eval(
xyz.reshape(1,-1),
cell.reshape(1,-1),
atype
)
data = DeepMDData(e[0], f[0])
return data
def _remap_atom_types(self, atype_raw):
"""Remap the atom types to the DeepMD type map.
Parameters
----------
atype_raw : list
The list of atom types in the system.
Returns
-------
atype : list
The remapped atom types corresponding to the DeepMD type map.
"""
type_map = self.calculator.get_type_map()
atype = []
for atom in atype_raw:
if atom in type_map:
loc = type_map.index(atom)
atype.append(loc)
else:
raise ValueError(f"Atom type {atom} not found in type map.")
return atype
[docs]
class DeepMDData:
""" This class is used to store the energy and forces calculated by DeepMD.
This is a simple data class that holds the energy and forces of the system. It is used to
store the results of the DeepMD calculations and is returned by the DeepMDInterface class, to have similar
attributes to outcomes of ASE calculators.
Parameters
----------
e : float
The energy of the system.
f : list
The forces acting on the atoms in the system.
Attributes
----------
e : float
The energy of the system.
f : list
The forces acting on the atoms in the system.
"""
def __init__(self, e, f):
self.e = e
self.f = f
[docs]
def get_potential_energy(self):
return self.e
[docs]
def get_forces(self):
return self.f