change bias method to compute_stats_do_not_distinguish_types

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -459,17 +459,15 @@ def change_out_bias(
             )
             self._store_out_stat(delta_bias, out_std, add=True)
         elif bias_adjust_mode == "set-by-statistic":
-            property_name = (
-                self.fitting_net.property_name if "property" in self.bias_keys else None
-            )
             bias_out, std_out = compute_output_stats(
                 sample_merged,
                 self.get_ntypes(),
                 keys=list(self.atomic_output_def().keys()),
                 stat_file_path=stat_file_path,
                 rcond=self.rcond,
                 preset_bias=self.preset_out_bias,
-                property_name=property_name,
+                stats_do_not_distinguish_types=self.fitting_net.get_compute_stats_do_not_distinguish_types(),
+                intensive=self.fitting_net.intensive,
             )
             self._store_out_stat(bias_out, std_out)
         else:

deepmd/pt/model/model/property_model.py

@@ -62,8 +62,8 @@ def forward(
             do_atomic_virial=do_atomic_virial,
         )
         model_predict = {}
-        model_predict["atom_property"] = model_ret["property"]
-        model_predict["property"] = model_ret["property_redu"]
+        model_predict["atom_property"] = model_ret[self.get_property_name()]
+        model_predict["property"] = model_ret[f"{self.get_property_name()}_redu"]
         if "mask" in model_ret:
             model_predict["mask"] = model_ret["mask"]
         return model_predict
@@ -107,8 +107,8 @@ def forward_lower(
             extra_nlist_sort=self.need_sorted_nlist_for_lower(),
         )
         model_predict = {}
-        model_predict["atom_property"] = model_ret["property"]
-        model_predict["property"] = model_ret["property_redu"]
+        model_predict["atom_property"] = model_ret[self.get_property_name()]
+        model_predict["property"] = model_ret[f"{self.get_property_name()}_redu"]
         if "mask" in model_ret:
             model_predict["mask"] = model_ret["mask"]
         return model_predict

deepmd/pt/model/task/fitting.py

@@ -367,6 +367,12 @@ def set_case_embd(self, case_idx: int):
             case_idx
         ]
 
+    def get_compute_stats_do_not_distinguish_types(self) -> bool:
+        """
+        Get whether the fitting net computes stats which are not distinguished between different types of atoms.
+        """
+        return False
+
     def __setitem__(self, key, value) -> None:
         if key in ["bias_atom_e"]:
             value = value.view([self.ntypes, self._net_out_dim()])

deepmd/pt/model/task/property.py

@@ -95,7 +95,7 @@ def __init__(
         self.intensive = intensive
         self.property_name = property_name
         super().__init__(
-            var_name="property",
+            var_name=self.property_name,
             ntypes=ntypes,
             dim_descrpt=dim_descrpt,
             dim_out=task_dim,
@@ -126,6 +126,12 @@ def output_def(self) -> FittingOutputDef:
             ]
         )
 
+    def get_compute_stats_do_not_distinguish_types(self) -> bool:
+        """
+        Get whether the fitting net computes stats which are not distinguished between different types of atoms.
+        """
+        return True
+
     @classmethod
     def deserialize(cls, data: dict) -> "PropertyFittingNet":
         data = data.copy()

deepmd/pt/utils/stat.py

@@ -26,7 +26,7 @@
 from deepmd.utils.out_stat import (
     compute_stats_from_atomic,
     compute_stats_from_redu,
-    compute_stats_property,
+    compute_stats_do_not_distinguish_types,
 )
 from deepmd.utils.path import (
     DPPath,
@@ -240,7 +240,8 @@ def compute_output_stats(
     rcond: Optional[float] = None,
     preset_bias: Optional[dict[str, list[Optional[np.ndarray]]]] = None,
     model_forward: Optional[Callable[..., torch.Tensor]] = None,
-    property_name: Optional[str] = None,
+    stats_do_not_distinguish_types: bool = False,
+    intensive: bool = False,
 ):
     """
     Compute the output statistics (e.g. energy bias) for the fitting net from packed data.
@@ -289,9 +290,6 @@ def compute_output_stats(
 
         # remove the keys that are not in the sample
         keys = [keys] if isinstance(keys, str) else keys
-        if property_name is not None:
-            del keys
-            keys = [property_name]
         assert isinstance(keys, list)
         new_keys = [
             ii
@@ -365,7 +363,8 @@ def compute_output_stats(
             rcond,
             preset_bias,
             model_pred_g,
-            property_name,
+            stats_do_not_distinguish_types,
+            intensive,
         )
         bias_atom_a, std_atom_a = compute_output_stats_atomic(
             sampled,
@@ -376,7 +375,6 @@ def compute_output_stats(
 
         # merge global/atomic bias
         bias_atom_e, std_atom_e = {}, {}
-        keys = ["property"] if (property_name is not None) else keys
         for kk in keys:
             # use atomic bias whenever available
             if kk in bias_atom_a:
@@ -409,7 +407,8 @@ def compute_output_stats_global(
     rcond: Optional[float] = None,
     preset_bias: Optional[dict[str, list[Optional[np.ndarray]]]] = None,
     model_pred: Optional[dict[str, np.ndarray]] = None,
-    property_name: Optional[str] = None,
+    stats_do_not_distinguish_types: bool = False,
+    intensive: bool = False,
 ):
     """This function only handle stat computation from reduced global labels."""
     # return directly if model predict is empty for global
@@ -478,25 +477,30 @@ def compute_output_stats_global(
 
     bias_atom_e = {}
     std_atom_e = {}
-    for kk in keys:
-        if kk in stats_input:
-            if property_name is not None:
-                assert len(keys) == 1
-                bias_atom_e["property"], std_atom_e["property"] = (
-                    compute_stats_property(
+    if stats_do_not_distinguish_types:
+        for kk in keys:
+            if kk in stats_input:
+                bias_atom_e[kk], std_atom_e[kk] = (
+                    compute_stats_do_not_distinguish_types(
                         stats_input[kk],
                         merged_natoms[kk],
                         assigned_bias=assigned_atom_ener[kk],
+                        intensive=intensive,
                     )
                 )
-                return bias_atom_e, std_atom_e
             else:
-                bias_atom_e[kk], std_atom_e[kk] = compute_stats_from_redu(
-                    stats_input[kk],
-                    merged_natoms[kk],
-                    assigned_bias=assigned_atom_ener[kk],
-                    rcond=rcond,
-                )
+                # this key does not have global labels, skip it.
+                continue
+        return bias_atom_e, std_atom_e
+
+    for kk in keys:
+        if kk in stats_input:
+            bias_atom_e[kk], std_atom_e[kk] = compute_stats_from_redu(
+                stats_input[kk],
+                merged_natoms[kk],
+                assigned_bias=assigned_atom_ener[kk],
+                rcond=rcond,
+            )
         else:
             # this key does not have global labels, skip it.
             continue

deepmd/utils/out_stat.py

@@ -132,10 +132,11 @@ def compute_stats_from_atomic(
     return output_bias, output_std
 
 
-def compute_stats_property(
+def compute_stats_do_not_distinguish_types(
     output_redu: np.ndarray,
     natoms: np.ndarray,
     assigned_bias: Optional[np.ndarray] = None,
+    intensive: bool = False,
 ) -> tuple[np.ndarray, np.ndarray]:
     """Compute the mean value and standard deviation of reduced output.
 
@@ -152,6 +153,8 @@ def compute_stats_property(
         The assigned output bias, shape is [ntypes, *(odim0, odim1, ...)].
         Set to a tensor of shape (odim0, odim1, ...) filled with nan if the bias
         of the type is not assigned.
+    intensive
+        Whether the output is intensive or extensive.
 
     Returns
     -------