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,61 @@
+# 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)
+                for i in range(ntypes):
+                    temp[i] = np.mean(np.diagonal(out_bias[kk][i].reshape(3, 3)))
+                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.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.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

@@ -8,6 +8,9 @@
 from deepmd.dpmodel.descriptor.base_descriptor import (
     BaseDescriptor,
 )
+from deepmd.dpmodel.fitting.dos_fitting import (
+    DOSFittingNet,
+)
 from deepmd.dpmodel.fitting.ener_fitting import (
     EnergyFittingNet,
 )
@@ -54,8 +57,15 @@
             mixed_types=descriptor.mixed_types(),
             **data["fitting_net"],
         )
+    elif fitting_type == "dos":
+        fitting = DOSFittingNet(
+            ntypes=descriptor.get_ntypes(),
+            dim_descrpt=descriptor.get_dim_out(),
+            mixed_types=descriptor.mixed_types(),
+            **data["fitting_net"],
+        )
     else:
-        raise ValueError(f"Unknown fitting type {fitting_type}")
+        raise ValueError(f"Unknown fitting type {fitting_type}")  # fix
     return EnergyModel(
         descriptor=descriptor,
         fitting=fitting,

deepmd/dpmodel/model/polar_model.py

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

deepmd/pt/model/atomic_model/dipole_atomic_model.py

@@ -12,8 +12,11 @@
 
 
 class DPDipoleAtomicModel(DPAtomicModel):
-    def __init__(self, descriptor, fitting, type_map, **kwargs) -> None:
-        assert isinstance(fitting, DipoleFittingNet)
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        if not isinstance(fitting, DipoleFittingNet):
+            raise TypeError(
+                "fitting must be an instance of DipoleFittingNet for DPDipoleAtomicModel"
+            )
         super().__init__(descriptor, fitting, type_map, **kwargs)
 
     def apply_out_stat(

deepmd/pt/model/atomic_model/dos_atomic_model.py

@@ -9,6 +9,9 @@
 
 
 class DPDOSAtomicModel(DPAtomicModel):
-    def __init__(self, descriptor, fitting, type_map, **kwargs) -> None:
-        assert isinstance(fitting, DOSFittingNet)
+    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/pt/model/atomic_model/energy_atomic_model.py

@@ -11,10 +11,13 @@
 
 
 class DPEnergyAtomicModel(DPAtomicModel):
-    def __init__(self, descriptor, fitting, type_map, **kwargs) -> None:
-        assert (
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        if not (
             isinstance(fitting, EnergyFittingNet)
             or isinstance(fitting, EnergyFittingNetDirect)
             or isinstance(fitting, InvarFitting)
-        )
+        ):
+            raise TypeError(
+                "fitting must be an instance of EnergyFittingNet, EnergyFittingNetDirect or InvarFitting for DPEnergyAtomicModel"
+            )
         super().__init__(descriptor, fitting, type_map, **kwargs)

deepmd/pt/model/atomic_model/polar_atomic_model.py

@@ -12,8 +12,11 @@
 
 
 class DPPolarAtomicModel(DPAtomicModel):
-    def __init__(self, descriptor, fitting, type_map, **kwargs) -> None:
-        assert isinstance(fitting, PolarFittingNet)
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        if not isinstance(fitting, PolarFittingNet):
+            raise TypeError(
+                "fitting must be an instance of PolarFittingNet for DPPolarAtomicModel"
+            )
         super().__init__(descriptor, fitting, type_map, **kwargs)
 
     def apply_out_stat(

deepmd/pt/model/atomic_model/property_atomic_model.py

@@ -12,8 +12,11 @@
 
 
 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)
 
     def apply_out_stat(