enable consistent tests

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>

Author: njzjz

deepmd/tf/fit/dipole.py

@@ -423,7 +423,6 @@ def serialize(self, suffix: str) -> dict:
             "dim_descrpt": self.dim_descrpt,
             "embedding_width": self.dim_rot_mat_1,
             "mixed_types": self.mixed_types,
-            "dim_out": 3,
             "neuron": self.n_neuron,
             "resnet_dt": self.resnet_dt,
             "numb_fparam": self.numb_fparam,
@@ -458,6 +457,15 @@ def serialize(self, suffix: str) -> dict:
                 ),
             },
             "type_map": self.type_map,
+            "var_name": "dipole",
+            "rcond": None,
+            "tot_ener_zero": False,
+            "trainable": self.trainable,
+            "layer_name": None,
+            "use_aparam_as_mask": False,
+            "spin": None,
+            "r_differentiable": True,
+            "c_differentiable": True,
         }
         return data
 

deepmd/tf/fit/dos.py

@@ -750,6 +750,11 @@ def serialize(self, suffix: str = "") -> dict:
                 "case_embd": None,
             },
             "type_map": self.type_map,
+            "tot_ener_zero": False,
+            "layer_name": None,
+            "use_aparam_as_mask": False,
+            "spin": None,
+            "atom_ener": None,
         }
         return data
 

deepmd/tf/fit/polar.py

@@ -6,6 +6,9 @@
 
 import numpy as np
 
+from deepmd.env import (
+    GLOBAL_NP_FLOAT_PRECISION,
+)
 from deepmd.tf.common import (
     cast_precision,
     get_activation_func,
@@ -641,7 +644,6 @@ def serialize(self, suffix: str) -> dict:
             "dim_descrpt": self.dim_descrpt,
             "embedding_width": self.dim_rot_mat_1,
             "mixed_types": self.mixed_types,
-            "dim_out": 3,
             "neuron": self.n_neuron,
             "resnet_dt": self.resnet_dt,
             "numb_fparam": self.numb_fparam,
@@ -652,7 +654,6 @@ def serialize(self, suffix: str) -> dict:
             "precision": self.fitting_precision.name,
             "exclude_types": [],
             "fit_diag": self.fit_diag,
-            "scale": list(self.scale),
             "shift_diag": self.shift_diag,
             "nets": self.serialize_network(
                 ntypes=self.ntypes,
@@ -674,8 +675,18 @@ def serialize(self, suffix: str) -> dict:
                 "case_embd": None,
                 "scale": self.scale.reshape(-1, 1),
                 "constant_matrix": self.constant_matrix.reshape(-1),
+                "bias_atom_e": np.zeros(
+                    (self.ntypes, self.dim_rot_mat_1), dtype=GLOBAL_NP_FLOAT_PRECISION
+                ),
             },
             "type_map": self.type_map,
+            "var_name": "polar",
+            "rcond": None,
+            "tot_ener_zero": False,
+            "trainable": self.trainable,
+            "layer_name": None,
+            "use_aparam_as_mask": False,
+            "spin": None,
         }
         return data
 

deepmd/tf/model/tensor.py

@@ -4,7 +4,13 @@
     Union,
 )
 
+import numpy as np
+
+from deepmd.env import (
+    GLOBAL_NP_FLOAT_PRECISION,
+)
 from deepmd.tf.env import (
+    GLOBAL_TF_FLOAT_PRECISION,
     MODEL_VERSION,
     global_cvt_2_ener_float,
     tf,
@@ -169,10 +175,35 @@ def build(
         )
 
         # Apply out_bias and out_std directly to tensor output
-        atype_selected = self._get_selected_atype(atype, natoms)
-        output = self._apply_out_bias_std(
-            output, atype, natoms, coord, selected_atype=atype_selected
-        )
+        # dipole not applying bias but polar does, per dpmodel
+        if self.model_type in {"polar"} and self.fitting.shift_diag:
+            v_constant_matrix = np.zeros(
+                self.ntypes,
+                dtype=GLOBAL_NP_FLOAT_PRECISION,
+            )
+            for itype in range(len(self.get_sel_type())):
+                v_constant_matrix[self.get_sel_type()[itype]] = np.mean(
+                    np.diagonal(self.out_bias[0, itype].reshape((3, 3)))
+                )
+            nframes = input_dict["nframes"]
+            nloc_mask = tf.reshape(
+                tf.tile(tf.repeat(self.fitting.sel_mask, natoms[2:]), [nframes]),
+                [nframes, -1],
+            )
+            constant_matrix = tf.reshape(
+                tf.reshape(
+                    tf.tile(tf.repeat(v_constant_matrix, natoms[2:]), [nframes]),
+                    [nframes, -1],
+                )[nloc_mask],
+                [nframes, -1],
+            )
+
+            # nf x nloc x odims, out_bias: ntypes x odims
+            output = output + tf.reshape(
+                tf.expand_dims(tf.expand_dims(constant_matrix, -1), -1)
+                * tf.eye(3, batch_shape=[1, 1], dtype=GLOBAL_TF_FLOAT_PRECISION),
+                tf.shape(output),
+            )
         framesize = nout if "global" in self.model_type else natomsel * nout
         output = tf.reshape(
             output, [-1, framesize], name="o_" + self.model_type + suffix

source/tests/consistent/common.py

@@ -354,7 +354,7 @@ def test_tf_consistent_with_ref(self) -> None:
         data1.pop("@version")
         data2.pop("@version")
 
-        if tf_obj.__class__.__name__.startswith("Polar"):
+        if tf_obj.__class__.__name__.startswith("PolarFitting"):
             data1["@variables"].pop("bias_atom_e")
 
         np.testing.assert_equal(data1, data2)

source/tests/consistent/model/common.py

@@ -77,8 +77,8 @@ def build_tf_model(
             ]
         elif ret_key == "polar":
             ret_list = [
-                ret["polar"],
                 ret["global_polar"],
+                ret["polar"],
             ]
         else:
             raise NotImplementedError

source/tests/consistent/model/test_dipole.py

@@ -73,6 +73,7 @@ def data(self) -> dict:
     pt_class = DipoleModelPT
     jax_class = DipoleModelJAX
     args = model_args()
+    atol = 1e-8
 
     def get_reference_backend(self):
         """Get the reference backend.
@@ -89,7 +90,7 @@ def get_reference_backend(self):
 
     @property
     def skip_tf(self):
-        return True  # need to fix tf consistency
+        return not INSTALLED_TF
 
     @property
     def skip_jax(self) -> bool:

source/tests/consistent/model/test_dos.py

@@ -90,7 +90,7 @@ def get_reference_backend(self):
 
     @property
     def skip_tf(self):
-        return True  # need to fix tf consistency
+        return not INSTALLED_TF
 
     @property
     def skip_jax(self) -> bool:

source/tests/consistent/model/test_polar.py

@@ -89,7 +89,7 @@ def get_reference_backend(self):
 
     @property
     def skip_tf(self):
-        return True  # need to fix tf consistency
+        return not INSTALLED_TF
 
     @property
     def skip_jax(self) -> bool: