Chore: refactor dpmodel

deepmd/dpmodel/__init__.py

@@ -5,7 +5,7 @@
     NativeOP,
 )
 from .model import (
-    DPModel,
+    DPModelCommon,
 )
 from .output_def import (
     FittingOutputDef,
@@ -19,7 +19,7 @@
 )
 
 __all__ = [
-    "DPModel",
+    "DPModelCommon",
     "PRECISION_DICT",
     "DEFAULT_PRECISION",
     "NativeOP",

deepmd/dpmodel/model/__init__.py

@@ -13,7 +13,7 @@
 """
 
 from .dp_model import (
-    DPModel,
+    DPModelCommon,
 )
 from .make_model import (
     make_model,
@@ -23,7 +23,7 @@
 )
 
 __all__ = [
-    "DPModel",
+    "DPModelCommon",
     "SpinModel",
     "make_model",
 ]

deepmd/dpmodel/model/base_model.py

@@ -35,7 +35,7 @@
 
         def __new__(cls, *args, **kwargs):
             if inspect.isabstract(cls):
-                cls = cls.get_class_by_type(kwargs.get("type", "standard"))
+                cls = cls.get_class_by_type(kwargs.get("type", "ener"))
             return super().__new__(cls)
 
         @abstractmethod
@@ -151,7 +151,7 @@
             local_jdata : dict
                 The local data refer to the current class
             """
-            cls = cls.get_class_by_type(local_jdata.get("type", "standard"))
+            cls = cls.get_class_by_type(local_jdata.get("type", "ener"))
             return cls.update_sel(global_jdata, local_jdata)
 
     return BaseBaseModel

deepmd/dpmodel/model/dp_model.py

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

deepmd/dpmodel/model/ener_model.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.dpmodel.atomic_model.dp_atomic_model import (
+    DPAtomicModel,
+)
+from deepmd.dpmodel.model.base_model import (
+    BaseModel,
+)
+
+from .dp_model import (
+    DPModelCommon,
+)
+from .make_model import (
+    make_model,
+)
+
+DPEnergyModel_ = make_model(DPAtomicModel)
+
+
+@BaseModel.register("ener")
+class EnergyModel(DPModelCommon, DPEnergyModel_):
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        DPModelCommon.__init__(self)
+        DPEnergyModel_.__init__(self, *args, **kwargs)

deepmd/dpmodel/model/model.py

@@ -5,8 +5,8 @@
 from deepmd.dpmodel.fitting.ener_fitting import (
     EnergyFittingNet,
 )
-from deepmd.dpmodel.model.dp_model import (
-    DPModel,
+from deepmd.dpmodel.model.ener_model import (
+    EnergyModel,
 )
 from deepmd.dpmodel.model.spin_model import (
     SpinModel,
@@ -16,8 +16,8 @@
 )
 
 
-def get_standard_model(data: dict) -> DPModel:
-    """Get a standard DPModel from a dictionary.
+def get_standard_model(data: dict) -> EnergyModel:
+    """Get a EnergyModel from a dictionary.
 
     Parameters
     ----------
@@ -41,7 +41,7 @@
         )
     else:
         raise ValueError(f"Unknown fitting type {fitting_type}")
-    return DPModel(
+    return EnergyModel(
         descriptor=descriptor,
         fitting=fitting,
         type_map=data["type_map"],

deepmd/dpmodel/model/spin_model.py

@@ -7,8 +7,11 @@
 
 import numpy as np
 
-from deepmd.dpmodel.model.dp_model import (
-    DPModel,
+from deepmd.dpmodel.atomic_model.dp_atomic_model import (
+    DPAtomicModel,
+)
+from deepmd.dpmodel.model.make_model import (
+    make_model,
 )
 from deepmd.utils.spin import (
     Spin,
@@ -259,7 +262,9 @@
 
     @classmethod
     def deserialize(cls, data) -> "SpinModel":
-        backbone_model_obj = DPModel.deserialize(data["backbone_model"])
+        backbone_model_obj = make_model(DPAtomicModel).deserialize(
+            data["backbone_model"]
+        )
         spin = Spin.deserialize(data["spin"])
         return cls(
             backbone_model=backbone_model_obj,

deepmd/pt/model/atomic_model/__init__.py

@@ -17,21 +17,37 @@
 from .base_atomic_model import (
     BaseAtomicModel,
 )
+from .dipole_atomic_model import (
+    DPDipoleAtomicModel,
+)
+from .dos_atomic_model import (
+    DPDOSAtomicModel,
+)
 from .dp_atomic_model import (
     DPAtomicModel,
 )
+from .energy_atomic_model import (
+    DPEnergyAtomicModel,
+)
 from .linear_atomic_model import (
     DPZBLLinearEnergyAtomicModel,
     LinearEnergyAtomicModel,
 )
 from .pairtab_atomic_model import (
     PairTabAtomicModel,
 )
+from .polar_atomic_model import (
+    DPPolarAtomicModel,
+)
 
 __all__ = [
     "BaseAtomicModel",
     "DPAtomicModel",
+    "DPDOSAtomicModel",
+    "DPEnergyAtomicModel",
     "PairTabAtomicModel",
     "LinearEnergyAtomicModel",
+    "DPPolarAtomicModel",
+    "DPDipoleAtomicModel",
     "DPZBLLinearEnergyAtomicModel",
 ]

deepmd/pt/model/atomic_model/dipole_atomic_model.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Dict,
+)
+
+import torch
+
+from deepmd.pt.model.task.dipole import (
+    DipoleFittingNet,
+)
+
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+
+class DPDipoleAtomicModel(DPAtomicModel):
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        assert isinstance(fitting, DipoleFittingNet)
+        super().__init__(descriptor, fitting, type_map, **kwargs)
+
+    def apply_out_stat(
+        self,
+        ret: Dict[str, torch.Tensor],
+        atype: torch.Tensor,
+    ):
+        # dipole not applying bias
+        return ret

deepmd/pt/model/atomic_model/dos_atomic_model.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.pt.model.task.dos import (
+    DOSFittingNet,
+)
+
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+
+class DPDOSAtomicModel(DPAtomicModel):
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        assert isinstance(fitting, DOSFittingNet)
+        super().__init__(descriptor, fitting, type_map, **kwargs)

deepmd/pt/model/atomic_model/energy_atomic_model.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.pt.model.task.ener import (
+    EnergyFittingNet,
+    EnergyFittingNetDirect,
+    InvarFitting,
+)
+
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+
+class DPEnergyAtomicModel(DPAtomicModel):
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        assert (
+            isinstance(fitting, EnergyFittingNet)
+            or isinstance(fitting, EnergyFittingNetDirect)
+            or isinstance(fitting, InvarFitting)
+        )
+        super().__init__(descriptor, fitting, type_map, **kwargs)

deepmd/pt/model/atomic_model/polar_atomic_model.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Dict,
+)
+
+import torch
+
+from deepmd.pt.model.task.polarizability import (
+    PolarFittingNet,
+)
+
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+
+class DPPolarAtomicModel(DPAtomicModel):
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        assert isinstance(fitting, PolarFittingNet)
+        super().__init__(descriptor, fitting, type_map, **kwargs)
+
+    def apply_out_stat(
+        self,
+        ret: Dict[str, torch.Tensor],
+        atype: torch.Tensor,
+    ):
+        # TODO: migrate bias
+        return ret

deepmd/pt/model/model/__init__.py

@@ -30,8 +30,14 @@
     Spin,
 )
 
+from .dipole_model import (
+    DipoleModel,
+)
+from .dos_model import (
+    DOSModel,
+)
 from .dp_model import (
-    DPModel,
+    DPModelCommon,
 )
 from .dp_zbl_model import (
     DPZBLModel,
@@ -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"] in ["ener", "direct_force_ener"]:
+        modelcls = EnergyModel
+    else:
+        raise RuntimeError(f"Unknown fitting type: {fitting_net['type']}")
+
+    model = modelcls(
         descriptor=descriptor,
         fitting=fitting,
         type_map=model_params["type_map"],
@@ -183,7 +203,7 @@
 __all__ = [
     "BaseModel",
     "get_model",
-    "DPModel",
+    "DPModelCommon",
     "EnergyModel",
     "FrozenModel",
     "SpinModel",

deepmd/pt/model/model/dipole_model.py

@@ -6,20 +6,34 @@
 
 import torch
 
+from deepmd.pt.model.atomic_model import (
+    DPDipoleAtomicModel,
+)
+from deepmd.pt.model.model.model import (
+    BaseModel,
+)
+
 from .dp_model import (
-    DPModel,
+    DPModelCommon,
 )
+from .make_model import (
+    make_model,
+)
+
+DPDOSModel_ = make_model(DPDipoleAtomicModel)
 
 
-class DipoleModel(DPModel):
+@BaseModel.register("dipole")
+class DipoleModel(DPModelCommon, DPDOSModel_):
     model_type = "dipole"
 
     def __init__(
         self,
         *args,
         **kwargs,
     ):
-        super().__init__(*args, **kwargs)
+        DPModelCommon.__init__(self)
+        DPDOSModel_.__init__(self, *args, **kwargs)
 
     def forward(
         self,