From f591c6b0a1570d0fa8314b3685be917a6412dd27 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Tue, 2 Apr 2024 14:33:54 +0800 Subject: [PATCH] feat: add dos stat UT --- .../dpmodel/atomic_model/dp_atomic_model.py | 8 +++-- .../atomic_model/make_base_atomic_model.py | 6 ++-- .../model/atomic_model/base_atomic_model.py | 4 +-- .../pt/model/atomic_model/dp_atomic_model.py | 6 ++-- deepmd/pt/utils/stat.py | 2 +- source/tests/pt/model/test_permutation.py | 4 +-- source/tests/pt/test_finetune.py | 31 ++++++++++++++++--- source/tests/pt/test_stat.py | 17 ++++++++++ 8 files changed, 60 insertions(+), 18 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 8a40f8d238..aef561510e 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -83,21 +83,23 @@ def mixed_types(self) -> bool: """ return self.descriptor.mixed_types() - def set_out_bias(self, out_bias: np.ndarray, add=False) -> None: + def set_out_bias(self, out_bias: Dict[str, np.ndarray], add=False) -> None: """ Modify the output bias for the atomic model. Parameters ---------- - out_bias : np.ndarray + out_bias : Dict[str, np.ndarray] The new bias to be applied. add : bool, optional Whether to add the new bias to the existing one. If False, the output bias will be directly replaced by the new bias. If True, the new bias will be added to the existing one. """ + #TODO: refactor for multiple properties + bias_keys = list(self.fitting_output_def().keys()) self.fitting["bias_atom_e"] = ( - out_bias + self.fitting["bias_atom_e"] if add else out_bias + out_bias[bias_keys[0]] + self.fitting["bias_atom_e"] if add else out_bias ) def get_out_bias(self) -> np.ndarray: diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py index 3e02a5d076..c005ce8e9b 100644 --- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py @@ -96,14 +96,14 @@ def get_sel_type(self) -> List[int]: """ @abstractmethod - def set_out_bias(self, out_bias: t_tensor, add=False) -> None: + def set_out_bias(self, out_bias: Dict[str,t_tensor], add=False) -> None: """ Modify the output bias for the atomic model. Parameters ---------- - out_bias : t_tensor - The new bias to be applied. + out_bias : Dict[str, t_tensor] + The new bias to be applied for multiple output properties. add : bool, optional Whether to add the new bias to the existing one. If False, the output bias will be directly replaced by the new bias. diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index 83be49d1c1..1feb2caede 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -285,14 +285,14 @@ def change_out_bias( self.get_ntypes(), keys=list(self.fitting_output_def().keys()), model_forward=self.get_forward_wrapper_func(), - )[list(self.fitting_output_def().keys())[0]] + ) self.set_out_bias(delta_bias, add=True) elif bias_adjust_mode == "set-by-statistic": bias_atom = compute_output_stats( merged, self.get_ntypes(), keys=list(self.fitting_output_def().keys()), - )[list(self.fitting_output_def().keys())[0]] + ) self.set_out_bias(bias_atom) else: raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index 13b8f09a79..1062766ceb 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -223,7 +223,7 @@ def wrapped_sampler(): if self.fitting_net is not None: self.fitting_net.compute_output_stats(wrapped_sampler, stat_file_path) - def set_out_bias(self, out_bias: torch.Tensor, add=False) -> None: + def set_out_bias(self, out_bias: Dict[str, torch.Tensor], add=False) -> None: """ Modify the output bias for the atomic model. @@ -236,8 +236,10 @@ def set_out_bias(self, out_bias: torch.Tensor, add=False) -> None: If False, the output bias will be directly replaced by the new bias. If True, the new bias will be added to the existing one. """ + #TODO: refactor for multiple properties + bias_keys = list(self.fitting_output_def().keys()) self.fitting_net["bias_atom_e"] = ( - out_bias + self.fitting_net["bias_atom_e"] if add else out_bias + out_bias[bias_keys[0]] + self.fitting_net["bias_atom_e"] if add else out_bias ) def get_out_bias(self) -> torch.Tensor: diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 007e4f2983..ae0af841f6 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -423,7 +423,7 @@ def model_forward_auto_batch_size(*args, **kwargs): sys_bias_redu, sys_type_count, rcond=rcond )[0] else: - bias_diff = sys_bias_redu - model_predict[kk][idx].sum(-1) + bias_diff = sys_bias_redu - model_predict[kk][idx].sum(axis=0) sys_bias = compute_stats_from_redu( bias_diff, sys_type_count, rcond=rcond )[0] diff --git a/source/tests/pt/model/test_permutation.py b/source/tests/pt/model/test_permutation.py index 5e395eb8c0..527798c757 100644 --- a/source/tests/pt/model/test_permutation.py +++ b/source/tests/pt/model/test_permutation.py @@ -37,10 +37,10 @@ } model_dos = { - "type_map": ["O", "H", "B"], + "type_map": ["O", "H"], "descriptor": { "type": "se_e2_a", - "sel": [46, 92, 4], + "sel": [46, 92], "rcut_smth": 0.50, "rcut": 4.00, "neuron": [25, 50, 100], diff --git a/source/tests/pt/test_finetune.py b/source/tests/pt/test_finetune.py index 79f8c57cb8..410e9e2da6 100644 --- a/source/tests/pt/test_finetune.py +++ b/source/tests/pt/test_finetune.py @@ -31,9 +31,11 @@ model_dpa2, model_se_e2_a, model_zbl, + model_dos, ) from .test_stat import ( energy_data_requirement, + dos_data_requirement, ) @@ -44,7 +46,7 @@ def test_finetune_change_out_bias(self): fitting_net = model.get_fitting_net() fitting_net["bias_atom_e"] = torch.rand_like(fitting_net["bias_atom_e"]) energy_bias_before = deepcopy( - to_numpy_array(fitting_net["bias_atom_e"]).reshape(-1) + to_numpy_array(fitting_net["bias_atom_e"]) ) # prepare original model for test @@ -63,7 +65,7 @@ def test_finetune_change_out_bias(self): full_type_map=full_type_map, ) energy_bias_after = deepcopy( - to_numpy_array(fitting_net["bias_atom_e"]).reshape(-1) + to_numpy_array(fitting_net["bias_atom_e"]) ) # get ground-truth energy bias change @@ -81,17 +83,33 @@ def test_finetune_change_out_bias(self): to_numpy_array(self.sampled[0]["box"][:ntest]), to_numpy_array(self.sampled[0]["atype"][0]), )[0] - energy_diff = to_numpy_array(self.sampled[0]["energy"][:ntest]) - energy + energy_diff = to_numpy_array(self.sampled[0][self.var_name][:ntest]) - energy finetune_shift = ( energy_bias_after[idx_type_map] - energy_bias_before[idx_type_map] ) ground_truth_shift = np.linalg.lstsq(atom_nums, energy_diff, rcond=None)[ 0 - ].reshape(-1) - + ] # check values np.testing.assert_almost_equal(finetune_shift, ground_truth_shift, decimal=10) +class TestDOSModelSeA(unittest.TestCase, FinetuneTest): + def setUp(self): + self.data_file = [str(Path(__file__).parent / "dos/data/global_system")] + self.model_config = model_dos + self.data = DpLoaderSet( + self.data_file, + batch_size=1, + type_map=self.model_config["type_map"], + ) + self.data.add_data_requirement(dos_data_requirement) + self.sampled = make_stat_input( + self.data.systems, + self.data.dataloaders, + nbatches=1, + ) + self.var_name = "dos" + class TestEnergyModelSeA(unittest.TestCase, FinetuneTest): def setUp(self): @@ -108,6 +126,7 @@ def setUp(self): self.data.dataloaders, nbatches=1, ) + self.var_name = "energy" @unittest.skip("change bias not implemented yet.") @@ -126,6 +145,7 @@ def setUp(self): self.data.dataloaders, nbatches=1, ) + self.var_name = "energy" class TestEnergyModelDPA2(unittest.TestCase, FinetuneTest): @@ -143,6 +163,7 @@ def setUp(self): self.data.dataloaders, nbatches=1, ) + self.var_name = "energy" if __name__ == "__main__": diff --git a/source/tests/pt/test_stat.py b/source/tests/pt/test_stat.py index 2362821dfa..81cf524d91 100644 --- a/source/tests/pt/test_stat.py +++ b/source/tests/pt/test_stat.py @@ -62,6 +62,23 @@ CUR_DIR = os.path.dirname(__file__) +dos_data_requirement = [ + DataRequirementItem( + "dos", + ndof=250, + atomic=False, + must=False, + high_prec=True, + ), + DataRequirementItem( + "atom_dos", + ndof=250, + atomic=True, + must=False, + high_prec=False, + ), +] + energy_data_requirement = [ DataRequirementItem( "energy",