Chore: refactor dpmodel

deepmd/pt/model/atomic_model/__init__.py

@@ -17,6 +17,9 @@
 from .base_atomic_model import (
     BaseAtomicModel,
 )
+from .dipole_atomic_model import (
+    DPDipoleAtomicModel,
+)
 from .dp_atomic_model import (
     DPAtomicModel,
 )
@@ -27,11 +30,16 @@
 from .pairtab_atomic_model import (
     PairTabAtomicModel,
 )
+from .polar_atomic_model import (
+    DPPolarAtomicModel,
+)
 
 __all__ = [
     "BaseAtomicModel",
     "DPAtomicModel",
     "PairTabAtomicModel",
     "LinearEnergyAtomicModel",
+    "DPPolarAtomicModel",
+    "DPDipoleAtomicModel",
     "DPZBLLinearEnergyAtomicModel",
 ]

deepmd/pt/model/atomic_model/dipole_atomic_model.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+
+class DPDipoleAtomicModel(DPAtomicModel):
+    def apply_out_stat(self, ret, atype):
+        # dipole not applying bias
+        pass

deepmd/pt/model/atomic_model/polar_atomic_model.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+
+class DPPolarAtomicModel(DPAtomicModel):
+    def apply_out_stat(self, ret, atype):
+        # TODO: migrate bias
+        pass

deepmd/pt/model/model/__init__.py

@@ -30,6 +30,12 @@
     Spin,
 )
 
+from .dipole_model import (
+    DipoleModel,
+)
+from .dos_model import (
+    DOSModel,
+)
 from .dp_model import (
     DPModel,
 )
@@ -51,6 +57,9 @@
 from .model import (
     BaseModel,
 )
+from .polar_model import (
+    PolarModel,
+)
 from .spin_model import (
     SpinEnergyModel,
     SpinModel,
@@ -160,7 +169,18 @@
     atom_exclude_types = model_params.get("atom_exclude_types", [])
     pair_exclude_types = model_params.get("pair_exclude_types", [])
 
-    model = DPModel(
+    if fitting_net["type"] == "dipole":
+        modelcls = DipoleModel
+    elif fitting_net["type"] == "polar":
+        modelcls = PolarModel
+    elif fitting_net["type"] == "dos":
+        modelcls = DOSModel
+    elif fitting_net["type"] == "ener":
+        modelcls = EnergyModel
+    else:
+        pass
+
+    model = modelcls(
         descriptor=descriptor,
         fitting=fitting,
         type_map=model_params["type_map"],

deepmd/pt/model/model/dipole_model.py

@@ -6,12 +6,23 @@
 
 import torch
 
+from deepmd.pt.model.atomic_model import (
+    DPDipoleAtomicModel,
+)
+from deepmd.pt.model.model.model import (
+    BaseModel,
+)
+
 from .dp_model import (
     DPModel,
 )
+from .make_model import (
+    make_model,
+)
 
 
-class DipoleModel(DPModel):
+@BaseModel.register("standard")
+class DipoleModel(DPModel, make_model(DPDipoleAtomicModel)):
     model_type = "dipole"
 
     def __init__(
@@ -57,6 +68,14 @@
             model_predict["updated_coord"] += coord
         return model_predict
 
+    def get_fitting_net(self):
+        """Get the fitting network."""
+        return self.atomic_model.fitting_net
+
+    def get_descriptor(self):
+        """Get the descriptor."""
+        return self.atomic_model.descriptor
+
     @torch.jit.export
     def forward_lower(
         self,

deepmd/pt/model/model/dos_model.py

@@ -6,12 +6,23 @@
 
 import torch
 
+from deepmd.pt.model.atomic_model import (
+    DPAtomicModel,
+)
+from deepmd.pt.model.model.model import (
+    BaseModel,
+)
+
 from .dp_model import (
     DPModel,
 )
+from .make_model import (
+    make_model,
+)
 
 
-class DOSModel(DPModel):
+@BaseModel.register("standard")
+class DOSModel(DPModel, make_model(DPAtomicModel)):
     model_type = "dos"
 
     def __init__(
@@ -50,6 +61,14 @@
             model_predict["updated_coord"] += coord
         return model_predict
 
+    def get_fitting_net(self):
+        """Get the fitting network."""
+        return self.atomic_model.fitting_net
+
+    def get_descriptor(self):
+        """Get the descriptor."""
+        return self.atomic_model.descriptor
+
     @torch.jit.export
     def get_numb_dos(self) -> int:
         """Get the number of  DOS for DOSFittingNet."""

deepmd/pt/model/model/dp_model.py

@@ -1,83 +1,11 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from typing import (
-    Dict,
-    Optional,
-)
-
-import torch
-
-from deepmd.pt.model.atomic_model import (
-    DPAtomicModel,
-)
 from deepmd.pt.model.descriptor.base_descriptor import (
     BaseDescriptor,
 )
-from deepmd.pt.model.model.model import (
-    BaseModel,
-)
-from deepmd.pt.model.task.dipole import (
-    DipoleFittingNet,
-)
-from deepmd.pt.model.task.dos import (
-    DOSFittingNet,
-)
-from deepmd.pt.model.task.ener import (
-    EnergyFittingNet,
-    EnergyFittingNetDirect,
-)
-from deepmd.pt.model.task.polarizability import (
-    PolarFittingNet,
-)
 
-from .make_model import (
-    make_model,
-)
 
-
-@BaseModel.register("standard")
-class DPModel(make_model(DPAtomicModel)):
-    def __new__(
-        cls,
-        descriptor=None,
-        fitting=None,
-        *args,
-        # disallow positional atomic_model_
-        atomic_model_: Optional[DPAtomicModel] = None,
-        **kwargs,
-    ):
-        from deepmd.pt.model.model.dipole_model import (
-            DipoleModel,
-        )
-        from deepmd.pt.model.model.dos_model import (
-            DOSModel,
-        )
-        from deepmd.pt.model.model.ener_model import (
-            EnergyModel,
-        )
-        from deepmd.pt.model.model.polar_model import (
-            PolarModel,
-        )
-
-        if atomic_model_ is not None:
-            fitting = atomic_model_.fitting_net
-        else:
-            assert fitting is not None, "fitting network is not provided"
-
-        # according to the fitting network to decide the type of the model
-        if cls is DPModel:
-            # map fitting to model
-            if isinstance(fitting, EnergyFittingNet) or isinstance(
-                fitting, EnergyFittingNetDirect
-            ):
-                cls = EnergyModel
-            elif isinstance(fitting, DipoleFittingNet):
-                cls = DipoleModel
-            elif isinstance(fitting, PolarFittingNet):
-                cls = PolarModel
-            elif isinstance(fitting, DOSFittingNet):
-                cls = DOSModel
-            # else: unknown fitting type, fall back to DPModel
-        return super().__new__(cls)
+class DPModel:
+    """A base class to implement common methods for all the Models."""
 
     @classmethod
     def update_sel(cls, global_jdata: dict, local_jdata: dict):
@@ -95,30 +23,3 @@
             global_jdata, local_jdata["descriptor"]
         )
         return local_jdata_cpy
-
-    def get_fitting_net(self):
-        """Get the fitting network."""
-        return self.atomic_model.fitting_net
-
-    def get_descriptor(self):
-        """Get the descriptor."""
-        return self.atomic_model.descriptor
-
-    def forward(
-        self,
-        coord,
-        atype,
-        box: Optional[torch.Tensor] = None,
-        fparam: Optional[torch.Tensor] = None,
-        aparam: Optional[torch.Tensor] = None,
-        do_atomic_virial: bool = False,
-    ) -> Dict[str, torch.Tensor]:
-        # directly call the forward_common method when no specific transform rule
-        return self.forward_common(
-            coord,
-            atype,
-            box,
-            fparam=fparam,
-            aparam=aparam,
-            do_atomic_virial=do_atomic_virial,
-        )

deepmd/pt/model/model/dp_zbl_model.py

@@ -6,16 +6,16 @@
 
 import torch
 
-from deepmd.dpmodel.model.dp_model import (
-    DPModel,
-)
 from deepmd.pt.model.atomic_model import (
     DPZBLLinearEnergyAtomicModel,
 )
 from deepmd.pt.model.model.model import (
     BaseModel,
 )
 
+from .dp_model import (
+    DPModel,
+)
 from .make_model import (
     make_model,
 )
@@ -24,7 +24,7 @@
 
 
 @BaseModel.register("zbl")
-class DPZBLModel(DPZBLModel_):
+class DPZBLModel(DPModel, DPZBLModel_):
     model_type = "ener"
 
     def __init__(
@@ -103,20 +103,3 @@
             assert model_ret["dforce"] is not None
             model_predict["dforce"] = model_ret["dforce"]
         return model_predict
-
-    @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
-        """Update the selection and perform neighbor statistics.
-
-        Parameters
-        ----------
-        global_jdata : dict
-            The global data, containing the training section
-        local_jdata : dict
-            The local data refer to the current class
-        """
-        local_jdata_cpy = local_jdata.copy()
-        local_jdata_cpy["dpmodel"] = DPModel.update_sel(
-            global_jdata, local_jdata["dpmodel"]
-        )
-        return local_jdata_cpy

deepmd/pt/model/model/ener_model.py

@@ -6,12 +6,23 @@
 
 import torch
 
+from deepmd.pt.model.atomic_model import (
+    DPAtomicModel,
+)
+from deepmd.pt.model.model.model import (
+    BaseModel,
+)
+
 from .dp_model import (
     DPModel,
 )
+from .make_model import (
+    make_model,
+)
 
 
-class EnergyModel(DPModel):
+@BaseModel.register("standard")
+class EnergyModel(DPModel, make_model(DPAtomicModel)):
     model_type = "ener"
 
     def __init__(
@@ -59,6 +70,14 @@
             model_predict["updated_coord"] += coord
         return model_predict
 
+    def get_fitting_net(self):
+        """Get the fitting network."""
+        return self.atomic_model.fitting_net
+
+    def get_descriptor(self):
+        """Get the descriptor."""
+        return self.atomic_model.descriptor
+
     @torch.jit.export
     def forward_lower(
         self,