Chore: migrate and refactor polar and dos bias

deepmd/pt/model/atomic_model/polar_atomic_model.py

@@ -39,9 +39,10 @@ def apply_out_stat(
         if self.fitting_net.shift_diag:
             nframes, nloc = atype.shape
             device = out_bias[self.bias_keys[0]].device
+            dtype = out_bias[self.bias_keys[0]].dtype
             for kk in self.bias_keys:
                 ntypes = out_bias[kk].shape[0]
-                temp = torch.zeros(ntypes, device=device)
+                temp = torch.zeros(ntypes, dtype=dtype, device=device)
                 for i in range(ntypes):
                     temp[i] = torch.mean(torch.diagonal(out_bias[kk][i].reshape(3, 3)))
                 modified_bias = temp[atype]
@@ -51,7 +52,7 @@ def apply_out_stat(
                     modified_bias.unsqueeze(-1) * self.fitting_net.scale[atype]
                 )
 
-                eye = torch.eye(3, device=device)
+                eye = torch.eye(3, dtype=dtype, device=device)
                 eye = eye.repeat(nframes, nloc, 1, 1)
                 # (nframes, nloc, 3, 3)
                 modified_bias = modified_bias.unsqueeze(-1) * eye

deepmd/pt/utils/stat.py

@@ -567,9 +567,8 @@ def compute_output_stats_atomic(
                 nan_padding = np.empty((missing_types, bias_atom_e[kk].shape[1]))
                 nan_padding.fill(np.nan)
                 bias_atom_e[kk] = np.concatenate([bias_atom_e[kk], nan_padding], axis=0)
-                std_atom_e[kk] = np.concatenate([bias_atom_e[kk], nan_padding], axis=0)
+                std_atom_e[kk] = np.concatenate([std_atom_e[kk], nan_padding], axis=0)
         else:
             # this key does not have atomic labels, skip it.
             continue
-    bias_atom_e, std_atom_e = _post_process_stat(bias_atom_e, std_atom_e)
     return bias_atom_e, std_atom_e

source/tests/pt/model/test_atomic_model_atomic_stat.py

@@ -212,8 +212,15 @@ def cvt_ret(x):
             self.merged_output_stat, stat_file_path=self.stat_file_path
         )
         ret1 = md0.forward_common_atomic(*args)
-        expected_std = np.ones((2, 2, 2))  # 2 keys, 2 atypes, 2 max dims.
-        np.testing.assert_almost_equal(to_numpy_array(md0.out_std), expected_std)
+        expected_std = np.ones(
+            (2, 2, 2), dtype=np.float64
+        )  # 2 keys, 2 atypes, 2 max dims.
+        expected_std[0, :, :1] = np.array([0.0, 0.816496]).reshape(
+            2, 1
+        )  # updating std for foo based on [5.0, 5.0, 5.0], [5.0, 6.0, 7.0]]
+        np.testing.assert_almost_equal(
+            to_numpy_array(md0.out_std), expected_std, decimal=4
+        )
         ret1 = cvt_ret(ret1)
         # nt x odim
         foo_bias = np.array([5.0, 6.0]).reshape(2, 1)
@@ -233,7 +240,9 @@ def raise_error():
         ret2 = cvt_ret(ret2)
         for kk in ["foo", "bar"]:
             np.testing.assert_almost_equal(ret1[kk], ret2[kk])
-        np.testing.assert_almost_equal(to_numpy_array(md0.out_std), expected_std)
+        np.testing.assert_almost_equal(
+            to_numpy_array(md0.out_std), expected_std, decimal=4
+        )
 
         # 4. test change bias
         BaseAtomicModel.change_out_bias(
@@ -249,15 +258,19 @@ def raise_error():
         ]
         ret3 = md0.forward_common_atomic(*args)
         ret3 = cvt_ret(ret3)
-
+        expected_std[0, :, :1] = np.array([1.24722, 0.47140]).reshape(
+            2, 1
+        )  # updating std for foo based on [4.0, 3.0, 2.0], [1.0, 1.0, 1.0]]
         expected_ret3 = {}
         # new bias [2.666, 1.333]
         expected_ret3["foo"] = np.array(
             [[3.6667, 4.6667, 4.3333], [6.6667, 6.3333, 7.3333]]
         ).reshape(2, 3, 1)
         for kk in ["foo"]:
             np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk], decimal=4)
-        np.testing.assert_almost_equal(to_numpy_array(md0.out_std), expected_std)
+        np.testing.assert_almost_equal(
+            to_numpy_array(md0.out_std), expected_std, decimal=4
+        )
 
 
 class TestAtomicModelStatMergeGlobalAtomic(

source/tests/pt/model/test_polar_atomic_model_stat.py

@@ -198,7 +198,10 @@ def cvt_ret(x):
         )
         ret1 = md0.forward_common_atomic(*args)
         ret1 = cvt_ret(ret1)
-        expected_std = np.ones((1, 2, 9))  # 1 keys, 2 atypes, 9 max dims.
+        expected_std = np.zeros(
+            (1, 2, 9), dtype=np.float64
+        )  # 1 keys, 2 atypes, 9 max dims.
+        expected_std[:, 1, :] = np.ones(9, dtype=np.float64) * 0.8164966  # updating std
         # nt x odim (dia)
         diagnoal_bias = np.array(
             [
@@ -239,6 +242,34 @@ def raise_error():
         ret3 = cvt_ret(ret3)
 
         expected_ret3 = {}
+        expected_std = np.array(
+            [
+                [
+                    [
+                        1.4142136,
+                        1.4142136,
+                        1.4142136,
+                        1.2472191,
+                        1.2472191,
+                        1.2472191,
+                        1.2472191,
+                        1.2472191,
+                        1.2472191,
+                    ],
+                    [
+                        0.4714045,
+                        0.4714045,
+                        0.4714045,
+                        0.8164966,
+                        0.8164966,
+                        0.8164966,
+                        2.6246693,
+                        2.6246693,
+                        2.6246693,
+                    ],
+                ]
+            ]
+        )
         # new bias [[[3.0000, -, -, -, 2.6667, -, -, -, 2.3333],
         # [1.6667, -, -, -, 2.0000, -, -, -, 1.3333]]]
         # which yields [2.667, 1.667]