feat: breaking: backend indepdent definition for dp model

deepmd/model_format/__init__.py

@@ -4,6 +4,12 @@
     PRECISION_DICT,
     NativeOP,
 )
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+from .dp_model import (
+    DPModel,
+)
 from .env_mat import (
     EnvMat,
 )
@@ -37,6 +43,8 @@
 )
 
 __all__ = [
+    "DPModel",
+    "DPAtomicModel",
     "InvarFitting",
     "DescrptSeA",
     "EnvMat",

deepmd/model_format/atomic_model.py

@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from abc import (
+    ABC,
+    abstractclassmethod,
+    abstractmethod,
+)
+from typing import (
+    Dict,
+    List,
+    Optional,
+)
+
+from .output_def import (
+    FittingOutputDef,
+)
+
+
+def make_base_atomic_model(T_Tensor):
+    class BAM(ABC):
+        """Base Atomic Model provides the interfaces of an atomic model."""
+
+        @abstractmethod
+        def get_fitting_output_def(self) -> FittingOutputDef:
+            raise NotImplementedError
+
+        @abstractmethod
+        def get_rcut(self) -> float:
+            raise NotImplementedError
+
+        @abstractmethod
+        def get_sel(self) -> List[int]:
+            raise NotImplementedError
+
+        @abstractmethod
+        def distinguish_types(self) -> bool:
+            raise NotImplementedError
+
+        @abstractmethod
+        def forward_atomic(
+            self,
+            extended_coord: T_Tensor,
+            extended_atype: T_Tensor,
+            nlist: T_Tensor,
+            mapping: Optional[T_Tensor] = None,
+            fparam: Optional[T_Tensor] = None,
+            aparam: Optional[T_Tensor] = None,
+        ) -> Dict[str, T_Tensor]:
+            raise NotImplementedError
+
+        @abstractmethod
+        def serialize(self) -> dict:
+            raise NotImplementedError
+
+        @abstractclassmethod
+        def deserialize(cls):
+            raise NotImplementedError
+
+    return BAM

deepmd/model_format/base_atomic_model.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import numpy as np
+
+from .atomic_model import (
+    make_base_atomic_model,
+)
+
+BaseAtomicModel = make_base_atomic_model(np.ndarray)

deepmd/model_format/common.py

@@ -12,6 +12,8 @@
     "half": np.float16,
     "single": np.float32,
     "double": np.float64,
+    "int32": np.int32,
+    "int64": np.int64,
 }
 DEFAULT_PRECISION = "float64"
 

deepmd/model_format/dp_atomic_model.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import copy
+import sys
+from typing import (
+    Dict,
+    List,
+    Optional,
+)
+
+import numpy as np
+
+from .base_atomic_model import (
+    BaseAtomicModel,
+)
+from .fitting import InvarFitting  # noqa # TODO: should import all fittings!
+from .output_def import (
+    FittingOutputDef,
+)
+from .se_e2_a import DescrptSeA  # noqa # TODO: should import all descriptors!
+
+
+class DPAtomicModel(BaseAtomicModel):
+    """Model give atomic prediction of some physical property.
+
+    Parameters
+    ----------
+    descriptor
+            Descriptor
+    fitting_net
+            Fitting net
+    type_map
+            Mapping atom type to the name (str) of the type.
+            For example `type_map[1]` gives the name of the type 1.
+
+    """
+
+    def __init__(
+        self,
+        descriptor,
+        fitting,
+        type_map: Optional[List[str]] = None,
+    ):
+        super().__init__()
+        self.type_map = type_map
+        self.descriptor = descriptor
+        self.fitting = fitting
+
+    def get_fitting_output_def(self) -> FittingOutputDef:
+        """Get the output def of the fitting net."""
+        return self.fitting.output_def()
+
+    def get_rcut(self) -> float:
+        """Get the cut-off radius."""
+        return self.descriptor.get_rcut()
+
+    def get_sel(self) -> List[int]:
+        """Get the neighbor selection."""
+        return self.descriptor.get_sel()
+
+    def distinguish_types(self) -> bool:
+        """If distinguish different types by sorting."""
+        return self.descriptor.distinguish_types()
+
+    def forward_atomic(
+        self,
+        extended_coord: np.ndarray,
+        extended_atype: np.ndarray,
+        nlist: np.ndarray,
+        mapping: Optional[np.ndarray] = None,
+        fparam: Optional[np.ndarray] = None,
+        aparam: Optional[np.ndarray] = None,
+    ) -> Dict[str, np.ndarray]:
+        """Models' atomic predictions.
+
+        Parameters
+        ----------
+        extended_coord
+            coodinates in extended region
+        extended_atype
+            atomic type in extended region
+        nlist
+            neighbor list. nf x nloc x nsel
+        mapping
+            mapps the extended indices to local indices. nf x nall
+        fparam
+            frame parameter. nf x ndf
+        aparam
+            atomic parameter. nf x nloc x nda
+
+        Returns
+        -------
+        result_dict
+            the result dict, defined by the `FittingOutputDef`.
+
+        """
+        nframes, nloc, nnei = nlist.shape
+        atype = extended_atype[:, :nloc]
+        descriptor, rot_mat, g2, h2, sw = self.descriptor(
+            extended_coord,
+            extended_atype,
+            nlist,
+            mapping=mapping,
+        )
+        ret = self.fitting(
+            descriptor,
+            atype,
+            gr=rot_mat,
+            g2=g2,
+            h2=h2,
+            fparam=fparam,
+            aparam=aparam,
+        )
+        return ret
+
+    def serialize(self) -> dict:
+        return {
+            "type_map": self.type_map,
+            "descriptor": self.descriptor.serialize(),
+            "fitting": self.fitting.serialize(),
+            "descriptor_name": self.descriptor.__class__.__name__,
+            "fitting_name": self.fitting.__class__.__name__,
+        }
+
+    @classmethod
+    def deserialize(cls, data) -> "DPAtomicModel":
+        data = copy.deepcopy(data)
+        descriptor_obj = getattr(
+            sys.modules[__name__], data["descriptor_name"]
+        ).deserialize(data["descriptor"])
+        fitting_obj = getattr(sys.modules[__name__], data["fitting_name"]).deserialize(
+            data["fitting"]
+        )
+        obj = cls(descriptor_obj, fitting_obj, type_map=data["type_map"])
+        return obj

deepmd/model_format/dp_model.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+from .make_model import (
+    make_model,
+)
+
+DPModel = make_model(DPAtomicModel)

deepmd/model_format/fitting.py

@@ -320,7 +320,7 @@ def call(
                     "which is not consistent with {self.numb_fparam}.",
                 )
             fparam = (fparam - self.fparam_avg) * self.fparam_inv_std
-            fparam = np.tile(fparam.reshape([nf, 1, -1]), [1, nloc, 1])
+            fparam = np.tile(fparam.reshape([nf, 1, self.numb_fparam]), [1, nloc, 1])
             xx = np.concatenate(
                 [xx, fparam],
                 axis=-1,
@@ -333,6 +333,7 @@ def call(
                     "get an input aparam of dim {aparam.shape[-1]}, ",
                     "which is not consistent with {self.numb_aparam}.",
                 )
+            aparam = aparam.reshape([nf, nloc, self.numb_aparam])
             aparam = (aparam - self.aparam_avg) * self.aparam_inv_std
             xx = np.concatenate(
                 [xx, aparam],