diff --git a/deepmd/dpmodel/fitting/dipole_fitting.py b/deepmd/dpmodel/fitting/dipole_fitting.py index 6d6324770c..ae9e7a3f0b 100644 --- a/deepmd/dpmodel/fitting/dipole_fitting.py +++ b/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, diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py index 5d75037137..eb31ae0ef1 100644 --- a/deepmd/dpmodel/fitting/polarizability_fitting.py +++ b/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, diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index cad6e12d2b..7b8c227ead 100644 --- a/deepmd/entrypoints/test.py +++ b/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]), ), diff --git a/deepmd/pt/loss/tensor.py b/deepmd/pt/loss/tensor.py index 3dd91d203e..34957815b5 100644 --- a/deepmd/pt/loss/tensor.py +++ b/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, diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index ca445c8588..3c32da9a10 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/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, diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py index efa6249b2a..634cfb23f3 100644 --- a/deepmd/pt/model/task/polarizability.py +++ b/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: diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index c25f725fc8..0a995eaf94 100644 --- a/deepmd/pt/utils/stat.py +++ b/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: diff --git a/source/tests/pt/model/test_polar_stat.py b/source/tests/pt/model/test_polar_stat.py index 3d72c6e8fa..d3f6334182 100644 --- a/source/tests/pt/model/test_polar_stat.py +++ b/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