refactor: split Model and AtomicModel

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/dp_model.py

@@ -1,14 +1,22 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
+import copy
+
 from deepmd.dpmodel.atomic_model import (
     DPAtomicModel,
 )
 from deepmd.dpmodel.descriptor.base_descriptor import (
     BaseDescriptor,
 )
+from deepmd.dpmodel.fitting.base_fitting import (
+    BaseFitting,
+)
 from deepmd.dpmodel.model.base_model import (
     BaseModel,
 )
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
 
 from .make_model import (
     make_model,
@@ -17,7 +25,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.
@@ -34,3 +42,15 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict):
             global_jdata, local_jdata["descriptor"]
         )
         return local_jdata_cpy
+
+    @classmethod
+    def deserialize(cls, data) -> "DPAtomicModel":
+        data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
+        data.pop("@class")
+        data.pop("type")
+        descriptor_obj = BaseDescriptor.deserialize(data.pop("descriptor"))
+        fitting_obj = BaseFitting.deserialize(data.pop("fitting"))
+        type_map = data.pop("type_map")
+        obj = cls(descriptor_obj, fitting_obj, type_map=type_map, **data)
+        return obj

deepmd/dpmodel/model/make_model.py

@@ -1,21 +1,32 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from abc import (
+    abstractmethod,
+)
 from typing import (
     Dict,
     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 +45,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 +68,13 @@
 
     """
 
-    class CM(T_AtomicModel, NativeOP):
+    class CM(NativeOP, BaseModel):
         def __init__(
             self,
             *args,
             **kwargs,
         ):
-            super().__init__(
-                *args,
-                **kwargs,
-            )
+            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 +216,7 @@
                 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 +385,94 @@
             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
+        @abstractmethod
+        def deserialize(cls, data) -> "CM":
+            pass
+
+        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 @@
         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 @@
             "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).