feat: breaking: backend indepdent definition for dp model (#3208)

Features:
- abstract base classes for atomic model, fitting and descriptor.
- dp model format for atomic models
- dp model format for models.
- torch support for atomic model format. 
- torch support `fparam` and `aparam`.

This pr also introduces the following updates:
- support region and nlist in numpy code.
- class decorator like `fitting_check_output` gives human readable class
names.
- support int types in precision dict.
- fix descriptor interfaces.
- refactor torch atomic model impl. introduces dirty hacks to be fixed. 
- provide `format_nlist` that format the nlist in forward_lower method. 

Known limitations:
- torch atomic model has dirty hacks
- interfaces for descriptor, fitting and model statistics was not
considered, should be fixed in future PRs.

Will be fixed
-  [x] dp model module path is a mess to be refactorized.
-  [x] nlist consistency should be checked. if not format nlist.
-  [x] doc strings.
-  [x] `fparam` and `aparam` support.

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

deepmd/dpmodel/__init__.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from .common import (
+    DEFAULT_PRECISION,
+    PRECISION_DICT,
+    NativeOP,
+)
+from .model import (
+    DPAtomicModel,
+    DPModel,
+)
+from .output_def import (
+    FittingOutputDef,
+    ModelOutputDef,
+    OutputVariableDef,
+    fitting_check_output,
+    get_deriv_name,
+    get_reduce_name,
+    model_check_output,
+)
+
+__all__ = [
+    "DPModel",
+    "DPAtomicModel",
+    "PRECISION_DICT",
+    "DEFAULT_PRECISION",
+    "NativeOP",
+    "ModelOutputDef",
+    "FittingOutputDef",
+    "OutputVariableDef",
+    "model_check_output",
+    "fitting_check_output",
+    "get_reduce_name",
+    "get_deriv_name",
+]

deepmd/model_format/common.py → deepmd/dpmodel/common.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from abc import (
     ABC,
+    abstractmethod,
 )
 
 import numpy as np
@@ -12,16 +13,19 @@
     "half": np.float16,
     "single": np.float32,
     "double": np.float64,
+    "int32": np.int32,
+    "int64": np.int64,
 }
 DEFAULT_PRECISION = "float64"
 
 
 class NativeOP(ABC):
     """The unit operation of a native model."""
 
+    @abstractmethod
     def call(self, *args, **kwargs):
         """Forward pass in NumPy implementation."""
-        raise NotImplementedError
+        pass
 
     def __call__(self, *args, **kwargs):
         """Forward pass in NumPy implementation."""

deepmd/dpmodel/descriptor/__init__.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from .make_base_descriptor import (
+    make_base_descriptor,
+)
+from .se_e2_a import (
+    DescrptSeA,
+)
+
+__all__ = [
+    "DescrptSeA",
+    "make_base_descriptor",
+]

deepmd/dpmodel/descriptor/base_descriptor.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import numpy as np
+
+from .make_base_descriptor import (
+    make_base_descriptor,
+)
+
+BaseDescriptor = make_base_descriptor(np.ndarray, "call")

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from abc import (
+    ABC,
+    abstractclassmethod,
+    abstractmethod,
+)
+from typing import (
+    List,
+    Optional,
+)
+
+
+def make_base_descriptor(
+    t_tensor,
+    fwd_method_name: str = "forward",
+):
+    """Make the base class for the descriptor.
+
+    Parameters
+    ----------
+    t_tensor
+        The type of the tensor. used in the type hint.
+    fwd_method_name
+        Name of the forward method. For dpmodels, it should be "call".
+        For torch models, it should be "forward".
+
+    """
+
+    class BD(ABC):
+        """Base descriptor provides the interfaces of descriptor."""
+
+        @abstractmethod
+        def get_rcut(self) -> float:
+            """Returns the cut-off radius."""
+            pass
+
+        @abstractmethod
+        def get_sel(self) -> List[int]:
+            """Returns the number of selected neighboring atoms for each type."""
+            pass
+
+        def get_nsel(self) -> int:
+            """Returns the total number of selected neighboring atoms in the cut-off radius."""
+            return sum(self.get_sel())
+
+        def get_nnei(self) -> int:
+            """Returns the total number of selected neighboring atoms in the cut-off radius."""
+            return self.get_nsel()
+
+        @abstractmethod
+        def get_ntypes(self) -> int:
+            """Returns the number of element types."""
+            pass
+
+        @abstractmethod
+        def get_dim_out(self) -> int:
+            """Returns the output descriptor dimension."""
+            pass
+
+        @abstractmethod
+        def get_dim_emb(self) -> int:
+            """Returns the embedding dimension of g2."""
+            pass
+
+        @abstractmethod
+        def distinguish_types(self) -> bool:
+            """Returns if the descriptor requires a neighbor list that distinguish different
+            atomic types or not.
+            """
+            pass
+
+        @abstractmethod
+        def compute_input_stats(self, merged):
+            """Update mean and stddev for descriptor elements."""
+            pass
+
+        @abstractmethod
+        def init_desc_stat(self, sumr, suma, sumn, sumr2, suma2):
+            """Initialize the model bias by the statistics."""
+            pass
+
+        @abstractmethod
+        def fwd(
+            self,
+            extended_coord,
+            extended_atype,
+            nlist,
+            mapping: Optional[t_tensor] = None,
+        ):
+            """Calculate descriptor."""
+            pass
+
+        @abstractmethod
+        def serialize(self) -> dict:
+            """Serialize the obj to dict."""
+            pass
+
+        @abstractclassmethod
+        def deserialize(cls):
+            """Deserialize from a dict."""
+            pass
+
+    setattr(BD, fwd_method_name, BD.fwd)
+    delattr(BD, "fwd")
+
+    return BD

deepmd/model_format/se_e2_a.py → deepmd/dpmodel/descriptor/se_e2_a.py

@@ -13,20 +13,22 @@
     Optional,
 )
 
-from .common import (
+from deepmd.dpmodel import (
     DEFAULT_PRECISION,
     NativeOP,
 )
-from .env_mat import (
-    EnvMat,
-)
-from .network import (
+from deepmd.dpmodel.utils import (
     EmbeddingNet,
+    EnvMat,
     NetworkCollection,
 )
 
+from .base_descriptor import (
+    BaseDescriptor,
+)
+
 
-class DescrptSeA(NativeOP):
+class DescrptSeA(NativeOP, BaseDescriptor):
     r"""DeepPot-SE constructed from all information (both angular and radial) of
     atomic configurations. The embedding takes the distance between atoms as input.
 
@@ -193,9 +195,43 @@ def __getitem__(self, key):
 
     @property
     def dim_out(self):
+        """Returns the output dimension of this descriptor."""
+        return self.get_dim_out()
+
+    def get_dim_out(self):
         """Returns the output dimension of this descriptor."""
         return self.neuron[-1] * self.axis_neuron
 
+    def get_dim_emb(self):
+        """Returns the embedding (g2) dimension of this descriptor."""
+        return self.neuron[-1]
+
+    def get_rcut(self):
+        """Returns cutoff radius."""
+        return self.rcut
+
+    def get_sel(self):
+        """Returns cutoff radius."""
+        return self.sel
+
+    def distinguish_types(self):
+        """Returns if the descriptor requires a neighbor list that distinguish different
+        atomic types or not.
+        """
+        return True
+
+    def get_ntypes(self) -> int:
+        """Returns the number of element types."""
+        return self.ntypes
+
+    def compute_input_stats(self, merged):
+        """Update mean and stddev for descriptor elements."""
+        raise NotImplementedError
+
+    def init_desc_stat(self, sumr, suma, sumn, sumr2, suma2):
+        """Initialize the model bias by the statistics."""
+        raise NotImplementedError
+
     def cal_g(
         self,
         ss,
@@ -212,6 +248,7 @@ def call(
         coord_ext,
         atype_ext,
         nlist,
+        mapping: Optional[np.ndarray] = None,
     ):
         """Compute the descriptor.
 
@@ -223,6 +260,8 @@ def call(
             The extended aotm types. shape: nf x nall
         nlist
             The neighbor list. shape: nf x nloc x nnei
+        mapping
+            The index mapping from extended to lcoal region. not used by this descriptor.
 
         Returns
         -------
@@ -240,6 +279,7 @@ def call(
         sw
             The smooth switch function.
         """
+        del mapping
         # nf x nloc x nnei x 4
         rr, ww = self.env_mat.call(coord_ext, atype_ext, nlist, self.davg, self.dstd)
         nf, nloc, nnei, _ = rr.shape

deepmd/dpmodel/fitting/__init__.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from .invar_fitting import (
+    InvarFitting,
+)
+from .make_base_fitting import (
+    make_base_fitting,
+)
+
+__all__ = [
+    "InvarFitting",
+    "make_base_fitting",
+]

deepmd/dpmodel/fitting/base_fitting.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import numpy as np
+
+from .make_base_fitting import (
+    make_base_fitting,
+)
+
+BaseFitting = make_base_fitting(np.ndarray, fwd_method_name="call")

deepmd/model_format/fitting.py → deepmd/dpmodel/fitting/invar_fitting.py

@@ -2,30 +2,35 @@
 import copy
 from typing import (
     Any,
+    Dict,
     List,
     Optional,
 )
 
 import numpy as np
 
-from .common import (
+from deepmd.dpmodel import (
     DEFAULT_PRECISION,
     NativeOP,
 )
-from .network import (
-    FittingNet,
-    NetworkCollection,
-)
-from .output_def import (
+from deepmd.dpmodel.output_def import (
     FittingOutputDef,
     OutputVariableDef,
     fitting_check_output,
 )
+from deepmd.dpmodel.utils import (
+    FittingNet,
+    NetworkCollection,
+)
+
+from .base_fitting import (
+    BaseFitting,
+)
 
 
 @fitting_check_output
-class InvarFitting(NativeOP):
-    r"""Fitting the energy (or a porperty of `dim_out`) of the system. The force and the virial can also be trained.
+class InvarFitting(NativeOP, BaseFitting):
+    r"""Fitting the energy (or a rotationally invariant porperty of `dim_out`) of the system. The force and the virial can also be trained.
 
     Lets take the energy fitting task as an example.
     The potential energy :math:`E` is a fitting network function of the descriptor :math:`\mathcal{D}`:
@@ -279,7 +284,7 @@ def call(
         h2: Optional[np.array] = None,
         fparam: Optional[np.array] = None,
         aparam: Optional[np.array] = None,
-    ):
+    ) -> Dict[str, np.array]:
         """Calculate the fitting.
 
         Parameters
@@ -320,7 +325,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 +338,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],