refactor: split Model and AtomicModel (#3438)

Model is not inherited from AtomicModel anymore.

---------

Signed-off-by: Jinzhe Zeng <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

deepmd/dpmodel/atomic_model/make_base_atomic_model.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from abc import (
     ABC,
-    abstractclassmethod,
     abstractmethod,
 )
 from typing import (
@@ -13,6 +12,10 @@
 from deepmd.dpmodel.output_def import (
     FittingOutputDef,
 )
+from deepmd.utils.plugin import (
+    PluginVariant,
+    make_plugin_registry,
+)
 
 
 def make_base_atomic_model(
@@ -31,7 +34,7 @@ def make_base_atomic_model(
 
     """
 
-    class BAM(ABC):
+    class BAM(ABC, PluginVariant, make_plugin_registry("atomic model")):
         """Base Atomic Model provides the interfaces of an atomic model."""
 
         @abstractmethod
@@ -128,8 +131,9 @@ def fwd(
         def serialize(self) -> dict:
             pass
 
-        @abstractclassmethod
-        def deserialize(cls):
+        @classmethod
+        @abstractmethod
+        def deserialize(cls, data: dict):
             pass
 
         def do_grad_r(

deepmd/dpmodel/model/__init__.py

@@ -8,6 +8,8 @@
 according to output variable definition
 `deepmd.dpmodel.OutputVariableDef`.
 
+All models should be inherited from :class:`deepmd.dpmodel.model.base_model.BaseModel`.
+Models generated by `make_model` have already done it.
 """
 
 from .dp_model import (

deepmd/dpmodel/model/dp_model.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
+
 from deepmd.dpmodel.atomic_model import (
     DPAtomicModel,
 )
@@ -17,7 +18,7 @@
 
 # use "class" to resolve "Variable not allowed in type expression"
 @BaseModel.register("standard")
-class DPModel(make_model(DPAtomicModel), BaseModel):
+class DPModel(make_model(DPAtomicModel)):
     @classmethod
     def update_sel(cls, global_jdata: dict, local_jdata: dict):
         """Update the selection and perform neighbor statistics.

deepmd/dpmodel/model/make_model.py

@@ -4,18 +4,26 @@
     List,
     Optional,
     Tuple,
+    Type,
 )
 
 import numpy as np
 
+from deepmd.dpmodel.atomic_model.base_atomic_model import (
+    BaseAtomicModel,
+)
 from deepmd.dpmodel.common import (
     GLOBAL_ENER_FLOAT_PRECISION,
     GLOBAL_NP_FLOAT_PRECISION,
     PRECISION_DICT,
     RESERVED_PRECISON_DICT,
     NativeOP,
 )
+from deepmd.dpmodel.model.base_model import (
+    BaseModel,
+)
 from deepmd.dpmodel.output_def import (
+    FittingOutputDef,
     ModelOutputDef,
     OutputVariableCategory,
     OutputVariableOperation,
@@ -34,7 +42,7 @@
 )
 
 
-def make_model(T_AtomicModel):
+def make_model(T_AtomicModel: Type[BaseAtomicModel]):
     """Make a model as a derived class of an atomic model.
 
     The model provide two interfaces.
@@ -57,16 +65,18 @@ def make_model(T_AtomicModel):
 
     """
 
-    class CM(T_AtomicModel, NativeOP):
+    class CM(NativeOP, BaseModel):
         def __init__(
             self,
             *args,
+            # underscore to prevent conflict with normal inputs
+            atomic_model_: Optional[T_AtomicModel] = None,
             **kwargs,
         ):
-            super().__init__(
-                *args,
-                **kwargs,
-            )
+            if atomic_model_ is not None:
+                self.atomic_model: T_AtomicModel = atomic_model_
+            else:
+                self.atomic_model: T_AtomicModel = T_AtomicModel(*args, **kwargs)
             self.precision_dict = PRECISION_DICT
             self.reverse_precision_dict = RESERVED_PRECISON_DICT
             self.global_np_float_precision = GLOBAL_NP_FLOAT_PRECISION
@@ -208,7 +218,7 @@ def call_lower(
                 extended_coord, fparam=fparam, aparam=aparam
             )
             del extended_coord, fparam, aparam
-            atomic_ret = self.forward_common_atomic(
+            atomic_ret = self.atomic_model.forward_common_atomic(
                 cc_ext,
                 extended_atype,
                 nlist,
@@ -377,4 +387,93 @@ def _format_nlist(
             assert ret.shape[-1] == nnei
             return ret
 
+        def do_grad_r(
+            self,
+            var_name: Optional[str] = None,
+        ) -> bool:
+            """Tell if the output variable `var_name` is r_differentiable.
+            if var_name is None, returns if any of the variable is r_differentiable.
+            """
+            return self.atomic_model.do_grad_r(var_name)
+
+        def do_grad_c(
+            self,
+            var_name: Optional[str] = None,
+        ) -> bool:
+            """Tell if the output variable `var_name` is c_differentiable.
+            if var_name is None, returns if any of the variable is c_differentiable.
+            """
+            return self.atomic_model.do_grad_c(var_name)
+
+        def serialize(self) -> dict:
+            return self.atomic_model.serialize()
+
+        @classmethod
+        def deserialize(cls, data) -> "CM":
+            return cls(atomic_model_=T_AtomicModel.deserialize(data))
+
+        def get_dim_fparam(self) -> int:
+            """Get the number (dimension) of frame parameters of this atomic model."""
+            return self.atomic_model.get_dim_fparam()
+
+        def get_dim_aparam(self) -> int:
+            """Get the number (dimension) of atomic parameters of this atomic model."""
+            return self.atomic_model.get_dim_aparam()
+
+        def get_sel_type(self) -> List[int]:
+            """Get the selected atom types of this model.
+
+            Only atoms with selected atom types have atomic contribution
+            to the result of the model.
+            If returning an empty list, all atom types are selected.
+            """
+            return self.atomic_model.get_sel_type()
+
+        def is_aparam_nall(self) -> bool:
+            """Check whether the shape of atomic parameters is (nframes, nall, ndim).
+
+            If False, the shape is (nframes, nloc, ndim).
+            """
+            return self.atomic_model.is_aparam_nall()
+
+        def get_rcut(self) -> float:
+            """Get the cut-off radius."""
+            return self.atomic_model.get_rcut()
+
+        def get_type_map(self) -> List[str]:
+            """Get the type map."""
+            return self.atomic_model.get_type_map()
+
+        def get_nsel(self) -> int:
+            """Returns the total number of selected neighboring atoms in the cut-off radius."""
+            return self.atomic_model.get_nsel()
+
+        def get_nnei(self) -> int:
+            """Returns the total number of selected neighboring atoms in the cut-off radius."""
+            return self.atomic_model.get_nnei()
+
+        def get_model_def_script(self) -> str:
+            """Get the model definition script."""
+            return self.atomic_model.get_model_def_script()
+
+        def get_sel(self) -> List[int]:
+            """Returns the number of selected atoms for each type."""
+            return self.atomic_model.get_sel()
+
+        def mixed_types(self) -> bool:
+            """If true, the model
+            1. assumes total number of atoms aligned across frames;
+            2. uses a neighbor list that does not distinguish different atomic types.
+
+            If false, the model
+            1. assumes total number of atoms of each atom type aligned across frames;
+            2. uses a neighbor list that distinguishes different atomic types.
+
+            """
+            return self.atomic_model.mixed_types()
+
+        def atomic_output_def(self) -> FittingOutputDef:
+            """Get the output def of the atomic model."""
+            return self.atomic_model.atomic_output_def()
+
     return CM

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -21,6 +21,9 @@
     AtomExcludeMask,
     PairExcludeMask,
 )
+from deepmd.utils.path import (
+    DPPath,
+)
 
 BaseAtomicModel_ = make_base_atomic_model(torch.Tensor)
 
@@ -55,12 +58,6 @@ def reinit_pair_exclude(
         else:
             self.pair_excl = PairExcludeMask(self.get_ntypes(), self.pair_exclude_types)
 
-    # export public methods that are not abstract
-    get_nsel = torch.jit.export(BaseAtomicModel_.get_nsel)
-    get_nnei = torch.jit.export(BaseAtomicModel_.get_nnei)
-    get_ntypes = torch.jit.export(BaseAtomicModel_.get_ntypes)
-
-    @torch.jit.export
     def get_model_def_script(self) -> str:
         return self.model_def_script
 
@@ -126,3 +123,25 @@ def serialize(self) -> dict:
             "atom_exclude_types": self.atom_exclude_types,
             "pair_exclude_types": self.pair_exclude_types,
         }
+
+    def compute_or_load_stat(
+        self,
+        sampled_func,
+        stat_file_path: Optional[DPPath] = None,
+    ):
+        """
+        Compute or load the statistics parameters of the model,
+        such as mean and standard deviation of descriptors or the energy bias of the fitting net.
+        When `sampled` is provided, all the statistics parameters will be calculated (or re-calculated for update),
+        and saved in the `stat_file_path`(s).
+        When `sampled` is not provided, it will check the existence of `stat_file_path`(s)
+        and load the calculated statistics parameters.
+
+        Parameters
+        ----------
+        sampled_func
+            The sampled data frames from different data systems.
+        stat_file_path
+            The path to the statistics files.
+        """
+        raise NotImplementedError

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -223,17 +223,14 @@ def wrapped_sampler():
         if self.fitting_net is not None:
             self.fitting_net.compute_output_stats(wrapped_sampler, stat_file_path)
 
-    @torch.jit.export
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""
         return self.fitting_net.get_dim_fparam()
 
-    @torch.jit.export
     def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this atomic model."""
         return self.fitting_net.get_dim_aparam()
 
-    @torch.jit.export
     def get_sel_type(self) -> List[int]:
         """Get the selected atom types of this model.
 
@@ -243,7 +240,6 @@ def get_sel_type(self) -> List[int]:
         """
         return self.fitting_net.get_sel_type()
 
-    @torch.jit.export
     def is_aparam_nall(self) -> bool:
         """Check whether the shape of atomic parameters is (nframes, nall, ndim).
 

deepmd/pt/model/atomic_model/linear_atomic_model.py

@@ -96,12 +96,10 @@ def mixed_types(self) -> bool:
         """
         return True
 
-    @torch.jit.export
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
         return max(self.get_model_rcuts())
 
-    @torch.jit.export
     def get_type_map(self) -> List[str]:
         """Get the type map."""
         return self.type_map
@@ -292,18 +290,15 @@ def _compute_weight(
         """This should be a list of user defined weights that matches the number of models to be combined."""
         raise NotImplementedError
 
-    @torch.jit.export
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""
         # tricky...
         return max([model.get_dim_fparam() for model in self.models])
 
-    @torch.jit.export
     def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this atomic model."""
         return max([model.get_dim_aparam() for model in self.models])
 
-    @torch.jit.export
     def get_sel_type(self) -> List[int]:
         """Get the selected atom types of this model.
 
@@ -324,7 +319,6 @@ def get_sel_type(self) -> List[int]:
             )
         ).tolist()
 
-    @torch.jit.export
     def is_aparam_nall(self) -> bool:
         """Check whether the shape of atomic parameters is (nframes, nall, ndim).
 

deepmd/pt/model/atomic_model/pairtab_atomic_model.py

@@ -139,11 +139,9 @@ def fitting_output_def(self) -> FittingOutputDef:
             ]
         )
 
-    @torch.jit.export
     def get_rcut(self) -> float:
         return self.rcut
 
-    @torch.jit.export
     def get_type_map(self) -> List[str]:
         return self.type_map
 
@@ -454,17 +452,14 @@ def _calculate_ener(coef: torch.Tensor, uu: torch.Tensor) -> torch.Tensor:
         ener = etmp * uu + a0  # this energy has the extrapolated value when rcut > rmax
         return ener
 
-    @torch.jit.export
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""
         return 0
 
-    @torch.jit.export
     def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this atomic model."""
         return 0
 
-    @torch.jit.export
     def get_sel_type(self) -> List[int]:
         """Get the selected atom types of this model.
 
@@ -474,7 +469,6 @@ def get_sel_type(self) -> List[int]:
         """
         return []
 
-    @torch.jit.export
     def is_aparam_nall(self) -> bool:
         """Check whether the shape of atomic parameters is (nframes, nall, ndim).