Chore: refactor atomic bias

deepmd/dpmodel/fitting/dipole_fitting.py

@@ -36,8 +36,6 @@
 
     Parameters
     ----------
-    var_name
-            The name of the output variable.
     ntypes
             The number of atom types.
     dim_descrpt
@@ -86,7 +84,6 @@
 
     def __init__(
         self,
-        var_name: str,
         ntypes: int,
         dim_descrpt: int,
         embedding_width: int,
@@ -124,7 +121,7 @@
         self.r_differentiable = r_differentiable
         self.c_differentiable = c_differentiable
         super().__init__(
-            var_name=var_name,
+            var_name="dipole",
             ntypes=ntypes,
             dim_descrpt=dim_descrpt,
             neuron=neuron,
@@ -161,6 +158,7 @@
     def deserialize(cls, data: dict) -> "GeneralFitting":
         data = copy.deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 1, 1)
+        data.pop("var_name", None)
         return super().deserialize(data)
 
     def output_def(self):

deepmd/dpmodel/fitting/polarizability_fitting.py

@@ -39,8 +39,6 @@
 
     Parameters
     ----------
-    var_name
-            The name of the output variable.
     ntypes
             The number of atom types.
     dim_descrpt
@@ -88,7 +86,6 @@
 
     def __init__(
         self,
-        var_name: str,
         ntypes: int,
         dim_descrpt: int,
         embedding_width: int,
@@ -145,7 +142,7 @@
         self.shift_diag = shift_diag
         self.constant_matrix = np.zeros(ntypes, dtype=GLOBAL_NP_FLOAT_PRECISION)
         super().__init__(
-            var_name=var_name,
+            var_name="polar",
             ntypes=ntypes,
             dim_descrpt=dim_descrpt,
             neuron=neuron,
@@ -201,6 +198,7 @@
     def deserialize(cls, data: dict) -> "GeneralFitting":
         data = copy.deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 2, 1)
+        data.pop("var_name", None)
         return super().deserialize(data)
 
     def output_def(self):

deepmd/entrypoints/test.py

@@ -870,7 +870,7 @@
         arrays with results and their shapes
     """
     data.add(
-        "polarizability" if not atomic else "atomic_polarizability",
+        "polarizability" if not atomic else "atom_polarizability",
         9,
         atomic=atomic,
         must=True,
@@ -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]),
@@ -1011,7 +1011,7 @@
         arrays with results and their shapes
     """
     data.add(
-        "dipole" if not atomic else "atomic_dipole",
+        "dipole" if not atomic else "atom_dipole",
         3,
         atomic=atomic,
         must=True,
@@ -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(
@@ -157,7 +157,7 @@
         if self.has_local_weight:
             label_requirement.append(
                 DataRequirementItem(
-                    "atomic_" + self.label_name,
+                    "atom_" + self.label_name,
                     ndof=self.tensor_size,
                     atomic=True,
                     must=False,

deepmd/pt/model/task/dipole.py

@@ -39,8 +39,6 @@
 
     Parameters
     ----------
-    var_name : str
-        The atomic property to fit, 'dipole'.
     ntypes : int
         Element count.
     dim_descrpt : int
@@ -97,7 +95,7 @@
         self.r_differentiable = r_differentiable
         self.c_differentiable = c_differentiable
         super().__init__(
-            var_name=kwargs.pop("var_name", "dipole"),
+            var_name="dipole",
             ntypes=ntypes,
             dim_descrpt=dim_descrpt,
             neuron=neuron,
@@ -131,6 +129,7 @@
     def deserialize(cls, data: dict) -> "GeneralFitting":
         data = copy.deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 1, 1)
+        data.pop("var_name", None)
         return super().deserialize(data)
 
     def output_def(self) -> FittingOutputDef:

deepmd/pt/model/task/polarizability.py

@@ -47,8 +47,6 @@
 
     Parameters
     ----------
-    var_name : str
-        The atomic property to fit, 'polar'.
     ntypes : int
         Element count.
     dim_descrpt : int
@@ -127,7 +125,7 @@
             ntypes, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE
         )
         super().__init__(
-            var_name=kwargs.pop("var_name", "polar"),
+            var_name="polar",
             ntypes=ntypes,
             dim_descrpt=dim_descrpt,
             neuron=neuron,
@@ -180,6 +178,7 @@
     def deserialize(cls, data: dict) -> "GeneralFitting":
         data = copy.deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 2, 1)
+        data.pop("var_name", None)
         return super().deserialize(data)
 
     def output_def(self) -> FittingOutputDef:
@@ -232,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: