fix: change atomic to atom

deepmd/dpmodel/fitting/dipole_fitting.py

@@ -36,8 +36,6 @@ class DipoleFitting(GeneralFitting):
 
     Parameters
     ----------
-    var_name
-            The name of the output variable.
     ntypes
             The number of atom types.
     dim_descrpt
@@ -86,7 +84,6 @@ class DipoleFitting(GeneralFitting):
 
     def __init__(
         self,
-        var_name: str,
         ntypes: int,
         dim_descrpt: int,
         embedding_width: int,
@@ -124,7 +121,7 @@ def __init__(
         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,

deepmd/dpmodel/fitting/polarizability_fitting.py

@@ -39,8 +39,6 @@ class PolarFitting(GeneralFitting):
 
     Parameters
     ----------
-    var_name
-            The name of the output variable.
     ntypes
             The number of atom types.
     dim_descrpt
@@ -88,7 +86,6 @@ class PolarFitting(GeneralFitting):
 
     def __init__(
         self,
-        var_name: str,
         ntypes: int,
         dim_descrpt: int,
         embedding_width: int,
@@ -145,7 +142,7 @@ def __init__(
         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,

deepmd/entrypoints/test.py

@@ -870,7 +870,7 @@ def test_polar(
         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 @@ def test_polar(
         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 @@ def test_polar(
             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 @@ def test_dipole(
         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 @@ def test_dipole(
         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 @@ def test_dipole(
             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 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False
         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 @@ def label_requirement(self) -> List[DataRequirementItem]:
         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 @@ class DipoleFittingNet(GeneralFitting):
 
     Parameters
     ----------
-    var_name : str
-        The atomic property to fit, 'dipole'.
     ntypes : int
         Element count.
     dim_descrpt : int
@@ -97,7 +95,7 @@ def __init__(
         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,

deepmd/pt/model/task/polarizability.py

@@ -47,8 +47,6 @@ class PolarFittingNet(GeneralFitting):
 
     Parameters
     ----------
-    var_name : str
-        The atomic property to fit, 'polar'.
     ntypes : int
         Element count.
     dim_descrpt : int
@@ -127,7 +125,7 @@ def __init__(
             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,
@@ -233,9 +231,9 @@ def compute_output_stats(
                 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/pt/utils/stat.py

@@ -83,37 +83,43 @@ def make_stat_input(datasets, dataloaders, nbatches):
 def compute_output_stats(
     merged: Union[Callable[[], List[dict]], List[dict]],
     ntypes: int,
+    keys: List[str],
     stat_file_path: Optional[DPPath] = None,
     rcond: Optional[float] = None,
     atom_ener: Optional[List[float]] = None,
     model_forward: Optional[Callable[..., torch.Tensor]] = None,
-    keys: Optional[str] = "energy",  # this is dict.keys()
+
 ):
-    if "energy" in keys:
-        return compute_output_stats_global_only(
-            merged=merged,
-            ntypes=ntypes,
-            stat_file_path=stat_file_path,
-            rcond=rcond,
-            atom_ener=atom_ener,
-            model_forward=model_forward,
-        )
-    elif (
-        len({"dos", "atom_dos", "polarizability", "atomic_polarizability"} & set(keys))
-        > 0
-    ):
-        return compute_output_stats_with_atomic(
-            merged=merged,
-            ntypes=ntypes,
-            keys=list(keys),
-            stat_file_path=stat_file_path,
-            rcond=rcond,
-            atom_ener=atom_ener,
-            model_forward=model_forward,
-        )
-    else:
-        # can add mode facade services.
-        pass
+    key_mapping = {
+        "polar":"polarizability",
+        "energy": "energy",
+        "dos": "dos",
+        "dipole": "dipole",
+    }
+
+    for out_put in keys:
+        if out_put == "energy":
+            return compute_output_stats_global_only(
+                merged=merged,
+                ntypes=ntypes,
+                stat_file_path=stat_file_path,
+                rcond=rcond,
+                atom_ener=atom_ener,
+                model_forward=model_forward,
+            )
+        elif out_put in ["dos", "polar"]:
+            return compute_output_stats_with_atomic(
+                merged=merged,
+                ntypes=ntypes,
+                key=key_mapping[out_put],
+                stat_file_path=stat_file_path,
+                rcond=rcond,
+                atom_ener=atom_ener,
+                model_forward=model_forward,
+            )
+        else:
+            # can add mode facade services.
+            pass
 
 
 def compute_output_stats_global_only(
@@ -248,7 +254,7 @@ def model_forward_auto_batch_size(*args, **kwargs):
 def compute_output_stats_with_atomic(
     merged: Union[Callable[[], List[dict]], List[dict]],
     ntypes: int,
-    keys: List[str],
+    key: str,
     stat_file_path: Optional[DPPath] = None,
     rcond: Optional[float] = None,
     atom_ener: Optional[List[float]] = None,
@@ -268,8 +274,8 @@ def compute_output_stats_with_atomic(
             the lazy function helps by only sampling once.
     ntypes : int
         The number of atom types.
-    keys : List[str]
-        The fitting output keys.
+    key : str
+        The var_name of the fitting net.
     stat_file_path : DPPath, optional
         The path to the stat file.
     rcond : float, optional
@@ -282,19 +288,11 @@ def compute_output_stats_with_atomic(
         which will be subtracted from the energy label of the data.
         The difference will then be used to calculate the delta complement energy bias for each type.
     """
-    if "dos" in keys or "atom_dos" in keys:
-        atomic_label_name = "atom_dos"
-        global_label_name = "dos"
-        file_label_name = "bias_dos"
-    elif "polarizability" in keys or "atomic_polarizability" in keys:
-        atomic_label_name = "atomic_polarizability"
-        global_label_name = "polarizability"
-        file_label_name = "constant_matrix"
-    else:
-        raise NotImplementedError
+
+    atomic_label_name, global_label_name = "atom_" + key, key
 
     if stat_file_path is not None:
-        stat_file_path = stat_file_path / file_label_name
+        stat_file_path = stat_file_path / "atomic_bias"
     if stat_file_path is not None and stat_file_path.is_file():
         total_bias = stat_file_path.load_numpy()
     else:
@@ -363,7 +361,7 @@ def model_forward_auto_batch_size(*args, **kwargs):
                 for itype in range(ntypes):
                     type_mask = system["atype"] == itype
                     sys_type_count[:, itype] = type_mask.sum(dim=1).numpy(force=True)
-                sys_bias_redu = system["dos"].numpy(force=True)
+                sys_bias_redu = system[global_label_name].numpy(force=True)
                 if property_predict is None:
                     sys_bias = compute_stats_from_redu(
                         sys_bias_redu, sys_type_count, rcond=rcond
@@ -374,20 +372,8 @@ def model_forward_auto_batch_size(*args, **kwargs):
                         bias_diff, sys_type_count, rcond=rcond
                     )[0]
 
-            if global_label_name == "dos":
-                total_bias.append(sys_bias)
-            elif global_label_name == "polarizability":
-                cur_constant_matrix = np.zeros(
-                    ntypes, dtype=env.GLOBAL_NP_FLOAT_PRECISION
-                )
-
-                for itype in range(ntypes):
-                    cur_constant_matrix[itype] = np.mean(
-                        np.diagonal(sys_bias[itype].reshape(3, 3))
-                    )
-                total_bias.append(cur_constant_matrix)
-            else:
-                raise NotImplementedError
+            total_bias.append(sys_bias)
+            # need to take care shift diag and add atom_ener
         total_bias = np.stack(total_bias).mean(axis=0)
         total_bias = np.nan_to_num(total_bias)
         if stat_file_path is not None:

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]["find_atom_polarizability"] = -1
         self.sampled[0]["polarizability"] = self.sampled[0][
-            "atomic_polarizability"
+            "atom_polarizability"
         ].sum(dim=1)
-        self.all_stat["find_atomic_polarizability"] = [-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