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,24 @@
+# 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):
+        assert isinstance(fitting, DipoleFitting)
+        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,14 @@
+# 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):
+        assert isinstance(fitting, DOSFittingNet)
+        super().__init__(descriptor, fitting, type_map, **kwargs)

deepmd/dpmodel/atomic_model/energy_atomic_model.py

@@ -0,0 +1,17 @@
+# 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):
+        assert isinstance(fitting, EnergyFittingNet) or isinstance(
+            fitting, InvarFitting
+        )
+        super().__init__(descriptor, fitting, type_map, **kwargs)

deepmd/dpmodel/atomic_model/polar_atomic_model.py

@@ -0,0 +1,58 @@
+# 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):
+        assert isinstance(fitting, PolarFitting)
+        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/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/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/model/dipole_model.py

@@ -22,11 +22,11 @@
     make_model,
 )
 
-DPDOSModel_ = make_model(DPDipoleAtomicModel)
+DPDipoleModel_ = make_model(DPDipoleAtomicModel)
 
 
 @BaseModel.register("dipole")
-class DipoleModel(DPModelCommon, DPDOSModel_):
+class DipoleModel(DPModelCommon, DPDipoleModel_):
     model_type = "dipole"
 
     def __init__(
@@ -35,7 +35,7 @@ def __init__(
         **kwargs,
     ):
         DPModelCommon.__init__(self)
-        DPDOSModel_.__init__(self, *args, **kwargs)
+        DPDipoleModel_.__init__(self, *args, **kwargs)
 
     def translated_output_def(self):
         out_def_data = self.model_output_def().get_data()

deepmd/pt/model/model/polar_model.py

@@ -22,11 +22,11 @@
     make_model,
 )
 
-DPDOSModel_ = make_model(DPPolarAtomicModel)
+DPPolarModel_ = make_model(DPPolarAtomicModel)
 
 
 @BaseModel.register("polar")
-class PolarModel(DPModelCommon, DPDOSModel_):
+class PolarModel(DPModelCommon, DPPolarModel_):
     model_type = "polar"
 
     def __init__(
@@ -35,7 +35,7 @@ def __init__(
         **kwargs,
     ):
         DPModelCommon.__init__(self)
-        DPDOSModel_.__init__(self, *args, **kwargs)
+        DPPolarModel_.__init__(self, *args, **kwargs)
 
     def translated_output_def(self):
         out_def_data = self.model_output_def().get_data()