Move model deviation and ase calculator to deepmd_utils

deepmd/calculator.py

@@ -1,144 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-"""ASE calculator interface module."""
-
-from pathlib import (
-    Path,
-)
-from typing import (
-    TYPE_CHECKING,
-    ClassVar,
-    Dict,
-    List,
-    Optional,
-    Union,
-)
-
-from ase.calculators.calculator import (
-    Calculator,
-    PropertyNotImplementedError,
-    all_changes,
-)
-
-from deepmd import (
-    DeepPotential,
+from deepmd_utils.calculator import (
+    DP,
 )
 
-if TYPE_CHECKING:
-    from ase import (
-        Atoms,
-    )
-
-__all__ = ["DP"]
-
-
-class DP(Calculator):
-    """Implementation of ASE deepmd calculator.
-
-    Implemented propertie are `energy`, `forces` and `stress`
-
-    Parameters
-    ----------
-    model : Union[str, Path]
-        path to the model
-    label : str, optional
-        calculator label, by default "DP"
-    type_dict : Dict[str, int], optional
-        mapping of element types and their numbers, best left None and the calculator
-        will infer this information from model, by default None
-    neighbor_list : ase.neighborlist.NeighborList, optional
-        The neighbor list object. If None, then build the native neighbor list.
-
-    Examples
-    --------
-    Compute potential energy
-
-    >>> from ase import Atoms
-    >>> from deepmd.calculator import DP
-    >>> water = Atoms('H2O',
-    >>>             positions=[(0.7601, 1.9270, 1),
-    >>>                        (1.9575, 1, 1),
-    >>>                        (1., 1., 1.)],
-    >>>             cell=[100, 100, 100],
-    >>>             calculator=DP(model="frozen_model.pb"))
-    >>> print(water.get_potential_energy())
-    >>> print(water.get_forces())
-
-    Run BFGS structure optimization
-
-    >>> from ase.optimize import BFGS
-    >>> dyn = BFGS(water)
-    >>> dyn.run(fmax=1e-6)
-    >>> print(water.get_positions())
-    """
-
-    name = "DP"
-    implemented_properties: ClassVar[List[str]] = [
-        "energy",
-        "free_energy",
-        "forces",
-        "virial",
-        "stress",
-    ]
-
-    def __init__(
-        self,
-        model: Union[str, "Path"],
-        label: str = "DP",
-        type_dict: Optional[Dict[str, int]] = None,
-        neighbor_list=None,
-        **kwargs,
-    ) -> None:
-        Calculator.__init__(self, label=label, **kwargs)
-        self.dp = DeepPotential(str(Path(model).resolve()), neighbor_list=neighbor_list)
-        if type_dict:
-            self.type_dict = type_dict
-        else:
-            self.type_dict = dict(
-                zip(self.dp.get_type_map(), range(self.dp.get_ntypes()))
-            )
-
-    def calculate(
-        self,
-        atoms: Optional["Atoms"] = None,
-        properties: List[str] = ["energy", "forces", "virial"],
-        system_changes: List[str] = all_changes,
-    ):
-        """Run calculation with deepmd model.
-
-        Parameters
-        ----------
-        atoms : Optional[Atoms], optional
-            atoms object to run the calculation on, by default None
-        properties : List[str], optional
-            unused, only for function signature compatibility,
-            by default ["energy", "forces", "stress"]
-        system_changes : List[str], optional
-            unused, only for function signature compatibility, by default all_changes
-        """
-        if atoms is not None:
-            self.atoms = atoms.copy()
-
-        coord = self.atoms.get_positions().reshape([1, -1])
-        if sum(self.atoms.get_pbc()) > 0:
-            cell = self.atoms.get_cell().reshape([1, -1])
-        else:
-            cell = None
-        symbols = self.atoms.get_chemical_symbols()
-        atype = [self.type_dict[k] for k in symbols]
-        e, f, v = self.dp.eval(coords=coord, cells=cell, atom_types=atype)
-        self.results["energy"] = e[0][0]
-        # see https://gitlab.com/ase/ase/-/merge_requests/2485
-        self.results["free_energy"] = e[0][0]
-        self.results["forces"] = f[0]
-        self.results["virial"] = v[0].reshape(3, 3)
-
-        # convert virial into stress for lattice relaxation
-        if "stress" in properties:
-            if sum(atoms.get_pbc()) > 0:
-                # the usual convention (tensile stress is positive)
-                # stress = -virial / volume
-                stress = -0.5 * (v[0].copy() + v[0].copy().T) / atoms.get_volume()
-                # Voigt notation
-                self.results["stress"] = stress.flat[[0, 4, 8, 5, 2, 1]]
-            else:
-                raise PropertyNotImplementedError
+__all__ = [
+    "DP",
+]