Feat: Add polar stat constant matrix calculation to PT

deepmd/dpmodel/fitting/dipole_fitting.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import copy
 from typing import (
     Any,
     Dict,
@@ -19,6 +20,9 @@
     OutputVariableDef,
     fitting_check_output,
 )
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
 
 from .general_fitting import (
     GeneralFitting,
@@ -153,6 +157,12 @@
         data["c_differentiable"] = self.c_differentiable
         return data
 
+    @classmethod
+    def deserialize(cls, data: dict) -> "GeneralFitting":
+        data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
+        return super().deserialize(data)
+
     def output_def(self):
         return FittingOutputDef(
             [

deepmd/dpmodel/fitting/ener_fitting.py

@@ -18,6 +18,9 @@
     from deepmd.dpmodel.fitting.general_fitting import (
         GeneralFitting,
     )
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
 
 
 @InvarFitting.register("ener")
@@ -69,6 +72,7 @@
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
         data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("var_name")
         data.pop("dim_out")
         return super().deserialize(data)

deepmd/dpmodel/fitting/general_fitting.py

@@ -21,9 +21,6 @@
     FittingNet,
     NetworkCollection,
 )
-from deepmd.utils.version import (
-    check_version_compatibility,
-)
 
 from .base_fitting import (
     BaseFitting,
@@ -256,7 +253,6 @@ def serialize(self) -> dict:
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
         data = copy.deepcopy(data)
-        check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class")
         data.pop("type")
         variables = data.pop("@variables")

deepmd/dpmodel/fitting/invar_fitting.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import copy
 from typing import (
     Any,
     Dict,
@@ -16,6 +17,9 @@
     OutputVariableDef,
     fitting_check_output,
 )
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
 
 from .general_fitting import (
     GeneralFitting,
@@ -169,6 +173,12 @@
         data["atom_ener"] = self.atom_ener
         return data
 
+    @classmethod
+    def deserialize(cls, data: dict) -> "GeneralFitting":
+        data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
+        return super().deserialize(data)
+
     def _net_out_dim(self):
         """Set the FittingNet output dim."""
         return self.dim_out

deepmd/dpmodel/fitting/polarizability_fitting.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import copy
 from typing import (
     Any,
     Dict,
@@ -22,6 +23,9 @@
     OutputVariableDef,
     fitting_check_output,
 )
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
 
 from .general_fitting import (
     GeneralFitting,
@@ -139,6 +143,7 @@
             ntypes, 1
         )
         self.shift_diag = shift_diag
+        self.constant_matrix = np.zeros(ntypes, dtype=GLOBAL_NP_FLOAT_PRECISION)
         super().__init__(
             var_name=var_name,
             ntypes=ntypes,
@@ -168,15 +173,36 @@
             else self.embedding_width * self.embedding_width
         )
 
+    def __setitem__(self, key, value):
+        if key in ["constant_matrix"]:
+            self.constant_matrix = value
+        else:
+            super().__setitem__(key, value)
+
+    def __getitem__(self, key):
+        if key in ["constant_matrix"]:
+            return self.constant_matrix
+        else:
+            super().__getitem__(key)
+
     def serialize(self) -> dict:
         data = super().serialize()
         data["type"] = "polar"
+        data["@version"] = 2
         data["embedding_width"] = self.embedding_width
         data["old_impl"] = self.old_impl
         data["fit_diag"] = self.fit_diag
+        data["shift_diag"] = self.shift_diag
         data["@variables"]["scale"] = self.scale
+        data["@variables"]["constant_matrix"] = self.constant_matrix
         return data
 
+    @classmethod
+    def deserialize(cls, data: dict) -> "GeneralFitting":
+        data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 2, 1)
+        return super().deserialize(data)
+
     def output_def(self):
         return FittingOutputDef(
             [
@@ -246,4 +272,13 @@
             "bim,bmj->bij", np.transpose(gr, axes=(0, 2, 1)), out
         )  # (nframes * nloc, 3, 3)
         out = out.reshape(nframes, nloc, 3, 3)
+        if self.shift_diag:
+            bias = self.constant_matrix[atype]
+            # (nframes, nloc, 1)
+            bias = np.expand_dims(bias, axis=-1) * self.scale[atype]
+            eye = np.eye(3)
+            eye = np.tile(eye, (nframes, nloc, 1, 1))
+            # (nframes, nloc, 3, 3)
+            bias = np.expand_dims(bias, axis=-1) * eye
+            out = out + bias
         return {self.var_name: out}

deepmd/pt/model/task/dipole.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import copy
 import logging
 from typing import (
     Callable,
@@ -25,6 +26,9 @@
 from deepmd.utils.path import (
     DPPath,
 )
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
 
 log = logging.getLogger(__name__)
 
@@ -123,6 +127,12 @@
         data["c_differentiable"] = self.c_differentiable
         return data
 
+    @classmethod
+    def deserialize(cls, data: dict) -> "GeneralFitting":
+        data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
+        return super().deserialize(data)
+
     def output_def(self) -> FittingOutputDef:
         return FittingOutputDef(
             [

deepmd/pt/model/task/ener.py

@@ -36,6 +36,9 @@
 from deepmd.utils.path import (
     DPPath,
 )
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
 
 dtype = env.GLOBAL_PT_FLOAT_PRECISION
 device = env.DEVICE
@@ -140,6 +143,12 @@
         data["atom_ener"] = self.atom_ener
         return data
 
+    @classmethod
+    def deserialize(cls, data: dict) -> "GeneralFitting":
+        data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
+        return super().deserialize(data)
+
     def compute_output_stats(
         self,
         merged: Union[Callable[[], List[dict]], List[dict]],
@@ -241,6 +250,7 @@
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
         data = copy.deepcopy(data)
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("var_name")
         data.pop("dim_out")
         return super().deserialize(data)

deepmd/pt/model/task/fitting.py

@@ -49,9 +49,6 @@
 from deepmd.utils.finetune import (
     change_energy_bias_lower,
 )
-from deepmd.utils.version import (
-    check_version_compatibility,
-)
 
 dtype = env.GLOBAL_PT_FLOAT_PRECISION
 device = env.DEVICE
@@ -371,7 +368,6 @@ def serialize(self) -> dict:
     @classmethod
     def deserialize(cls, data: dict) -> "GeneralFitting":
         data = copy.deepcopy(data)
-        check_version_compatibility(data.pop("@version", 1), 1, 1)
         variables = data.pop("@variables")
         nets = data.pop("nets")
         obj = cls(**data)