From fa8e6453b5cc88ae575284c4c0ded06758c6de82 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Wed, 6 Mar 2024 17:23:29 +0800 Subject: [PATCH] Feat: add zbl training (#3398) Signed-off-by: Anyang Peng <137014849+anyangml@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- deepmd/dpmodel/atomic_model/__init__.py | 8 +- .../atomic_model/linear_atomic_model.py | 62 ++++++++--- .../atomic_model/make_base_atomic_model.py | 11 +- .../atomic_model/pairtab_atomic_model.py | 11 +- deepmd/pt/model/atomic_model/__init__.py | 8 +- .../model/atomic_model/linear_atomic_model.py | 100 +++++++++++++++--- .../atomic_model/pairtab_atomic_model.py | 85 ++++++++++++++- deepmd/pt/model/model/dp_zbl_model.py | 4 +- deepmd/pt/model/task/ener.py | 49 ++------- deepmd/pt/train/training.py | 29 +++-- deepmd/pt/utils/stat.py | 81 ++++++++++++++ .../dpmodel/test_linear_atomic_model.py | 8 +- .../pt/model/test_linear_atomic_model.py | 32 ++++-- .../pt/model/test_pairtab_atomic_model.py | 2 +- source/tests/pt/model/water/zbl.json | 92 ++++++++++++++++ source/tests/pt/test_finetune.py | 51 +++++++-- source/tests/pt/test_training.py | 18 ++++ 17 files changed, 519 insertions(+), 132 deletions(-) create mode 100644 source/tests/pt/model/water/zbl.json diff --git a/deepmd/dpmodel/atomic_model/__init__.py b/deepmd/dpmodel/atomic_model/__init__.py index 2cd20f54c1..e51ca0a65e 100644 --- a/deepmd/dpmodel/atomic_model/__init__.py +++ b/deepmd/dpmodel/atomic_model/__init__.py @@ -22,8 +22,8 @@ DPAtomicModel, ) from .linear_atomic_model import ( - DPZBLLinearAtomicModel, - LinearAtomicModel, + DPZBLLinearEnergyAtomicModel, + LinearEnergyAtomicModel, ) from .make_base_atomic_model import ( make_base_atomic_model, @@ -37,6 +37,6 @@ "BaseAtomicModel", "DPAtomicModel", "PairTabAtomicModel", - "LinearAtomicModel", - "DPZBLLinearAtomicModel", + "LinearEnergyAtomicModel", + "DPZBLLinearEnergyAtomicModel", ] diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index b1a32cdaa5..ac2a73a381 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -38,7 +38,7 @@ ) -class LinearAtomicModel(BaseAtomicModel): +class LinearEnergyAtomicModel(BaseAtomicModel): """Linear model make linear combinations of several existing models. Parameters @@ -59,14 +59,16 @@ def __init__( self.models = models sub_model_type_maps = [md.get_type_map() for md in models] err_msg = [] + self.mapping_list = [] common_type_map = set(type_map) + self.type_map = type_map for tpmp in sub_model_type_maps: if not common_type_map.issubset(set(tpmp)): err_msg.append( f"type_map {tpmp} is not a subset of type_map {type_map}" ) + self.mapping_list.append(self.remap_atype(tpmp, self.type_map)) assert len(err_msg) == 0, "\n".join(err_msg) - self.type_map = type_map self.mixed_types_list = [model.mixed_types() for model in self.models] super().__init__(**kwargs) @@ -163,17 +165,20 @@ def forward_atomic( self.mixed_types_list, raw_nlists, self.get_model_sels() ) ] - ener_list = [ - model.forward_atomic( - extended_coord, - extended_atype, - nl, - mapping, - fparam, - aparam, - )["energy"] - for model, nl in zip(self.models, nlists_) - ] + ener_list = [] + + for i, model in enumerate(self.models): + mapping = self.mapping_list[i] + ener_list.append( + model.forward_atomic( + extended_coord, + mapping[extended_atype], + nlists_[i], + mapping, + fparam, + aparam, + )["energy"] + ) self.weights = self._compute_weight(extended_coord, extended_atype, nlists_) self.atomic_bias = None if self.atomic_bias is not None: @@ -184,6 +189,29 @@ def forward_atomic( } # (nframes, nloc, 1) return fit_ret + @staticmethod + def remap_atype(ori_map: List[str], new_map: List[str]) -> np.ndarray: + """ + This method is used to map the atype from the common type_map to the original type_map of + indivial AtomicModels. + + Parameters + ---------- + ori_map : List[str] + The original type map of an AtomicModel. + new_map : List[str] + The common type map of the DPZBLLinearEnergyAtomicModel, created by the `get_type_map` method, + must be a subset of the ori_map. + + Returns + ------- + np.ndarray + """ + type_2_idx = {atp: idx for idx, atp in enumerate(ori_map)} + # this maps the atype in the new map to the original map + mapping = np.array([type_2_idx[new_map[idx]] for idx in range(len(new_map))]) + return mapping + def fitting_output_def(self) -> FittingOutputDef: return FittingOutputDef( [ @@ -261,7 +289,7 @@ def is_aparam_nall(self) -> bool: return False -class DPZBLLinearAtomicModel(LinearAtomicModel): +class DPZBLLinearEnergyAtomicModel(LinearEnergyAtomicModel): """Model linearly combine a list of AtomicModels. Parameters @@ -308,7 +336,7 @@ def serialize(self) -> dict: "@class": "Model", "type": "zbl", "@version": 1, - "models": LinearAtomicModel.serialize( + "models": LinearEnergyAtomicModel.serialize( [self.dp_model, self.zbl_model], self.type_map ), "sw_rmin": self.sw_rmin, @@ -319,7 +347,7 @@ def serialize(self) -> dict: return dd @classmethod - def deserialize(cls, data) -> "DPZBLLinearAtomicModel": + def deserialize(cls, data) -> "DPZBLLinearEnergyAtomicModel": data = copy.deepcopy(data) check_version_compatibility(data.pop("@version", 1), 1, 1) data.pop("@class") @@ -328,7 +356,7 @@ def deserialize(cls, data) -> "DPZBLLinearAtomicModel": sw_rmax = data.pop("sw_rmax") smin_alpha = data.pop("smin_alpha") - ([dp_model, zbl_model], type_map) = LinearAtomicModel.deserialize( + ([dp_model, zbl_model], type_map) = LinearEnergyAtomicModel.deserialize( data.pop("models") ) diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py index 5548147d54..ce1a6708e6 100644 --- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py @@ -54,18 +54,13 @@ def get_rcut(self) -> float: pass @abstractmethod - def get_type_map(self) -> Optional[List[str]]: + def get_type_map(self) -> List[str]: """Get the type map.""" + pass def get_ntypes(self) -> int: """Get the number of atom types.""" - tmap = self.get_type_map() - if tmap is not None: - return len(tmap) - else: - raise ValueError( - "cannot infer the number of types from a None type map" - ) + return len(self.get_type_map()) @abstractmethod def get_sel(self) -> List[int]: diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index c858179939..46ec808ad4 100644 --- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py @@ -66,9 +66,17 @@ def __init__( self.type_map = type_map self.tab = PairTab(self.tab_file, rcut=rcut) + self.type_map = type_map + self.ntypes = len(type_map) if self.tab_file is not None: self.tab_info, self.tab_data = self.tab.get() + nspline, ntypes_tab = self.tab_info[-2:].astype(int) + self.tab_data = self.tab_data.reshape(ntypes_tab, ntypes_tab, nspline, 4) + if self.ntypes != ntypes_tab: + raise ValueError( + "The `type_map` provided does not match the number of columns in the table." + ) else: self.tab_info, self.tab_data = None, None @@ -145,7 +153,8 @@ def deserialize(cls, data) -> "PairTabAtomicModel": tab_model = cls(None, rcut, sel, type_map, **data) tab_model.tab = tab tab_model.tab_info = tab_model.tab.tab_info - tab_model.tab_data = tab_model.tab.tab_data + nspline, ntypes = tab_model.tab_info[-2:].astype(int) + tab_model.tab_data = tab_model.tab.tab_data.reshape(ntypes, ntypes, nspline, 4) return tab_model def forward_atomic( diff --git a/deepmd/pt/model/atomic_model/__init__.py b/deepmd/pt/model/atomic_model/__init__.py index 75c1ce3c2e..a747f28556 100644 --- a/deepmd/pt/model/atomic_model/__init__.py +++ b/deepmd/pt/model/atomic_model/__init__.py @@ -21,8 +21,8 @@ DPAtomicModel, ) from .linear_atomic_model import ( - DPZBLLinearAtomicModel, - LinearAtomicModel, + DPZBLLinearEnergyAtomicModel, + LinearEnergyAtomicModel, ) from .pairtab_atomic_model import ( PairTabAtomicModel, @@ -32,6 +32,6 @@ "BaseAtomicModel", "DPAtomicModel", "PairTabAtomicModel", - "LinearAtomicModel", - "DPZBLLinearAtomicModel", + "LinearEnergyAtomicModel", + "DPZBLLinearEnergyAtomicModel", ] diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 0dd1b13723..5e1a80087e 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -25,6 +25,9 @@ get_multiple_nlist_key, nlist_distinguish_types, ) +from deepmd.utils.path import ( + DPPath, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -40,7 +43,7 @@ ) -class LinearAtomicModel(torch.nn.Module, BaseAtomicModel): +class LinearEnergyAtomicModel(torch.nn.Module, BaseAtomicModel): """Linear model make linear combinations of several existing models. Parameters @@ -62,14 +65,17 @@ def __init__( self.models = torch.nn.ModuleList(models) sub_model_type_maps = [md.get_type_map() for md in models] err_msg = [] + self.mapping_list = [] common_type_map = set(type_map) + self.type_map = type_map for tpmp in sub_model_type_maps: if not common_type_map.issubset(set(tpmp)): err_msg.append( f"type_map {tpmp} is not a subset of type_map {type_map}" ) + self.mapping_list.append(self.remap_atype(tpmp, self.type_map)) assert len(err_msg) == 0, "\n".join(err_msg) - self.type_map = type_map + self.atomic_bias = None self.mixed_types_list = [model.mixed_types() for model in self.models] BaseAtomicModel.__init__(self, **kwargs) @@ -117,8 +123,8 @@ def _sort_rcuts_sels(self, device: torch.device) -> Tuple[List[float], List[int] nsels = torch.tensor(self.get_model_nsels(), device=device) zipped = torch.stack( [ - torch.tensor(rcuts, device=device), - torch.tensor(nsels, device=device), + rcuts, + nsels, ], dim=0, ).T @@ -185,10 +191,11 @@ def forward_atomic( ener_list = [] for i, model in enumerate(self.models): + mapping = self.mapping_list[i] ener_list.append( model.forward_atomic( extended_coord, - extended_atype, + mapping[extended_atype], nlists_[i], mapping, fparam, @@ -198,16 +205,48 @@ def forward_atomic( weights = self._compute_weight(extended_coord, extended_atype, nlists_) - if self.atomic_bias is not None: - raise NotImplementedError("Need to add bias in a future PR.") - else: - fit_ret = { - "energy": torch.sum( - torch.stack(ener_list) * torch.stack(weights), dim=0 - ), - } # (nframes, nloc, 1) + atype = extended_atype[:, :nloc] + for idx, model in enumerate(self.models): + # TODO: provide interfaces for atomic models to access bias_atom_e + if isinstance(model, DPAtomicModel): + bias_atom_e = model.fitting_net.bias_atom_e + elif isinstance(model, PairTabAtomicModel): + bias_atom_e = model.bias_atom_e + else: + bias_atom_e = None + if bias_atom_e is not None: + ener_list[idx] += bias_atom_e[atype] + + fit_ret = { + "energy": torch.sum(torch.stack(ener_list) * torch.stack(weights), dim=0), + } # (nframes, nloc, 1) return fit_ret + @staticmethod + def remap_atype(ori_map: List[str], new_map: List[str]) -> torch.Tensor: + """ + This method is used to map the atype from the common type_map to the original type_map of + indivial AtomicModels. It creates a index mapping for the conversion. + + Parameters + ---------- + ori_map : List[str] + The original type map of an AtomicModel. + new_map : List[str] + The common type map of the DPZBLLinearEnergyAtomicModel, created by the `get_type_map` method, + must be a subset of the ori_map. + + Returns + ------- + torch.Tensor + """ + type_2_idx = {atp: idx for idx, atp in enumerate(ori_map)} + # this maps the atype in the new map to the original map + mapping = torch.tensor( + [type_2_idx[new_map[idx]] for idx in range(len(new_map))], device=env.DEVICE + ) + return mapping + def fitting_output_def(self) -> FittingOutputDef: return FittingOutputDef( [ @@ -292,7 +331,7 @@ def is_aparam_nall(self) -> bool: return False -class DPZBLLinearAtomicModel(LinearAtomicModel): +class DPZBLLinearEnergyAtomicModel(LinearEnergyAtomicModel): """Model linearly combine a list of AtomicModels. Parameters @@ -336,6 +375,33 @@ def __init__( # this is a placeholder being updated in _compute_weight, to handle Jit attribute init error. self.zbl_weight = torch.empty(0, dtype=torch.float64, device=env.DEVICE) + def compute_or_load_stat( + self, + sampled_func, + stat_file_path: Optional[DPPath] = None, + ): + """ + Compute or load the statistics parameters of the model, + such as mean and standard deviation of descriptors or the energy bias of the fitting net. + When `sampled` is provided, all the statistics parameters will be calculated (or re-calculated for update), + and saved in the `stat_file_path`(s). + When `sampled` is not provided, it will check the existence of `stat_file_path`(s) + and load the calculated statistics parameters. + + Parameters + ---------- + sampled_func + The lazy sampled function to get data frames from different data systems. + stat_file_path + The dictionary of paths to the statistics files. + """ + self.dp_model.compute_or_load_stat(sampled_func, stat_file_path) + self.zbl_model.compute_or_load_stat(sampled_func, stat_file_path) + + def change_energy_bias(self): + # need to implement + pass + def serialize(self) -> dict: dd = BaseAtomicModel.serialize(self) dd.update( @@ -343,7 +409,7 @@ def serialize(self) -> dict: "@class": "Model", "@version": 1, "type": "zbl", - "models": LinearAtomicModel.serialize( + "models": LinearEnergyAtomicModel.serialize( [self.dp_model, self.zbl_model], self.type_map ), "sw_rmin": self.sw_rmin, @@ -354,14 +420,14 @@ def serialize(self) -> dict: return dd @classmethod - def deserialize(cls, data) -> "DPZBLLinearAtomicModel": + def deserialize(cls, data) -> "DPZBLLinearEnergyAtomicModel": data = copy.deepcopy(data) check_version_compatibility(data.pop("@version", 1), 1, 1) sw_rmin = data.pop("sw_rmin") sw_rmax = data.pop("sw_rmax") smin_alpha = data.pop("smin_alpha") - [dp_model, zbl_model], type_map = LinearAtomicModel.deserialize( + [dp_model, zbl_model], type_map = LinearEnergyAtomicModel.deserialize( data.pop("models") ) diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py index bae4ea55e2..215bb25de5 100644 --- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py +++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import copy from typing import ( + Callable, Dict, List, Optional, @@ -13,9 +14,18 @@ FittingOutputDef, OutputVariableDef, ) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.stat import ( + compute_output_stats, +) from deepmd.utils.pair_tab import ( PairTab, ) +from deepmd.utils.path import ( + DPPath, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -47,9 +57,14 @@ class PairTabAtomicModel(torch.nn.Module, BaseAtomicModel): The cutoff radius. sel : int or list[int] The maxmum number of atoms in the cut-off radius. - type_map : list[str] + type_map : List[str] Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. + rcond : float, optional + The condition number for the regression of atomic energy. + atom_ener + Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set. + """ def __init__( @@ -58,6 +73,8 @@ def __init__( rcut: float, sel: Union[int, List[int]], type_map: List[str], + rcond: Optional[float] = None, + atom_ener: Optional[List[float]] = None, **kwargs, ): torch.nn.Module.__init__(self) @@ -65,8 +82,12 @@ def __init__( self.tab_file = tab_file self.rcut = rcut self.tab = self._set_pairtab(tab_file, rcut) - self.type_map = type_map + BaseAtomicModel.__init__(self, **kwargs) + self.rcond = rcond + self.atom_ener = atom_ener + self.type_map = type_map + self.ntypes = len(type_map) # handle deserialization with no input file if self.tab_file is not None: @@ -74,11 +95,22 @@ def __init__( tab_info, tab_data, ) = self.tab.get() # this returns -> Tuple[np.array, np.array] + nspline, ntypes_tab = tab_info[-2:].astype(int) self.register_buffer("tab_info", torch.from_numpy(tab_info)) - self.register_buffer("tab_data", torch.from_numpy(tab_data)) + self.register_buffer( + "tab_data", + torch.from_numpy(tab_data).reshape(ntypes_tab, ntypes_tab, nspline, 4), + ) + if self.ntypes != ntypes_tab: + raise ValueError( + "The `type_map` provided does not match the number of columns in the table." + ) else: self.register_buffer("tab_info", None) self.register_buffer("tab_data", None) + self.bias_atom_e = torch.zeros( + self.ntypes, 1, dtype=env.GLOBAL_PT_ENER_FLOAT_PRECISION, device=env.DEVICE + ) # self.model_type = "ener" # self.model_version = MODEL_VERSION ## this shoud be in the parent class @@ -145,6 +177,8 @@ def serialize(self) -> dict: "rcut": self.rcut, "sel": self.sel, "type_map": self.type_map, + "rcond": self.rcond, + "atom_ener": self.atom_ener, } ) return dd @@ -156,15 +190,56 @@ def deserialize(cls, data) -> "PairTabAtomicModel": rcut = data.pop("rcut") sel = data.pop("sel") type_map = data.pop("type_map") + rcond = data.pop("rcond") + atom_ener = data.pop("atom_ener") tab = PairTab.deserialize(data.pop("tab")) data.pop("@class", None) data.pop("type", None) - tab_model = cls(None, rcut, sel, type_map, **data) + tab_model = cls(None, rcut, sel, type_map, rcond, atom_ener, **data) + tab_model.tab = tab tab_model.register_buffer("tab_info", torch.from_numpy(tab_model.tab.tab_info)) - tab_model.register_buffer("tab_data", torch.from_numpy(tab_model.tab.tab_data)) + nspline, ntypes = tab_model.tab.tab_info[-2:].astype(int) + tab_model.register_buffer( + "tab_data", + torch.from_numpy(tab_model.tab.tab_data).reshape( + ntypes, ntypes, nspline, 4 + ), + ) return tab_model + def compute_or_load_stat( + self, + merged: Union[Callable[[], List[dict]], List[dict]], + stat_file_path: Optional[DPPath] = None, + ): + """ + Compute the output statistics (e.g. energy bias) for the fitting net from packed data. + + Parameters + ---------- + merged : Union[Callable[[], List[dict]], List[dict]] + - List[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], List[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + stat_file_path : Optional[DPPath] + The path to the stat file. + + """ + bias_atom_e = compute_output_stats( + merged, stat_file_path, self.rcond, self.atom_ener + ) + self.bias_atom_e.copy_( + torch.tensor(bias_atom_e, device=env.DEVICE).view([self.ntypes, 1]) + ) + + def change_energy_bias(self) -> None: + # need to implement + pass + def forward_atomic( self, extended_coord: torch.Tensor, diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py index dcf1c36e83..cacf59c16c 100644 --- a/deepmd/pt/model/model/dp_zbl_model.py +++ b/deepmd/pt/model/model/dp_zbl_model.py @@ -10,7 +10,7 @@ DPModel, ) from deepmd.pt.model.atomic_model import ( - DPZBLLinearAtomicModel, + DPZBLLinearEnergyAtomicModel, ) from deepmd.pt.model.model.model import ( BaseModel, @@ -20,7 +20,7 @@ make_model, ) -DPZBLModel_ = make_model(DPZBLLinearAtomicModel) +DPZBLModel_ = make_model(DPZBLLinearEnergyAtomicModel) @BaseModel.register("zbl") diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 8bf9cc1c90..a11f6410a4 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -30,11 +30,8 @@ from deepmd.pt.utils.env import ( DEFAULT_PRECISION, ) -from deepmd.pt.utils.utils import ( - to_numpy_array, -) -from deepmd.utils.out_stat import ( - compute_stats_from_redu, +from deepmd.pt.utils.stat import ( + compute_output_stats, ) from deepmd.utils.path import ( DPPath, @@ -84,8 +81,8 @@ class InvarFitting(GeneralFitting): Random seed. exclude_types: List[int] Atomic contributions of the excluded atom types are set zero. - atom_ener - Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set. + atom_ener: List[float], optional + Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set. """ @@ -164,41 +161,9 @@ def compute_output_stats( The path to the stat file. """ - if stat_file_path is not None: - stat_file_path = stat_file_path / "bias_atom_e" - if stat_file_path is not None and stat_file_path.is_file(): - bias_atom_e = stat_file_path.load_numpy() - else: - if callable(merged): - # only get data for once - sampled = merged() - else: - sampled = merged - energy = [item["energy"] for item in sampled] - data_mixed_type = "real_natoms_vec" in sampled[0] - if data_mixed_type: - input_natoms = [item["real_natoms_vec"] for item in sampled] - else: - input_natoms = [item["natoms"] for item in sampled] - # shape: (nframes, ndim) - merged_energy = to_numpy_array(torch.cat(energy)) - # shape: (nframes, ntypes) - merged_natoms = to_numpy_array(torch.cat(input_natoms)[:, 2:]) - if self.atom_ener is not None and len(self.atom_ener) > 0: - assigned_atom_ener = np.array( - [ee if ee is not None else np.nan for ee in self.atom_ener] - ) - else: - assigned_atom_ener = None - bias_atom_e, _ = compute_stats_from_redu( - merged_energy, - merged_natoms, - assigned_bias=assigned_atom_ener, - rcond=self.rcond, - ) - if stat_file_path is not None: - stat_file_path.save_numpy(bias_atom_e) - assert all(x is not None for x in [bias_atom_e]) + bias_atom_e = compute_output_stats( + merged, stat_file_path, self.rcond, self.atom_ener + ) self.bias_atom_e.copy_( torch.tensor(bias_atom_e, device=env.DEVICE).view( [self.ntypes, self.dim_out] diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index e5a7632ac4..2a80956b9d 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -29,7 +29,9 @@ TensorLoss, ) from deepmd.pt.model.model import ( + DPZBLModel, get_model, + get_zbl_model, ) from deepmd.pt.optimizer import ( KFOptimizerWrapper, @@ -247,7 +249,10 @@ def get_sample(): def get_single_model( _model_params, ): - model = get_model(deepcopy(_model_params)).to(DEVICE) + if "use_srtab" in _model_params: + model = get_zbl_model(deepcopy(_model_params)).to(DEVICE) + else: + model = get_model(deepcopy(_model_params)).to(DEVICE) return model def get_lr(lr_params): @@ -506,14 +511,20 @@ def get_loss(loss_params, start_lr, _ntypes, _model): model_params["type_map"], model_params["new_type_map"], ) - self.model.fitting_net.change_energy_bias( - config, - self.model, - old_type_map, - new_type_map, - ntest=ntest, - bias_shift=model_params.get("bias_shift", "delta"), - ) + if hasattr(self.model, "fitting_net"): + self.model.fitting_net.change_energy_bias( + config, + self.model, + old_type_map, + new_type_map, + ntest=ntest, + bias_shift=model_params.get("bias_shift", "delta"), + ) + elif isinstance(self.model, DPZBLModel): + # need to updated + self.model.change_energy_bias() + else: + raise NotImplementedError if init_frz_model is not None: frz_model = torch.jit.load(init_frz_model, map_location=DEVICE) self.model.load_state_dict(frz_model.state_dict()) diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 3b246a0ec2..63abccc75d 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -1,10 +1,27 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging +from typing import ( + Callable, + List, + Optional, + Union, +) +import numpy as np import torch +from deepmd.pt.utils import ( + env, +) from deepmd.pt.utils.utils import ( dict_to_device, + to_numpy_array, +) +from deepmd.utils.out_stat import ( + compute_stats_from_redu, +) +from deepmd.utils.path import ( + DPPath, ) log = logging.getLogger(__name__) @@ -50,3 +67,67 @@ def make_stat_input(datasets, dataloaders, nbatches): dict_to_device(sys_stat) lst.append(sys_stat) return lst + + +def compute_output_stats( + merged: Union[Callable[[], List[dict]], List[dict]], + stat_file_path: Optional[DPPath] = None, + rcond: Optional[float] = None, + atom_ener: Optional[List[float]] = None, +): + """ + Compute the output statistics (e.g. energy bias) for the fitting net from packed data. + + Parameters + ---------- + merged : Union[Callable[[], List[dict]], List[dict]] + - List[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], List[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + stat_file_path : DPPath, optional + The path to the stat file. + rcond : float, optional + The condition number for the regression of atomic energy. + atom_ener : List[float], optional + Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set. + + """ + if stat_file_path is not None: + stat_file_path = stat_file_path / "bias_atom_e" + if stat_file_path is not None and stat_file_path.is_file(): + bias_atom_e = stat_file_path.load_numpy() + else: + if callable(merged): + # only get data for once + sampled = merged() + else: + sampled = merged + energy = [item["energy"] for item in sampled] + data_mixed_type = "real_natoms_vec" in sampled[0] + if data_mixed_type: + input_natoms = [item["real_natoms_vec"] for item in sampled] + else: + input_natoms = [item["natoms"] for item in sampled] + # shape: (nframes, ndim) + merged_energy = to_numpy_array(torch.cat(energy)) + # shape: (nframes, ntypes) + merged_natoms = to_numpy_array(torch.cat(input_natoms)[:, 2:]) + if atom_ener is not None and len(atom_ener) > 0: + assigned_atom_ener = np.array( + [ee if ee is not None else np.nan for ee in atom_ener] + ) + else: + assigned_atom_ener = None + bias_atom_e, _ = compute_stats_from_redu( + merged_energy, + merged_natoms, + assigned_bias=assigned_atom_ener, + rcond=rcond, + ) + if stat_file_path is not None: + stat_file_path.save_numpy(bias_atom_e) + assert all(x is not None for x in [bias_atom_e]) + return torch.tensor(bias_atom_e, device=env.DEVICE) diff --git a/source/tests/common/dpmodel/test_linear_atomic_model.py b/source/tests/common/dpmodel/test_linear_atomic_model.py index cc08a3b3dd..832d1de106 100644 --- a/source/tests/common/dpmodel/test_linear_atomic_model.py +++ b/source/tests/common/dpmodel/test_linear_atomic_model.py @@ -10,7 +10,7 @@ DPAtomicModel, ) from deepmd.dpmodel.atomic_model.linear_atomic_model import ( - DPZBLLinearAtomicModel, + DPZBLLinearEnergyAtomicModel, ) from deepmd.dpmodel.atomic_model.pairtab_atomic_model import ( PairTabAtomicModel, @@ -58,7 +58,7 @@ def test_pairwise(self, mock_loadtxt): ) dp_model = DPAtomicModel(ds, ft, type_map=type_map) - wgt_model = DPZBLLinearAtomicModel( + wgt_model = DPZBLLinearEnergyAtomicModel( dp_model, zbl_model, sw_rmin=0.1, @@ -151,14 +151,14 @@ def setUp(self, mock_loadtxt): zbl_model = PairTabAtomicModel( file_path, self.rcut, sum(self.sel), type_map=type_map ) - self.md0 = DPZBLLinearAtomicModel( + self.md0 = DPZBLLinearEnergyAtomicModel( dp_model, zbl_model, sw_rmin=0.1, sw_rmax=0.25, type_map=type_map, ) - self.md1 = DPZBLLinearAtomicModel.deserialize(self.md0.serialize()) + self.md1 = DPZBLLinearEnergyAtomicModel.deserialize(self.md0.serialize()) def test_self_consistency(self): ret0 = self.md0.forward_atomic(self.coord_ext, self.atype_ext, self.nlist) diff --git a/source/tests/pt/model/test_linear_atomic_model.py b/source/tests/pt/model/test_linear_atomic_model.py index e0904097e3..adc682a41f 100644 --- a/source/tests/pt/model/test_linear_atomic_model.py +++ b/source/tests/pt/model/test_linear_atomic_model.py @@ -8,11 +8,11 @@ import torch from deepmd.dpmodel.atomic_model import ( - DPZBLLinearAtomicModel as DPDPZBLLinearAtomicModel, + DPZBLLinearEnergyAtomicModel as DPDPZBLLinearEnergyAtomicModel, ) from deepmd.pt.model.atomic_model import ( DPAtomicModel, - DPZBLLinearAtomicModel, + DPZBLLinearEnergyAtomicModel, PairTabAtomicModel, ) from deepmd.pt.model.descriptor.se_a import ( @@ -70,10 +70,10 @@ def test_pairwise(self, mock_loadtxt): type_map = ["foo", "bar"] zbl_model = PairTabAtomicModel( - tab_file=file_path, rcut=0.3, sel=2, type_map=type_map + tab_file=file_path, rcut=0.3, sel=2, type_map=type_map[::-1] ) dp_model = DPAtomicModel(ds, ft, type_map=type_map).to(env.DEVICE) - wgt_model = DPZBLLinearAtomicModel( + wgt_model = DPZBLLinearEnergyAtomicModel( dp_model, zbl_model, sw_rmin=0.1, @@ -145,17 +145,17 @@ def setUp(self, mock_loadtxt): zbl_model = PairTabAtomicModel( file_path, self.rcut, sum(self.sel), type_map=type_map ) - self.md0 = DPZBLLinearAtomicModel( + self.md0 = DPZBLLinearEnergyAtomicModel( dp_model, zbl_model, sw_rmin=0.1, sw_rmax=0.25, type_map=type_map, ).to(env.DEVICE) - self.md1 = DPZBLLinearAtomicModel.deserialize(self.md0.serialize()).to( + self.md1 = DPZBLLinearEnergyAtomicModel.deserialize(self.md0.serialize()).to( env.DEVICE ) - self.md2 = DPDPZBLLinearAtomicModel.deserialize(self.md0.serialize()) + self.md2 = DPDPZBLLinearEnergyAtomicModel.deserialize(self.md0.serialize()) self.md3 = DPZBLModel( dp_model, zbl_model, sw_rmin=0.1, sw_rmax=0.25, type_map=type_map ) @@ -185,5 +185,23 @@ def test_jit(self): self.assertEqual(md3.get_type_map(), ["foo", "bar"]) +class TestRemmapMethod(unittest.TestCase): + def test_valid(self): + atype = torch.randint(0, 3, (4, 20), device=env.DEVICE) + commonl = ["H", "O", "S"] + originl = ["Si", "H", "O", "S"] + mapping = DPZBLLinearEnergyAtomicModel.remap_atype(originl, commonl) + new_atype = mapping[atype] + + def trans(atype, map): + idx = atype.flatten().tolist() + res = [] + for i in idx: + res.append(map[i]) + return res + + assert trans(atype, commonl) == trans(new_atype, originl) + + if __name__ == "__main__": unittest.main(warnings="ignore") diff --git a/source/tests/pt/model/test_pairtab_atomic_model.py b/source/tests/pt/model/test_pairtab_atomic_model.py index 0576f89910..322de51a2c 100644 --- a/source/tests/pt/model/test_pairtab_atomic_model.py +++ b/source/tests/pt/model/test_pairtab_atomic_model.py @@ -229,7 +229,7 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: ) model = PairTabAtomicModel( - tab_file=file_path, rcut=rcut, sel=2, type_map=["S"] + tab_file=file_path, rcut=rcut, sel=2, type_map=["H"] ) results.append( model.forward_atomic(extended_coord, extended_atype, nlist)["energy"] diff --git a/source/tests/pt/model/water/zbl.json b/source/tests/pt/model/water/zbl.json new file mode 100644 index 0000000000..cb5602d92d --- /dev/null +++ b/source/tests/pt/model/water/zbl.json @@ -0,0 +1,92 @@ +{ + "_comment1": " model parameters", + "model": { + "use_srtab": "H2O_tab_potential.txt", + "smin_alpha": 0.1, + "sw_rmin": 0.8, + "sw_rmax": 1.0, + "type_map": [ + "O", + "H" + ], + "descriptor": { + "type": "se_e2_a", + "sel": [ + 46, + 92 + ], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [ + 25, + 50, + 100 + ], + "resnet_dt": false, + "axis_neuron": 16, + "type_one_side": true, + "precision": "float64", + "seed": 1, + "_comment2": " that's all" + }, + "fitting_net": { + "neuron": [ + 240, + 240, + 240 + ], + "resnet_dt": true, + "precision": "float64", + "seed": 1, + "_comment3": " that's all" + }, + "_comment4": " that's all" + }, + + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 3.51e-8, + "_comment5": "that's all" + }, + + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + "_comment6": " that's all" + }, + + "training": { + "training_data": { + "systems": [ + "../data/data_0/", + "../data/data_1/", + "../data/data_2/" + ], + "batch_size": "auto", + "_comment7": "that's all" + }, + "validation_data": { + "systems": [ + "../data/data_3" + ], + "batch_size": 1, + "numb_btch": 3, + "_comment8": "that's all" + }, + "numb_steps": 1000000, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 1000, + "_comment9": "that's all" + }, + + "_comment10": "that's all" +} diff --git a/source/tests/pt/test_finetune.py b/source/tests/pt/test_finetune.py index 226fed3c65..d21a44acc7 100644 --- a/source/tests/pt/test_finetune.py +++ b/source/tests/pt/test_finetune.py @@ -17,7 +17,10 @@ DeepEval, ) from deepmd.pt.model.model import ( + DPZBLModel, + EnergyModel, get_model, + get_zbl_model, ) from deepmd.utils.data_system import ( DeepmdDataSystem, @@ -27,23 +30,48 @@ ) from .model.test_permutation import ( - model_dpa1, model_dpa2, model_se_e2_a, + model_zbl, ) class FinetuneTest: def test_finetune_change_energy_bias(self): # get model - model = get_model(self.model_config) - model.fitting_net.bias_atom_e = torch.rand_like(model.fitting_net.bias_atom_e) - energy_bias_before = deepcopy( - model.fitting_net.bias_atom_e.detach().cpu().numpy().reshape(-1) - ) - bias_atom_e_input = deepcopy( - model.fitting_net.bias_atom_e.detach().cpu().numpy().reshape(-1) - ) + if "use_srtab" in self.model_config: + model = get_zbl_model(self.model_config) + else: + model = get_model(self.model_config) + if isinstance(model, EnergyModel): + model.fitting_net.bias_atom_e = torch.rand_like( + model.fitting_net.bias_atom_e + ) + energy_bias_before = deepcopy( + model.fitting_net.bias_atom_e.detach().cpu().numpy().reshape(-1) + ) + bias_atom_e_input = deepcopy( + model.fitting_net.bias_atom_e.detach().cpu().numpy().reshape(-1) + ) + elif isinstance(model, DPZBLModel): + model.dp_model.fitting_net.bias_atom_e = torch.rand_like( + model.dp_model.fitting_net.bias_atom_e + ) + energy_bias_before = deepcopy( + model.dp_model.fitting_net.bias_atom_e.detach() + .cpu() + .numpy() + .reshape(-1) + ) + bias_atom_e_input = deepcopy( + model.dp_model.fitting_net.bias_atom_e.detach() + .cpu() + .numpy() + .reshape(-1) + ) + else: + bias_atom_e_input = None + model = torch.jit.script(model) tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth") torch.jit.save(model, tmp_model.name) @@ -109,7 +137,8 @@ def tearDown(self) -> None: FinetuneTest.tearDown(self) -class TestEnergyModelDPA1(unittest.TestCase, FinetuneTest): +@unittest.skip("change bias not implemented yet.") +class TestEnergyZBLModelSeA(unittest.TestCase, FinetuneTest): def setUp(self): self.data_file = [str(Path(__file__).parent / "water/data/data_0")] self.data = DeepmdDataSystem( @@ -118,7 +147,7 @@ def setUp(self): test_size=1, ) self.data.add("energy", ndof=1, atomic=False, must=True, high_prec=True) - self.model_config = model_dpa1 + self.model_config = model_zbl def tearDown(self) -> None: FinetuneTest.tearDown(self) diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py index bcb8c6c188..db69a1bcea 100644 --- a/source/tests/pt/test_training.py +++ b/source/tests/pt/test_training.py @@ -21,6 +21,7 @@ model_dpa2, model_hybrid, model_se_e2_a, + model_zbl, ) @@ -66,6 +67,7 @@ def test_trainable(self): torch.testing.assert_close( model_dict_before_training[key], model_dict_after_training[key] ) + self.tearDown() def tearDown(self): @@ -94,6 +96,22 @@ def tearDown(self) -> None: DPTrainTest.tearDown(self) +class TestEnergyZBLModelSeA(unittest.TestCase, DPTrainTest): + def setUp(self): + input_json = str(Path(__file__).parent / "water/zbl.json") + with open(input_json) as f: + self.config = json.load(f) + data_file = [str(Path(__file__).parent / "water/data/data_0")] + self.config["training"]["training_data"]["systems"] = data_file + self.config["training"]["validation_data"]["systems"] = data_file + self.config["model"] = deepcopy(model_zbl) + self.config["training"]["numb_steps"] = 1 + self.config["training"]["save_freq"] = 1 + + def tearDown(self) -> None: + DPTrainTest.tearDown(self) + + class TestFparam(unittest.TestCase, DPTrainTest): """Test if `fparam` can be loaded correctly."""