Chore: migrate and refactor polar and dos bias

deepmd/entrypoints/test.py

@@ -897,7 +897,7 @@
         polar = polar.reshape((polar.shape[0], -1, 9))[:, sel_mask, :].reshape(
             (polar.shape[0], -1)
         )
-        rmse_f = rmse(polar - test_data["atomic_polarizability"][:numb_test])
+        rmse_f = rmse(polar - test_data["atom_polarizability"][:numb_test])
 
     log.info(f"# number of test data : {numb_test:d} ")
     log.info(f"Polarizability  RMSE       : {rmse_f:e}")
@@ -926,7 +926,7 @@
             pe = np.concatenate(
                 (
                     np.reshape(
-                        test_data["atomic_polarizability"][:numb_test],
+                        test_data["atom_polarizability"][:numb_test],
                         [-1, 9 * sel_natoms],
                     ),
                     np.reshape(polar, [-1, 9 * sel_natoms]),
@@ -1037,7 +1037,7 @@
         dipole = dipole.reshape((dipole.shape[0], -1, 3))[:, sel_mask, :].reshape(
             (dipole.shape[0], -1)
         )
-        rmse_f = rmse(dipole - test_data["atomic_dipole"][:numb_test])
+        rmse_f = rmse(dipole - test_data["atom_dipole"][:numb_test])
 
     log.info(f"# number of test data : {numb_test:d}")
     log.info(f"Dipole  RMSE       : {rmse_f:e}")
@@ -1061,7 +1061,7 @@
             pe = np.concatenate(
                 (
                     np.reshape(
-                        test_data["atomic_dipole"][:numb_test], [-1, 3 * sel_natoms]
+                        test_data["atom_dipole"][:numb_test], [-1, 3 * sel_natoms]
                     ),
                     np.reshape(dipole, [-1, 3 * sel_natoms]),
                 ),

deepmd/pt/loss/tensor.py

@@ -93,14 +93,14 @@
         if (
             self.has_local_weight
             and self.tensor_name in model_pred
-            and "atomic_" + self.label_name in label
+            and "atom_" + self.label_name in label
         ):
-            find_local = label.get("find_" + "atomic_" + self.label_name, 0.0)
+            find_local = label.get("find_" + "atom_" + self.label_name, 0.0)
             local_weight = self.local_weight * find_local
             local_tensor_pred = model_pred[self.tensor_name].reshape(
                 [-1, natoms, self.tensor_size]
             )
-            local_tensor_label = label["atomic_" + self.label_name].reshape(
+            local_tensor_label = label["atom_" + self.label_name].reshape(
                 [-1, natoms, self.tensor_size]
             )
             diff = (local_tensor_pred - local_tensor_label).reshape(

deepmd/pt/model/task/polarizability.py

@@ -231,9 +231,9 @@
                 for sys in range(len(sampled)):
                     nframs = sampled[sys]["atype"].shape[0]
 
-                    if sampled[sys]["find_atomic_polarizability"] > 0.0:
+                    if sampled[sys]["find_atom_polarizability"] > 0.0:
                         sys_atom_polar = compute_stats_from_atomic(
-                            sampled[sys]["atomic_polarizability"].numpy(force=True),
+                            sampled[sys]["atom_polarizability"].numpy(force=True),
                             sampled[sys]["atype"].numpy(force=True),
                         )[0]
                     else:

deepmd/tf/fit/polar.py

@@ -193,7 +193,7 @@
                 atom_has_polar = [
                     w for w in all_stat["type"][ss][0] if (w in self.sel_type)
                 ]  # select atom with polar
-                if all_stat["find_atomic_polarizability"][ss] > 0.0:
+                if all_stat["find_atom_polarizability"][ss] > 0.0:
                     for itype in range(
                         len(self.sel_type)
                     ):  # Atomic polar mode, should specify the atoms
@@ -208,7 +208,7 @@
 
                         polar_bias.append(
                             np.sum(
-                                all_stat["atomic_polarizability"][ss].reshape(
+                                all_stat["atom_polarizability"][ss].reshape(
                                     nframes, len(atom_has_polar), -1
                                 )[:, index_lis, :]
                                 / nframes,

deepmd/tf/loss/tensor.py

@@ -70,11 +70,11 @@
 
     def build(self, learning_rate, natoms, model_dict, label_dict, suffix):
         polar_hat = label_dict[self.label_name]
-        atomic_polar_hat = label_dict["atomic_" + self.label_name]
+        atomic_polar_hat = label_dict["atom_" + self.label_name]
         polar = tf.reshape(model_dict[self.tensor_name], [-1])
 
         find_global = label_dict["find_" + self.label_name]
-        find_atomic = label_dict["find_atomic_" + self.label_name]
+        find_atomic = label_dict["find_atom_" + self.label_name]
 
         # YHT: added for global / local dipole combination
         l2_loss = global_cvt_2_tf_float(0.0)

deepmd/utils/data.py

@@ -577,6 +577,8 @@
         else:
             data["type"] = np.tile(self.atom_type[self.idx_map], (nframes, 1))
 
+        # standardize keys
+        data = {kk.replace("atomic", "atom"): vv for kk, vv in data.items()}
         return data
 
     def _load_data(

source/tests/pt/model/test_polar_stat.py

@@ -31,8 +31,8 @@ def setUp(self) -> None:
         self.sampled = [
             {
                 "atype": types,
-                "find_atomic_polarizability": find_atomic_polarizability,
-                "atomic_polarizability": atomic_polarizability,
+                "find_atom_polarizability": find_atomic_polarizability,
+                "atom_polarizability": atomic_polarizability,
                 "polarizability": polarizability,
                 "find_polarizability": find_polarizability,
             }
@@ -61,13 +61,13 @@ def test_atomic_consistency(self):
         np.testing.assert_allclose(tfbias, to_numpy_array(ptbias))
 
     def test_global_consistency(self):
-        self.sampled[0]["find_atomic_polarizability"] = -1
-        self.sampled[0]["polarizability"] = self.sampled[0][
-            "atomic_polarizability"
-        ].sum(dim=1)
-        self.all_stat["find_atomic_polarizability"] = [-1]
+        self.sampled[0]["find_atom_polarizability"] = -1
+        self.sampled[0]["polarizability"] = self.sampled[0]["atom_polarizability"].sum(
+            dim=1
+        )
+        self.all_stat["find_atom_polarizability"] = [-1]
         self.all_stat["polarizability"] = [
-            self.all_stat["atomic_polarizability"][0].sum(axis=1)
+            self.all_stat["atom_polarizability"][0].sum(axis=1)
         ]
         self.tfpolar.compute_output_stats(self.all_stat)
         tfbias = self.tfpolar.constant_matrix

source/tests/tf/test_deepmd_data.py

@@ -190,8 +190,8 @@ def test_load_set(self):
             self.assertEqual(data["type"][ii][1], 1)
         self.assertEqual(data["find_coord"], 1)
         self._comp_np_mat2(data["coord"], self.coord)
-        self.assertEqual(data["find_test_atomic"], 1)
-        self._comp_np_mat2(data["test_atomic"], self.test_atomic)
+        self.assertEqual(data["find_test_atom"], 1)
+        self._comp_np_mat2(data["test_atom"], self.test_atomic)
         self.assertEqual(data["find_test_frame"], 1)
         self._comp_np_mat2(data["test_frame"], self.test_frame)
         self.assertEqual(data["find_test_null"], 0)
@@ -207,7 +207,7 @@ def test_shuffle(self):
         data_bk = copy.deepcopy(data)
         data, idx = dd._shuffle_data(data)
         self._comp_np_mat2(data_bk["coord"][idx, :], data["coord"])
-        self._comp_np_mat2(data_bk["test_atomic"][idx, :], data["test_atomic"])
+        self._comp_np_mat2(data_bk["test_atom"][idx, :], data["test_atom"])
         self._comp_np_mat2(data_bk["test_frame"][idx, :], data["test_frame"])
 
     def test_shuffle_with_numb_copy(self):
@@ -224,25 +224,25 @@ def test_shuffle_with_numb_copy(self):
         data, idx = dd._shuffle_data(data)
         assert idx.size == np.sum(prob)
         self._comp_np_mat2(data_bk["coord"][idx, :], data["coord"])
-        self._comp_np_mat2(data_bk["test_atomic"][idx, :], data["test_atomic"])
+        self._comp_np_mat2(data_bk["test_atom"][idx, :], data["test_atom"])
         self._comp_np_mat2(data_bk["test_frame"][idx, :], data["test_frame"])
 
     def test_reduce(self):
         dd = DeepmdData(self.data_name).add("test_atomic", 7, atomic=True, must=True)
         dd.reduce("redu", "test_atomic")
         data = dd._load_set(os.path.join(self.data_name, "set.foo"))
-        self.assertEqual(data["find_test_atomic"], 1)
-        self._comp_np_mat2(data["test_atomic"], self.test_atomic)
+        self.assertEqual(data["find_test_atom"], 1)
+        self._comp_np_mat2(data["test_atom"], self.test_atomic)
         self.assertEqual(data["find_redu"], 1)
         self._comp_np_mat2(data["redu"], self.redu_atomic)
 
     def test_reduce_null(self):
         dd = DeepmdData(self.data_name).add("test_atomic_1", 7, atomic=True, must=False)
         dd.reduce("redu", "test_atomic_1")
         data = dd._load_set(os.path.join(self.data_name, "set.foo"))
-        self.assertEqual(data["find_test_atomic_1"], 0)
+        self.assertEqual(data["find_test_atom_1"], 0)
         self._comp_np_mat2(
-            data["test_atomic_1"], np.zeros([self.nframes, self.natoms * 7])
+            data["test_atom_1"], np.zeros([self.nframes, self.natoms * 7])
         )
         self.assertEqual(data["find_redu"], 0)
         self._comp_np_mat2(data["redu"], np.zeros([self.nframes, 7]))