From eff7eb1b82a4633d9482de5baeafccd83642b223 Mon Sep 17 00:00:00 2001 From: root <2000011006@stu.pku.edu.cn> Date: Mon, 23 Dec 2024 12:53:46 +0800 Subject: [PATCH] change bias method to compute_stats_do_not_distinguish_types --- .../model/atomic_model/base_atomic_model.py | 6 +-- deepmd/pt/model/model/property_model.py | 8 ++-- deepmd/pt/model/task/fitting.py | 6 +++ deepmd/pt/model/task/property.py | 8 +++- deepmd/pt/utils/stat.py | 46 ++++++++++--------- deepmd/utils/out_stat.py | 5 +- 6 files changed, 48 insertions(+), 31 deletions(-) diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index e54e748025..3001470d1b 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -459,9 +459,6 @@ 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(), @@ -469,7 +466,8 @@ def change_out_bias( 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: diff --git a/deepmd/pt/model/model/property_model.py b/deepmd/pt/model/model/property_model.py index 79302adb5b..2da7995c80 100644 --- a/deepmd/pt/model/model/property_model.py +++ b/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 diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 2486ab576f..6eee56c8a0 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/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()]) diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index b954e51c29..f0cfa6f5d9 100644 --- a/deepmd/pt/model/task/property.py +++ b/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() diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index a9fc1a2edf..f16d31a215 100644 --- a/deepmd/pt/utils/stat.py +++ b/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 diff --git a/deepmd/utils/out_stat.py b/deepmd/utils/out_stat.py index 4a368e76e8..4198427414 100644 --- a/deepmd/utils/out_stat.py +++ b/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 -------