Chore: refactor dpmodel model

deepmd/dpmodel/atomic_model/__init__.py

@@ -17,9 +17,18 @@
 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,
@@ -30,12 +39,19 @@
 from .pairtab_atomic_model import (
     PairTabAtomicModel,
 )
+from .polar_atomic_model import (
+    DPPolarAtomicModel,
+)
 
 __all__ = [
     "make_base_atomic_model",
     "BaseAtomicModel",
     "DPAtomicModel",
+    "DPEnergyAtomicModel",
     "PairTabAtomicModel",
     "LinearEnergyAtomicModel",
     "DPZBLLinearEnergyAtomicModel",
+    "DPDOSAtomicModel",
+    "DPPolarAtomicModel",
+    "DPDipoleAtomicModel",
 ]

deepmd/dpmodel/atomic_model/dipole_atomic_model.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import numpy as np
+
+from deepmd.dpmodel.fitting.dipole_fitting import (
+    DipoleFitting,
+)
+
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+
+class DPDipoleAtomicModel(DPAtomicModel):
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        if not isinstance(fitting, DipoleFitting):
+            raise TypeError(
+                "fitting must be an instance of DipoleFitting for DPDipoleAtomicModel"
+            )
+        super().__init__(descriptor, fitting, type_map, **kwargs)
+
+    def apply_out_stat(
+        self,
+        ret: dict[str, np.ndarray],
+        atype: np.ndarray,
+    ):
+        # dipole not applying bias
+        return ret

deepmd/dpmodel/atomic_model/dos_atomic_model.py

@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.dpmodel.fitting.dos_fitting import (
+    DOSFittingNet,
+)
+
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+
+class DPDOSAtomicModel(DPAtomicModel):
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        if not isinstance(fitting, DOSFittingNet):
+            raise TypeError(
+                "fitting must be an instance of DOSFittingNet for DPDOSAtomicModel"
+            )
+        super().__init__(descriptor, fitting, type_map, **kwargs)

deepmd/dpmodel/atomic_model/energy_atomic_model.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.dpmodel.fitting.ener_fitting import (
+    EnergyFittingNet,
+    InvarFitting,
+)
+
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+
+class DPEnergyAtomicModel(DPAtomicModel):
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        if not (
+            isinstance(fitting, EnergyFittingNet) or isinstance(fitting, InvarFitting)
+        ):
+            raise TypeError(
+                "fitting must be an instance of EnergyFittingNet or InvarFitting for DPEnergyAtomicModel"
+            )
+        super().__init__(descriptor, fitting, type_map, **kwargs)

deepmd/dpmodel/atomic_model/polar_atomic_model.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+
+import numpy as np
+
+from deepmd.dpmodel.fitting.polarizability_fitting import (
+    PolarFitting,
+)
+
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+
+class DPPolarAtomicModel(DPAtomicModel):
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        if not isinstance(fitting, PolarFitting):
+            raise TypeError(
+                "fitting must be an instance of PolarFitting for DPPolarAtomicModel"
+            )
+        super().__init__(descriptor, fitting, type_map, **kwargs)
+
+    def apply_out_stat(
+        self,
+        ret: dict[str, np.ndarray],
+        atype: np.ndarray,
+    ):
+        """Apply the stat to each atomic output.
+
+        Parameters
+        ----------
+        ret
+            The returned dict by the forward_atomic method
+        atype
+            The atom types. nf x nloc
+
+        """
+        out_bias, out_std = self._fetch_out_stat(self.bias_keys)
+
+        if self.fitting_net.shift_diag:
+            nframes, nloc = atype.shape
+            dtype = out_bias[self.bias_keys[0]].dtype
+            for kk in self.bias_keys:
+                ntypes = out_bias[kk].shape[0]
+                temp = np.zeros(ntypes, dtype=dtype)
+                temp = np.mean(
+                    np.diagonal(out_bias[kk].reshape(ntypes, 3, 3), axis1=1, axis2=2),
+                    axis=1,
+                )
+                modified_bias = temp[atype]
+
+                # (nframes, nloc, 1)
+                modified_bias = (
+                    modified_bias[..., np.newaxis] * (self.fitting_net.scale[atype])
+                )
+
+                eye = np.eye(3, dtype=dtype)
+                eye = np.tile(eye, (nframes, nloc, 1, 1))
+                # (nframes, nloc, 3, 3)
+                modified_bias = modified_bias[..., np.newaxis] * eye
+
+                # nf x nloc x odims, out_bias: ntypes x odims
+                ret[kk] = ret[kk] + modified_bias
+        return ret

deepmd/dpmodel/atomic_model/property_atomic_model.py

@@ -9,6 +9,9 @@
 
 
 class DPPropertyAtomicModel(DPAtomicModel):
-    def __init__(self, descriptor, fitting, type_map, **kwargs) -> None:
-        assert isinstance(fitting, PropertyFittingNet)
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        if not isinstance(fitting, PropertyFittingNet):
+            raise TypeError(
+                "fitting must be an instance of PropertyFittingNet for DPPropertyAtomicModel"
+            )
         super().__init__(descriptor, fitting, type_map, **kwargs)

deepmd/dpmodel/model/dipole_model.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+
+
+from deepmd.dpmodel.atomic_model import (
+    DPDipoleAtomicModel,
+)
+from deepmd.dpmodel.model.base_model import (
+    BaseModel,
+)
+
+from .dp_model import (
+    DPModelCommon,
+)
+from .make_model import (
+    make_model,
+)
+
+DPDipoleModel_ = make_model(DPDipoleAtomicModel)
+
+
+@BaseModel.register("dipole")
+class DipoleModel(DPModelCommon, DPDipoleModel_):
+    model_type = "dipole"
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        DPModelCommon.__init__(self)
+        DPDipoleModel_.__init__(self, *args, **kwargs)

deepmd/dpmodel/model/dos_model.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+
+from deepmd.dpmodel.atomic_model import (
+    DPDOSAtomicModel,
+)
+from deepmd.dpmodel.model.base_model import (
+    BaseModel,
+)
+
+from .dp_model import (
+    DPModelCommon,
+)
+from .make_model import (
+    make_model,
+)
+
+DPDOSModel_ = make_model(DPDOSAtomicModel)
+
+
+@BaseModel.register("dos")
+class DOSModel(DPModelCommon, DPDOSModel_):
+    model_type = "dos"
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        DPModelCommon.__init__(self)
+        DPDOSModel_.__init__(self, *args, **kwargs)

deepmd/dpmodel/model/ener_model.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from deepmd.dpmodel.atomic_model.dp_atomic_model import (
-    DPAtomicModel,
+from deepmd.dpmodel.atomic_model import (
+    DPEnergyAtomicModel,
 )
 from deepmd.dpmodel.model.base_model import (
     BaseModel,
@@ -13,7 +13,7 @@
     make_model,
 )
 
-DPEnergyModel_ = make_model(DPAtomicModel)
+DPEnergyModel_ = make_model(DPEnergyAtomicModel)
 
 
 @BaseModel.register("ener")

deepmd/dpmodel/model/model.py

@@ -1,4 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import copy
+
 from deepmd.dpmodel.atomic_model.dp_atomic_model import (
     DPAtomicModel,
 )
@@ -8,18 +10,33 @@
 from deepmd.dpmodel.descriptor.base_descriptor import (
     BaseDescriptor,
 )
+from deepmd.dpmodel.fitting.base_fitting import (
+    BaseFitting,
+)
 from deepmd.dpmodel.fitting.ener_fitting import (
     EnergyFittingNet,
 )
 from deepmd.dpmodel.model.base_model import (
     BaseModel,
 )
+from deepmd.dpmodel.model.dipole_model import (
+    DipoleModel,
+)
+from deepmd.dpmodel.model.dos_model import (
+    DOSModel,
+)
 from deepmd.dpmodel.model.dp_zbl_model import (
     DPZBLModel,
 )
 from deepmd.dpmodel.model.ener_model import (
     EnergyModel,
 )
+from deepmd.dpmodel.model.polar_model import (
+    PolarModel,
+)
+from deepmd.dpmodel.model.property_model import (
+    PropertyModel,
+)
 from deepmd.dpmodel.model.spin_model import (
     SpinModel,
 )
@@ -28,6 +45,29 @@
 )
 
 
+def _get_standard_model_components(data, ntypes):
+    # descriptor
+    data["descriptor"]["ntypes"] = ntypes
+    data["descriptor"]["type_map"] = copy.deepcopy(data["type_map"])
+    descriptor = BaseDescriptor(**data["descriptor"])
+    # fitting
+    fitting_net = data.get("fitting_net", {})
+    fitting_net["type"] = fitting_net.get("type", "ener")
+    fitting_net["ntypes"] = descriptor.get_ntypes()
+    fitting_net["type_map"] = copy.deepcopy(data["type_map"])
+    fitting_net["mixed_types"] = descriptor.mixed_types()
+    if fitting_net["type"] in ["dipole", "polar"]:
+        fitting_net["embedding_width"] = descriptor.get_dim_emb()
+    fitting_net["dim_descrpt"] = descriptor.get_dim_out()
+    grad_force = "direct" not in fitting_net["type"]
+    if not grad_force:
+        fitting_net["out_dim"] = descriptor.get_dim_emb()
+        if "ener" in fitting_net["type"]:
+            fitting_net["return_energy"] = True
+    fitting = BaseFitting(**fitting_net)
+    return descriptor, fitting, fitting_net["type"]
+
+
 def get_standard_model(data: dict) -> EnergyModel:
     """Get a EnergyModel from a dictionary.
 
@@ -40,29 +80,33 @@
         raise ValueError(
             "In the DP backend, type_embedding is not at the model level, but within the descriptor. See type embedding documentation for details."
         )
-    data["descriptor"]["type_map"] = data["type_map"]
-    data["descriptor"]["ntypes"] = len(data["type_map"])
-    fitting_type = data["fitting_net"].pop("type")
-    data["fitting_net"]["type_map"] = data["type_map"]
-    descriptor = BaseDescriptor(
-        **data["descriptor"],
-    )
-    if fitting_type == "ener":
-        fitting = EnergyFittingNet(
-            ntypes=descriptor.get_ntypes(),
-            dim_descrpt=descriptor.get_dim_out(),
-            mixed_types=descriptor.mixed_types(),
-            **data["fitting_net"],
-        )
+    data = copy.deepcopy(data)
+    ntypes = len(data["type_map"])
+    descriptor, fitting, fitting_net_type = _get_standard_model_components(data, ntypes)
+    atom_exclude_types = data.get("atom_exclude_types", [])
+    pair_exclude_types = data.get("pair_exclude_types", [])
+
+    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
+    elif fitting_net_type == "property":
+        modelcls = PropertyModel
     else:
-        raise ValueError(f"Unknown fitting type {fitting_type}")
-    return EnergyModel(
+        raise RuntimeError(f"Unknown fitting type: {fitting_net_type}")
+
+    model = modelcls(
         descriptor=descriptor,
         fitting=fitting,
         type_map=data["type_map"],
-        atom_exclude_types=data.get("atom_exclude_types", []),
-        pair_exclude_types=data.get("pair_exclude_types", []),
+        atom_exclude_types=atom_exclude_types,
+        pair_exclude_types=pair_exclude_types,
     )
+    return model
 
 
 def get_zbl_model(data: dict) -> DPZBLModel: