diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 3c9c595221..ecfd08b61a 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import functools import math from collections.abc import ( Callable, @@ -52,6 +53,7 @@ def __init__( pair_exclude_types: list[tuple[int, int]] = [], rcond: float | None = None, preset_out_bias: dict[str, Array] | None = None, + data_stat_protect: float = 1e-2, ) -> None: super().__init__() self.type_map = type_map @@ -59,6 +61,7 @@ def __init__( self.reinit_pair_exclude(pair_exclude_types) self.rcond = rcond self.preset_out_bias = preset_out_bias + self.data_stat_protect = data_stat_protect def init_out_stat(self) -> None: """Initialize the output bias.""" @@ -77,6 +80,14 @@ def init_out_stat(self) -> None: self.out_bias = out_bias_data self.out_std = out_std_data + def get_out_bias(self) -> Array: + """Get the output bias.""" + return self.out_bias + + def set_out_bias(self, out_bias: Array) -> None: + """Set the output bias.""" + self.out_bias = out_bias + def __setitem__(self, key: str, value: Array) -> None: if key in ["out_bias"]: self.out_bias = value @@ -287,6 +298,57 @@ def compute_or_load_out_stat( bias_adjust_mode="set-by-statistic", ) + def _make_wrapped_sampler( + self, + sampled_func: Callable[[], list[dict]], + ) -> Callable[[], list[dict]]: + """Wrap the sampled function with exclusion types and default fparam. + + The returned callable is cached so that the sampling (which may be + expensive) is performed at most once. + + Parameters + ---------- + sampled_func + The lazy sampled function to get data frames from different data + systems. + + Returns + ------- + Callable[[], list[dict]] + A cached wrapper around *sampled_func* that additionally sets + ``pair_exclude_types``, ``atom_exclude_types`` and default + ``fparam`` on every sample dict when applicable. + """ + + @functools.lru_cache + def wrapped_sampler() -> list[dict]: + sampled = sampled_func() + if self.pair_excl is not None: + pair_exclude_types = self.pair_excl.get_exclude_types() + for sample in sampled: + sample["pair_exclude_types"] = list(pair_exclude_types) + if self.atom_excl is not None: + atom_exclude_types = self.atom_excl.get_exclude_types() + for sample in sampled: + sample["atom_exclude_types"] = list(atom_exclude_types) + if ( + "find_fparam" not in sampled[0] + and "fparam" not in sampled[0] + and self.has_default_fparam() + ): + default_fparam = self.get_default_fparam() + if default_fparam is not None: + default_fparam_np = np.array(default_fparam) + for sample in sampled: + nframe = sample["atype"].shape[0] + sample["fparam"] = np.tile( + default_fparam_np.reshape(1, -1), (nframe, 1) + ) + return sampled + + return wrapped_sampler + def change_out_bias( self, sample_merged: Callable[[], list[dict]] | list[dict], diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index a6c71b0b85..73447de955 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, ) @@ -15,6 +18,9 @@ from deepmd.dpmodel.output_def import ( FittingOutputDef, ) +from deepmd.utils.path import ( + DPPath, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -62,17 +68,16 @@ def __init__( **kwargs: Any, ) -> None: super().__init__(type_map, **kwargs) - self.type_map = type_map self.descriptor = descriptor - self.fitting = fitting - if hasattr(self.fitting, "reinit_exclude"): - self.fitting.reinit_exclude(self.atom_exclude_types) + self.fitting_net = fitting + if hasattr(self.fitting_net, "reinit_exclude"): + self.fitting_net.reinit_exclude(self.atom_exclude_types) self.type_map = type_map super().init_out_stat() def fitting_output_def(self) -> FittingOutputDef: """Get the output def of the fitting net.""" - return self.fitting.output_def() + return self.fitting_net.output_def() def get_rcut(self) -> float: """Get the cut-off radius.""" @@ -87,7 +92,7 @@ def set_case_embd(self, case_idx: int) -> None: Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. """ - self.fitting.set_case_embd(case_idx) + self.fitting_net.set_case_embd(case_idx) def mixed_types(self) -> bool: """If true, the model @@ -180,7 +185,7 @@ def forward_atomic( nlist, mapping=mapping, ) - ret = self.fitting( + ret = self.fitting_net( descriptor, atype, gr=rot_mat, @@ -191,6 +196,37 @@ def forward_atomic( ) return ret + def compute_or_load_stat( + self, + sampled_func: Callable[[], list[dict]], + stat_file_path: DPPath | None = None, + compute_or_load_out_stat: bool = True, + ) -> 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. + + Parameters + ---------- + sampled_func + The lazy sampled function to get data frames from different data systems. + stat_file_path + The path to the stat file. + compute_or_load_out_stat : bool + Whether to compute the output statistics. + If False, it will only compute the input statistics + (e.g. mean and standard deviation of descriptors). + """ + if stat_file_path is not None and self.type_map is not None: + stat_file_path /= " ".join(self.type_map) + + wrapped_sampler = self._make_wrapped_sampler(sampled_func) + self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path) + self.fitting_net.compute_input_stats( + wrapped_sampler, stat_file_path=stat_file_path + ) + if compute_or_load_out_stat: + self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) + def change_type_map( self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: @@ -207,7 +243,31 @@ def change_type_map( if model_with_new_type_stat is not None else None, ) - self.fitting.change_type_map(type_map=type_map) + self.fitting_net.change_type_map(type_map=type_map) + + def compute_fitting_input_stat( + self, + sample_merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, + ) -> None: + """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. + + Parameters + ---------- + sample_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``: ``np.ndarray`` originating from the ``i``-th data system. + - Callable[[], list[dict]]: A lazy function that returns data samples + in the above format only when needed. + stat_file_path : Optional[DPPath] + The path to the stat file. + """ + self.fitting_net.compute_input_stats( + sample_merged, + protection=self.data_stat_protect, + stat_file_path=stat_file_path, + ) def serialize(self) -> dict: dd = super().serialize() @@ -218,7 +278,7 @@ def serialize(self) -> dict: "@version": 2, "type_map": self.type_map, "descriptor": self.descriptor.serialize(), - "fitting": self.fitting.serialize(), + "fitting": self.fitting_net.serialize(), } ) return dd @@ -244,19 +304,19 @@ def deserialize(cls, data: dict[str, Any]) -> "DPAtomicModel": def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" - return self.fitting.get_dim_fparam() + return self.fitting_net.get_dim_fparam() def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" - return self.fitting.get_dim_aparam() + return self.fitting_net.get_dim_aparam() def has_default_fparam(self) -> bool: """Check if the model has default frame parameters.""" - return self.fitting.has_default_fparam() + return self.fitting_net.has_default_fparam() def get_default_fparam(self) -> list[float] | None: """Get the default frame parameters.""" - return self.fitting.get_default_fparam() + return self.fitting_net.get_default_fparam() def get_sel_type(self) -> list[int]: """Get the selected atom types of this model. @@ -265,7 +325,7 @@ def get_sel_type(self) -> list[int]: to the result of the model. If returning an empty list, all atom types are selected. """ - return self.fitting.get_sel_type() + return self.fitting_net.get_sel_type() def is_aparam_nall(self) -> bool: """Check whether the shape of atomic parameters is (nframes, nall, ndim). diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index be483f97c4..06ed524ef6 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, ) @@ -17,6 +20,9 @@ from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) +from deepmd.utils.path import ( + DPPath, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -338,6 +344,38 @@ def deserialize(cls, data: dict) -> "LinearEnergyAtomicModel": data["models"] = models return super().deserialize(data) + def compute_or_load_stat( + self, + sampled_func: Callable[[], list[dict]], + stat_file_path: DPPath | None = None, + compute_or_load_out_stat: bool = True, + ) -> None: + """Compute or load the statistics parameters of the model. + + For LinearEnergyAtomicModel, this first computes input stats for each + sub-model (without output stats), then computes its own output stats. + + Parameters + ---------- + sampled_func + The lazy sampled function to get data frames from different data systems. + stat_file_path + The path to the stat file. + compute_or_load_out_stat : bool + Whether to compute the output statistics. + """ + for md in self.models: + md.compute_or_load_stat( + sampled_func, stat_file_path, compute_or_load_out_stat=False + ) + + if stat_file_path is not None and self.type_map is not None: + stat_file_path /= " ".join(self.type_map) + + if compute_or_load_out_stat: + wrapped_sampler = self._make_wrapped_sampler(sampled_func) + self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) + def _compute_weight( self, extended_coord: Array, @@ -523,4 +561,4 @@ def _compute_weight( # to handle masked atoms coef = xp.where(sigma != 0, coef, xp.zeros_like(coef)) self.zbl_weight = coef - return [1 - xp.expand_dims(coef, -1), xp.expand_dims(coef, -1)] + return [1 - xp.expand_dims(coef, axis=-1), xp.expand_dims(coef, axis=-1)] diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index 0e5049bd12..834ee19016 100644 --- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, NoReturn, @@ -21,6 +24,9 @@ from deepmd.utils.pair_tab import ( PairTab, ) +from deepmd.utils.path import ( + DPPath, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -75,20 +81,16 @@ def __init__( sel: int | list[int], type_map: list[str], rcond: float | None = None, - atom_ener: list[float] | None = None, **kwargs: Any, ) -> None: super().__init__(type_map, **kwargs) super().init_out_stat() self.tab_file = tab_file self.rcut = rcut - self.type_map = type_map self.tab = PairTab(self.tab_file, rcut=rcut) - self.type_map = type_map self.ntypes = len(type_map) self.rcond = rcond - self.atom_ener = atom_ener if self.tab_file is not None: tab_info, tab_data = self.tab.get() @@ -202,11 +204,37 @@ def deserialize(cls, data: dict) -> "PairTabAtomicModel": tab_model = super().deserialize(data) tab_model.tab = tab - tab_model.tab_info = tab_model.tab.tab_info - nspline, ntypes = tab_model.tab_info[-2:].astype(int) - tab_model.tab_data = tab_model.tab.tab_data.reshape(ntypes, ntypes, nspline, 4) + # Extract nspline/ntypes from the numpy source before setting on the + # model, because dpmodel_setattr may convert to torch tensor. + nspline, ntypes = tab.tab_info[-2:].astype(int) + tab_model.tab_info = tab.tab_info + tab_model.tab_data = tab.tab_data.reshape(ntypes, ntypes, nspline, 4) return tab_model + def compute_or_load_stat( + self, + sampled_func: Callable[[], list[dict]], + stat_file_path: DPPath | None = None, + compute_or_load_out_stat: bool = True, + ) -> None: + """Compute or load the statistics parameters of the model. + + PairTabAtomicModel has no descriptor or fitting net input stats, + so this only computes output stats (energy bias) when requested. + + Parameters + ---------- + sampled_func + The lazy sampled function to get data frames from different data systems. + stat_file_path + The path to the stat file. + compute_or_load_out_stat : bool + Whether to compute the output statistics. + """ + if compute_or_load_out_stat: + wrapped_sampler = self._make_wrapped_sampler(sampled_func) + self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) + def forward_atomic( self, extended_coord: Array, @@ -291,7 +319,7 @@ def _pair_tabulated_inter( hi = 1.0 / hh # jax jit does not support convert to a Python int, so we need to convert to xp.int64. - nspline = (self.tab_info[2] + 0.1).astype(xp.int64) + nspline = xp.astype(self.tab_info[2] + 0.1, xp.int64) uu = (rr - rmin) * hi # this is broadcasted to (nframes,nloc,nnei) @@ -348,7 +376,7 @@ def _get_pairwise_dist(coords: Array, nlist: Array) -> Array: neighbor_atoms = coords[batch_indices, nlist] loc_atoms = coords[:, : nlist.shape[1], :] pairwise_dr = loc_atoms[:, :, None, :] - neighbor_atoms - pairwise_rr = safe_for_sqrt(xp.sum(xp.power(pairwise_dr, 2), axis=-1)) + pairwise_rr = safe_for_sqrt(xp.sum(pairwise_dr**2, axis=-1)) return pairwise_rr @@ -394,16 +422,18 @@ def _extract_spline_coefficient( expanded_idx = xp.broadcast_to( idx[..., xp.newaxis, xp.newaxis], (*idx.shape, 1, 4) ) - clipped_indices = xp.clip(expanded_idx, 0, nspline - 1).astype(int) + clipped_indices = xp.astype(xp.clip(expanded_idx, 0, nspline - 1), xp.int64) # (nframes, nloc, nnei, 4) final_coef = xp.squeeze( - xp_take_along_axis(expanded_tab_data, clipped_indices, 3) + xp_take_along_axis(expanded_tab_data, clipped_indices, 3), axis=3 ) # when the spline idx is beyond the table, all spline coefficients are set to `0`, and the resulting ener corresponding to the idx is also `0`. final_coef = xp.where( - expanded_idx.squeeze() > nspline, xp.zeros_like(final_coef), final_coef + xp.squeeze(expanded_idx, axis=3) > nspline, + xp.zeros_like(final_coef), + final_coef, ) return final_coef diff --git a/deepmd/dpmodel/atomic_model/polar_atomic_model.py b/deepmd/dpmodel/atomic_model/polar_atomic_model.py index bdad31dcb1..76a221de46 100644 --- a/deepmd/dpmodel/atomic_model/polar_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/polar_atomic_model.py @@ -45,21 +45,21 @@ def apply_out_stat( xp = array_api_compat.array_namespace(atype) out_bias, out_std = self._fetch_out_stat(self.bias_keys) - if self.fitting.shift_diag: + if self.fitting_net.shift_diag: nframes, nloc = atype.shape dtype = out_bias[self.bias_keys[0]].dtype device = array_api_compat.device(out_bias[self.bias_keys[0]]) for kk in self.bias_keys: ntypes = out_bias[kk].shape[0] temp = xp.mean( - xp.diagonal(out_bias[kk].reshape(ntypes, 3, 3), axis1=1, axis2=2), + xp.diagonal(out_bias[kk].reshape(ntypes, 3, 3), 0, 1, 2), axis=1, ) modified_bias = temp[atype] # (nframes, nloc, 1) modified_bias = ( - modified_bias[..., xp.newaxis] * (self.fitting.scale[atype]) + modified_bias[..., xp.newaxis] * (self.fitting_net.scale[atype]) ) eye = xp.eye(3, dtype=dtype, device=device) diff --git a/deepmd/dpmodel/atomic_model/property_atomic_model.py b/deepmd/dpmodel/atomic_model/property_atomic_model.py index ec65f949e0..07dd00b109 100644 --- a/deepmd/dpmodel/atomic_model/property_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/property_atomic_model.py @@ -25,6 +25,14 @@ def __init__( ) super().__init__(descriptor, fitting, type_map, **kwargs) + def get_compute_stats_distinguish_types(self) -> bool: + """Get whether the fitting net computes stats which are not distinguished between different types of atoms.""" + return False + + def get_intensive(self) -> bool: + """Whether the fitting property is intensive.""" + return self.fitting_net.intensive + def apply_out_stat( self, ret: dict[str, Array], diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 8b0a8b9072..7aa2044075 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -707,6 +707,7 @@ def __init__( sel = [sel] self.sel = sel self.nnei = sum(sel) + self.ndescrpt = self.nnei * 4 self.ntypes = ntypes self.neuron = neuron self.filter_neuron = self.neuron diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py index d69732911d..4f45401511 100644 --- a/deepmd/dpmodel/descriptor/repflows.py +++ b/deepmd/dpmodel/descriptor/repflows.py @@ -438,6 +438,8 @@ def compute_input_stats( The path to the stat file. """ + if self.set_stddev_constant and self.set_davg_zero: + return env_mat_stat = EnvMatStatSe(self) if path is not None: path = path / env_mat_stat.get_hash() @@ -458,9 +460,10 @@ def compute_input_stats( self.mean = xp.asarray( mean, dtype=self.mean.dtype, copy=True, device=device ) - self.stddev = xp.asarray( - stddev, dtype=self.stddev.dtype, copy=True, device=device - ) + if not self.set_stddev_constant: + self.stddev = xp.asarray( + stddev, dtype=self.stddev.dtype, copy=True, device=device + ) def get_stats(self) -> dict[str, StatItem]: """Get the statistics of the descriptor.""" diff --git a/deepmd/dpmodel/fitting/dos_fitting.py b/deepmd/dpmodel/fitting/dos_fitting.py index 803f31b30f..c8f145ce15 100644 --- a/deepmd/dpmodel/fitting/dos_fitting.py +++ b/deepmd/dpmodel/fitting/dos_fitting.py @@ -15,6 +15,10 @@ from deepmd.dpmodel.fitting.invar_fitting import ( InvarFitting, ) +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + OutputVariableDef, +) if TYPE_CHECKING: from deepmd.dpmodel.fitting.general_fitting import ( @@ -75,6 +79,19 @@ def __init__( default_fparam=default_fparam, ) + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [self.dim_out], + reducible=True, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 260be619fd..c7372f05ac 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -35,10 +35,16 @@ from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) +from deepmd.utils.env_mat_stat import ( + StatItem, +) from deepmd.utils.finetune import ( get_index_between_two_maps, map_atom_exclude_types, ) +from deepmd.utils.path import ( + DPPath, +) from .base_fitting import ( BaseFitting, @@ -226,6 +232,7 @@ def compute_input_stats( self, merged: Callable[[], list[dict]] | list[dict], protection: float = 1e-2, + stat_file_path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. @@ -241,27 +248,48 @@ def compute_input_stats( the lazy function helps by only sampling once. protection : float Divided-by-zero protection + stat_file_path : Optional[DPPath] + The path to the stat file. """ if self.numb_fparam == 0 and self.numb_aparam == 0: # skip data statistics return - if callable(merged): - sampled = merged() - else: - sampled = merged # stat fparam if self.numb_fparam > 0: - cat_data = np.concatenate([frame["fparam"] for frame in sampled], axis=0) - cat_data = np.reshape(cat_data, [-1, self.numb_fparam]) - fparam_avg = np.mean(cat_data, axis=0) - fparam_std = np.std(cat_data, axis=0, ddof=0) # ddof=0 for population std - fparam_std = np.where( - fparam_std < protection, - np.array(protection, dtype=fparam_std.dtype), - fparam_std, + if ( + stat_file_path is not None + and stat_file_path.is_dir() + and (stat_file_path / "fparam").is_file() + ): + fparam_stats = self._load_param_stats_from_file( + stat_file_path, "fparam", self.numb_fparam + ) + else: + sampled = merged() if callable(merged) else merged + cat_data = np.concatenate( + [frame["fparam"] for frame in sampled], axis=0 + ) + cat_data = np.reshape(cat_data, [-1, self.numb_fparam]) + fparam_stats = [ + StatItem( + number=cat_data.shape[0], + sum=np.sum(cat_data[:, ii]), + squared_sum=np.sum(cat_data[:, ii] ** 2), + ) + for ii in range(self.numb_fparam) + ] + if stat_file_path is not None: + self._save_param_stats_to_file( + stat_file_path, "fparam", fparam_stats + ) + fparam_avg = np.array( + [s.compute_avg() for s in fparam_stats], dtype=np.float64 + ) + fparam_std = np.array( + [s.compute_std(protection=protection) for s in fparam_stats], + dtype=np.float64, ) fparam_inv_std = 1.0 / fparam_std - # Use array_api_compat to handle both numpy and torch xp = array_api_compat.array_namespace(self.fparam_avg) self.fparam_avg = xp.asarray( fparam_avg, @@ -275,26 +303,47 @@ def compute_input_stats( ) # stat aparam if self.numb_aparam > 0: - sys_sumv = [] - sys_sumv2 = [] - sys_sumn = [] - for ss_ in [frame["aparam"] for frame in sampled]: - ss = np.reshape(ss_, [-1, self.numb_aparam]) - sys_sumv.append(np.sum(ss, axis=0)) - sys_sumv2.append(np.sum(ss * ss, axis=0)) - sys_sumn.append(ss.shape[0]) - sumv = np.sum(np.stack(sys_sumv), axis=0) - sumv2 = np.sum(np.stack(sys_sumv2), axis=0) - sumn = sum(sys_sumn) - aparam_avg = sumv / sumn - aparam_std = np.sqrt(sumv2 / sumn - (sumv / sumn) ** 2) - aparam_std = np.where( - aparam_std < protection, - np.array(protection, dtype=aparam_std.dtype), - aparam_std, + if ( + stat_file_path is not None + and stat_file_path.is_dir() + and (stat_file_path / "aparam").is_file() + ): + aparam_stats = self._load_param_stats_from_file( + stat_file_path, "aparam", self.numb_aparam + ) + else: + sampled = merged() if callable(merged) else merged + sys_sumv = [] + sys_sumv2 = [] + sys_sumn = [] + for ss_ in [frame["aparam"] for frame in sampled]: + ss = np.reshape(ss_, [-1, self.numb_aparam]) + sys_sumv.append(np.sum(ss, axis=0)) + sys_sumv2.append(np.sum(ss * ss, axis=0)) + sys_sumn.append(ss.shape[0]) + sumv = np.sum(np.stack(sys_sumv), axis=0) + sumv2 = np.sum(np.stack(sys_sumv2), axis=0) + sumn = sum(sys_sumn) + aparam_stats = [ + StatItem( + number=sumn, + sum=sumv[ii], + squared_sum=sumv2[ii], + ) + for ii in range(self.numb_aparam) + ] + if stat_file_path is not None: + self._save_param_stats_to_file( + stat_file_path, "aparam", aparam_stats + ) + aparam_avg = np.array( + [s.compute_avg() for s in aparam_stats], dtype=np.float64 + ) + aparam_std = np.array( + [s.compute_std(protection=protection) for s in aparam_stats], + dtype=np.float64, ) aparam_inv_std = 1.0 / aparam_std - # Use array_api_compat to handle both numpy and torch xp = array_api_compat.array_namespace(self.aparam_avg) self.aparam_avg = xp.asarray( aparam_avg, @@ -307,6 +356,31 @@ def compute_input_stats( device=array_api_compat.device(self.aparam_inv_std), ) + @staticmethod + def _save_param_stats_to_file( + stat_file_path: DPPath, + name: str, + stats: list[StatItem], + ) -> None: + stat_file_path.mkdir(exist_ok=True, parents=True) + fp = stat_file_path / name + arr = np.array([[s.number, s.sum, s.squared_sum] for s in stats]) + fp.save_numpy(arr) + + @staticmethod + def _load_param_stats_from_file( + stat_file_path: DPPath, + name: str, + numb: int, + ) -> list[StatItem]: + fp = stat_file_path / name + arr = fp.load_numpy() + assert arr.shape == (numb, 3) + return [ + StatItem(number=arr[ii][0], sum=arr[ii][1], squared_sum=arr[ii][2]) + for ii in range(numb) + ] + @abstractmethod def _net_out_dim(self) -> int: """Set the FittingNet output dim.""" @@ -363,11 +437,14 @@ def change_type_map( self.ntypes = len(type_map) self.reinit_exclude(map_atom_exclude_types(self.exclude_types, remap_index)) if has_new_type: + xp = array_api_compat.array_namespace(self.bias_atom_e) extend_shape = [len(type_map), *list(self.bias_atom_e.shape[1:])] - extend_bias_atom_e = np.zeros(extend_shape, dtype=self.bias_atom_e.dtype) - self.bias_atom_e = np.concatenate( - [self.bias_atom_e, extend_bias_atom_e], axis=0 + extend_bias_atom_e = xp.zeros( + extend_shape, + dtype=self.bias_atom_e.dtype, + device=array_api_compat.device(self.bias_atom_e), ) + self.bias_atom_e = xp.concat([self.bias_atom_e, extend_bias_atom_e], axis=0) self.bias_atom_e = self.bias_atom_e[remap_index] def __setitem__(self, key: str, value: Any) -> None: diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py index b7868a9502..361f033a68 100644 --- a/deepmd/dpmodel/fitting/polarizability_fitting.py +++ b/deepmd/dpmodel/fitting/polarizability_fitting.py @@ -266,14 +266,21 @@ def change_type_map( remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) super().change_type_map(type_map=type_map) if has_new_type: + xp = array_api_compat.array_namespace(self.scale) extend_shape = [len(type_map), *list(self.scale.shape[1:])] - extend_scale = np.ones(extend_shape, dtype=self.scale.dtype) - self.scale = np.concatenate([self.scale, extend_scale], axis=0) + extend_scale = xp.ones( + extend_shape, + dtype=self.scale.dtype, + device=array_api_compat.device(self.scale), + ) + self.scale = xp.concat([self.scale, extend_scale], axis=0) extend_shape = [len(type_map), *list(self.constant_matrix.shape[1:])] - extend_constant_matrix = np.zeros( - extend_shape, dtype=self.constant_matrix.dtype + extend_constant_matrix = xp.zeros( + extend_shape, + dtype=self.constant_matrix.dtype, + device=array_api_compat.device(self.constant_matrix), ) - self.constant_matrix = np.concatenate( + self.constant_matrix = xp.concat( [self.constant_matrix, extend_constant_matrix], axis=0 ) self.scale = self.scale[remap_index] diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index 163cd62387..d87c0eb5b7 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -142,9 +142,10 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" pass + @abstractmethod def get_min_nbor_dist(self) -> float | None: """Get the minimum distance between two atoms.""" - return self.min_nbor_dist + pass @abstractmethod def get_nnei(self) -> int: @@ -190,6 +191,17 @@ def update_sel( cls = cls.get_class_by_type(model_type) return cls.update_sel(train_data, type_map, local_jdata) + @abstractmethod + def get_observed_type_list(self) -> list[str]: + """Get observed types (elements) of the model during data statistics. + + Returns + ------- + list[str] + A list of the observed type names in this model. + """ + pass + def enable_compression( self, table_extrapolate: float = 5, @@ -255,10 +267,4 @@ class BaseModel(make_base_model()): Backend-independent BaseModel class. """ - def __init__(self) -> None: - self.model_def_script = "" - self.min_nbor_dist = None - - def get_model_def_script(self) -> str: - """Get the model definition script.""" - return self.model_def_script + pass diff --git a/deepmd/dpmodel/model/dipole_model.py b/deepmd/dpmodel/model/dipole_model.py index 421dd1b10f..fa5a76e0af 100644 --- a/deepmd/dpmodel/model/dipole_model.py +++ b/deepmd/dpmodel/model/dipole_model.py @@ -9,6 +9,9 @@ from deepmd.dpmodel.atomic_model import ( DPDipoleAtomicModel, ) +from deepmd.dpmodel.common import ( + NativeOP, +) from deepmd.dpmodel.model.base_model import ( BaseModel, ) @@ -20,13 +23,11 @@ make_model, ) -DPDipoleModel_ = make_model(DPDipoleAtomicModel) +DPDipoleModel_ = make_model(DPDipoleAtomicModel, T_Bases=(NativeOP, BaseModel)) @BaseModel.register("dipole") class DipoleModel(DPModelCommon, DPDipoleModel_): - model_type = "dipole" - def __init__( self, *args: Any, diff --git a/deepmd/dpmodel/model/dos_model.py b/deepmd/dpmodel/model/dos_model.py index dded7b076a..b75c9a2bcc 100644 --- a/deepmd/dpmodel/model/dos_model.py +++ b/deepmd/dpmodel/model/dos_model.py @@ -9,6 +9,9 @@ from deepmd.dpmodel.atomic_model import ( DPDOSAtomicModel, ) +from deepmd.dpmodel.common import ( + NativeOP, +) from deepmd.dpmodel.model.base_model import ( BaseModel, ) @@ -20,13 +23,11 @@ make_model, ) -DPDOSModel_ = make_model(DPDOSAtomicModel) +DPDOSModel_ = make_model(DPDOSAtomicModel, T_Bases=(NativeOP, BaseModel)) @BaseModel.register("dos") class DOSModel(DPModelCommon, DPDOSModel_): - model_type = "dos" - def __init__( self, *args: Any, diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index b6db2272b7..5a1a253596 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/deepmd/dpmodel/model/dp_model.py @@ -53,7 +53,7 @@ def update_sel( def get_fitting_net(self) -> BaseFitting: """Get the fitting network.""" - return self.atomic_model.fitting + return self.atomic_model.fitting_net def get_descriptor(self) -> BaseDescriptor: """Get the descriptor.""" diff --git a/deepmd/dpmodel/model/dp_zbl_model.py b/deepmd/dpmodel/model/dp_zbl_model.py index 81c2476447..d864b3b61e 100644 --- a/deepmd/dpmodel/model/dp_zbl_model.py +++ b/deepmd/dpmodel/model/dp_zbl_model.py @@ -9,6 +9,9 @@ from deepmd.dpmodel.atomic_model.linear_atomic_model import ( DPZBLLinearEnergyAtomicModel, ) +from deepmd.dpmodel.common import ( + NativeOP, +) from deepmd.dpmodel.model.base_model import ( BaseModel, ) @@ -23,13 +26,11 @@ make_model, ) -DPZBLModel_ = make_model(DPZBLLinearEnergyAtomicModel) +DPZBLModel_ = make_model(DPZBLLinearEnergyAtomicModel, T_Bases=(NativeOP, BaseModel)) @BaseModel.register("zbl") class DPZBLModel(DPZBLModel_): - model_type = "zbl" - def __init__( self, *args: Any, diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py index 85fe98c7d8..57b518d75d 100644 --- a/deepmd/dpmodel/model/ener_model.py +++ b/deepmd/dpmodel/model/ener_model.py @@ -12,6 +12,9 @@ from deepmd.dpmodel.atomic_model import ( DPEnergyAtomicModel, ) +from deepmd.dpmodel.common import ( + NativeOP, +) from deepmd.dpmodel.model.base_model import ( BaseModel, ) @@ -26,7 +29,7 @@ make_model, ) -DPEnergyModel_ = make_model(DPEnergyAtomicModel) +DPEnergyModel_ = make_model(DPEnergyAtomicModel, T_Bases=(NativeOP, BaseModel)) @BaseModel.register("ener") diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 71169dc64f..f06451b3fa 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -20,12 +20,8 @@ GLOBAL_NP_FLOAT_PRECISION, PRECISION_DICT, RESERVED_PRECISION_DICT, - NativeOP, get_xp_precision, ) -from deepmd.dpmodel.model.base_model import ( - BaseModel, -) from deepmd.dpmodel.output_def import ( FittingOutputDef, ModelOutputDef, @@ -39,6 +35,9 @@ nlist_distinguish_types, normalize_coord, ) +from deepmd.utils.path import ( + DPPath, +) from .transform_output import ( communicate_extended_output, @@ -138,7 +137,10 @@ def model_call_from_call_lower( return model_predict -def make_model(T_AtomicModel: type[BaseAtomicModel]) -> type: +def make_model( + T_AtomicModel: type[BaseAtomicModel], + T_Bases: tuple[type, ...] = (), +) -> type: """Make a model as a derived class of an atomic model. The model provide two interfaces. @@ -153,6 +155,9 @@ def make_model(T_AtomicModel: type[BaseAtomicModel]) -> type: ---------- T_AtomicModel The atomic model. + T_Bases + Additional base classes for the returned model class. + Defaults to ``()``. For example, dpmodel passes ``(NativeOP,)``. Returns ------- @@ -161,7 +166,7 @@ def make_model(T_AtomicModel: type[BaseAtomicModel]) -> type: """ - class CM(NativeOP, BaseModel): + class CM(*T_Bases): def __init__( self, *args: Any, @@ -169,7 +174,8 @@ def __init__( atomic_model_: T_AtomicModel | None = None, **kwargs: Any, ) -> None: - BaseModel.__init__(self) + self.model_def_script = "" + self.min_nbor_dist = None if atomic_model_ is not None: self.atomic_model: T_AtomicModel = atomic_model_ else: @@ -370,11 +376,30 @@ def forward_common_atomic( def get_out_bias(self) -> Array: """Get the output bias.""" - return self.atomic_model.out_bias + return self.atomic_model.get_out_bias() + + def get_observed_type_list(self) -> list[str]: + """Get observed types (elements) of the model during data statistics. + + Returns + ------- + list[str] + A list of the observed type names in this model. + """ + type_map = self.get_type_map() + out_bias = self.get_out_bias()[0] + assert out_bias is not None, "No out_bias found in the model." + assert out_bias.ndim == 2, "The supported out_bias should be a 2D array." + assert out_bias.shape[0] == len(type_map), ( + "The out_bias shape does not match the type_map length." + ) + xp = array_api_compat.array_namespace(out_bias) + bias_mask = xp.any(xp.abs(out_bias) > 1e-6, axis=-1) + return [type_map[i] for i in range(len(type_map)) if bias_mask[i]] def set_out_bias(self, out_bias: Array) -> None: """Set the output bias.""" - self.atomic_model.out_bias = out_bias + self.atomic_model.set_out_bias(out_bias) def change_out_bias( self, @@ -555,7 +580,7 @@ def _format_nlist( m_real_nei = nlist >= 0 ret = xp.where(m_real_nei, nlist, 0) coord0 = extended_coord[:, :n_nloc, :] - index = ret.reshape(n_nf, n_nloc * n_nnei, 1).repeat(3, axis=2) + index = xp.tile(ret.reshape(n_nf, n_nloc * n_nnei, 1), (1, 1, 3)) coord1 = xp.take_along_axis(extended_coord, index, axis=1) coord1 = coord1.reshape(n_nf, n_nloc, n_nnei, 3) rr = xp.linalg.norm(coord0[:, :, None, :] - coord1, axis=-1) @@ -591,12 +616,17 @@ def do_grad_c( return self.atomic_model.do_grad_c(var_name) def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Any = None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - self.atomic_model.change_type_map(type_map=type_map) + self.atomic_model.change_type_map( + type_map=type_map, + model_with_new_type_stat=model_with_new_type_stat.atomic_model + if model_with_new_type_stat is not None + else None, + ) def serialize(self) -> dict: return self.atomic_model.serialize() @@ -684,6 +714,31 @@ def atomic_output_def(self) -> FittingOutputDef: """Get the output def of the atomic model.""" return self.atomic_model.atomic_output_def() + def compute_or_load_stat( + self, + sampled_func: Callable[[], Any], + stat_file_path: DPPath | None = None, + ) -> None: + """Compute or load the statistics parameters of the model. + + Parameters + ---------- + sampled_func + The lazy sampled function to get data frames from different + data systems. + stat_file_path + The path to the stat file. + """ + self.atomic_model.compute_or_load_stat(sampled_func, stat_file_path) + + def get_model_def_script(self) -> str: + """Get the model definition script.""" + return self.model_def_script + + def get_min_nbor_dist(self) -> float | None: + """Get the minimum distance between two atoms.""" + return self.min_nbor_dist + def get_ntypes(self) -> int: """Get the number of types.""" return len(self.get_type_map()) diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 339998aa89..b560366457 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -115,9 +115,12 @@ def get_standard_model(data: dict) -> EnergyModel: def get_zbl_model(data: dict) -> DPZBLModel: + data = copy.deepcopy(data) data["descriptor"]["ntypes"] = len(data["type_map"]) + data["descriptor"]["type_map"] = data["type_map"] descriptor = BaseDescriptor(**data["descriptor"]) fitting_type = data["fitting_net"].pop("type") + data["fitting_net"]["type_map"] = data["type_map"] if fitting_type == "ener": fitting = EnergyFittingNet( ntypes=descriptor.get_ntypes(), diff --git a/deepmd/dpmodel/model/polar_model.py b/deepmd/dpmodel/model/polar_model.py index 057410f280..5031166a5e 100644 --- a/deepmd/dpmodel/model/polar_model.py +++ b/deepmd/dpmodel/model/polar_model.py @@ -9,6 +9,9 @@ from deepmd.dpmodel.atomic_model import ( DPPolarAtomicModel, ) +from deepmd.dpmodel.common import ( + NativeOP, +) from deepmd.dpmodel.model.base_model import ( BaseModel, ) @@ -20,13 +23,11 @@ make_model, ) -DPPolarModel_ = make_model(DPPolarAtomicModel) +DPPolarModel_ = make_model(DPPolarAtomicModel, T_Bases=(NativeOP, BaseModel)) @BaseModel.register("polar") class PolarModel(DPModelCommon, DPPolarModel_): - model_type = "polar" - def __init__( self, *args: Any, diff --git a/deepmd/dpmodel/model/property_model.py b/deepmd/dpmodel/model/property_model.py index ea34609393..bc1657f0bd 100644 --- a/deepmd/dpmodel/model/property_model.py +++ b/deepmd/dpmodel/model/property_model.py @@ -9,6 +9,9 @@ from deepmd.dpmodel.atomic_model import ( DPPropertyAtomicModel, ) +from deepmd.dpmodel.common import ( + NativeOP, +) from deepmd.dpmodel.model.base_model import ( BaseModel, ) @@ -23,7 +26,7 @@ make_model, ) -DPPropertyModel_ = make_model(DPPropertyAtomicModel) +DPPropertyModel_ = make_model(DPPropertyAtomicModel, T_Bases=(NativeOP, BaseModel)) @BaseModel.register("property") diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index e62442887c..510ca5dce2 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -17,6 +17,9 @@ from deepmd.dpmodel.common import ( NativeOP, ) +from deepmd.dpmodel.model.base_model import ( + BaseModel, +) from deepmd.dpmodel.model.make_model import ( make_model, ) @@ -343,9 +346,9 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "SpinModel": - backbone_model_obj = make_model(DPAtomicModel).deserialize( - data["backbone_model"] - ) + backbone_model_obj = make_model( + DPAtomicModel, T_Bases=(NativeOP, BaseModel) + ).deserialize(data["backbone_model"]) spin = Spin.deserialize(data["spin"]) return cls( backbone_model=backbone_model_obj, diff --git a/deepmd/jax/model/hlo.py b/deepmd/jax/model/hlo.py index 7eb4e2c4b3..c79bc727cf 100644 --- a/deepmd/jax/model/hlo.py +++ b/deepmd/jax/model/hlo.py @@ -265,6 +265,10 @@ def deserialize(cls, data: dict) -> "BaseModel": """ raise NotImplementedError("Not implemented") + def get_observed_type_list(self) -> list[str]: + """Get observed types (elements) of the model during data statistics.""" + raise NotImplementedError("Not implemented for HLO model") + def get_model_def_script(self) -> str: """Get the model definition script.""" return self.model_def_script diff --git a/deepmd/pd/model/model/frozen.py b/deepmd/pd/model/model/frozen.py index 365202dd6c..c4d5762530 100644 --- a/deepmd/pd/model/model/frozen.py +++ b/deepmd/pd/model/model/frozen.py @@ -178,6 +178,10 @@ def update_sel( """ return local_jdata, None + def get_observed_type_list(self) -> list[str]: + """Get observed types (elements) of the model during data statistics.""" + return self.model.get_observed_type_list() + def model_output_type(self) -> str: """Get the output type for the model.""" return self.model.model_output_type() diff --git a/deepmd/pd/model/model/make_model.py b/deepmd/pd/model/model/make_model.py index a127c9c367..d98bee7ed9 100644 --- a/deepmd/pd/model/model/make_model.py +++ b/deepmd/pd/model/model/make_model.py @@ -202,6 +202,24 @@ def forward_common( def get_out_bias(self) -> paddle.Tensor: return self.atomic_model.get_out_bias() + def get_observed_type_list(self) -> list[str]: + """Get observed types (elements) of the model during data statistics. + + Returns + ------- + list[str] + A list of the observed type names in this model. + """ + type_map = self.get_type_map() + out_bias = self.get_out_bias()[0] + assert out_bias is not None, "No out_bias found in the model." + assert out_bias.ndim == 2, "The supported out_bias should be a 2D array." + assert out_bias.shape[0] == len(type_map), ( + "The out_bias shape does not match the type_map length." + ) + bias_mask = paddle.any(paddle.abs(out_bias) > 1e-6, axis=-1) + return [type_map[i] for i in range(len(type_map)) if bias_mask[i]] + def set_out_bias(self, out_bias: paddle.Tensor) -> None: self.atomic_model.set_out_bias(out_bias) diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index 0ccf539757..920b83d12b 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -104,7 +104,12 @@ def init_out_stat(self) -> None: self.register_buffer("out_bias", out_bias_data) self.register_buffer("out_std", out_std_data) + def get_out_bias(self) -> torch.Tensor: + """Get the output bias.""" + return self.out_bias + def set_out_bias(self, out_bias: torch.Tensor) -> None: + """Set the output bias.""" self.out_bias = out_bias def __setitem__(self, key: str, value: torch.Tensor) -> None: @@ -477,6 +482,8 @@ def change_out_bias( model_forward=self._get_forward_wrapper_func(), rcond=self.rcond, preset_bias=self.preset_out_bias, + stats_distinguish_types=self.get_compute_stats_distinguish_types(), + intensive=self.get_intensive(), ) self._store_out_stat(delta_bias, out_std, add=True) elif bias_adjust_mode == "set-by-statistic": diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index a71427d5e9..78fa0c3cf7 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -299,9 +299,6 @@ def forward_atomic( ) return fit_ret - def get_out_bias(self) -> torch.Tensor: - return self.out_bias - def compute_or_load_stat( self, sampled_func: Callable[[], list[dict]], diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 96b3baf6ec..de3acfcaca 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -126,9 +126,6 @@ def need_sorted_nlist_for_lower(self) -> bool: """Returns whether the atomic model needs sorted nlist when using `forward_lower`.""" return True - def get_out_bias(self) -> torch.Tensor: - return self.out_bias - def get_rcut(self) -> float: """Get the cut-off radius.""" return max(self.get_model_rcuts()) diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py index a77e5391f8..4424509776 100644 --- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py +++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py @@ -131,9 +131,6 @@ def fitting_output_def(self) -> FittingOutputDef: ] ) - def get_out_bias(self) -> torch.Tensor: - return self.out_bias - def get_rcut(self) -> float: return self.rcut diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index 7dcd035412..1680d1e258 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/deepmd/pt/model/model/ener_model.py @@ -44,32 +44,6 @@ def enable_hessian(self) -> None: self.requires_hessian("energy") self._hessian_enabled = True - @torch.jit.export - def get_observed_type_list(self) -> list[str]: - """Get observed types (elements) of the model during data statistics. - - Returns - ------- - observed_type_list: a list of the observed types in this model. - """ - type_map = self.get_type_map() - out_bias = self.atomic_model.get_out_bias()[0] - - assert out_bias is not None, "No out_bias found in the model." - assert out_bias.dim() == 2, "The supported out_bias should be a 2D tensor." - assert out_bias.size(0) == len(type_map), ( - "The out_bias shape does not match the type_map length." - ) - bias_mask = ( - torch.gt(torch.abs(out_bias), 1e-6).any(dim=-1).detach().cpu() - ) # 1e-6 for stability - - observed_type_list: list[str] = [] - for i in range(len(type_map)): - if bias_mask[i]: - observed_type_list.append(type_map[i]) - return observed_type_list - def translated_output_def(self) -> dict[str, Any]: out_def_data = self.model_output_def().get_data() output_def = { diff --git a/deepmd/pt/model/model/frozen.py b/deepmd/pt/model/model/frozen.py index 402e78e95c..d44c608ade 100644 --- a/deepmd/pt/model/model/frozen.py +++ b/deepmd/pt/model/model/frozen.py @@ -198,6 +198,10 @@ def update_sel( """ return local_jdata, None + def get_observed_type_list(self) -> list[str]: + """Get observed types (elements) of the model during data statistics.""" + return self.model.get_observed_type_list() + @torch.jit.export def model_output_type(self) -> str: """Get the output type for the model.""" diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index ad4c35bcc9..2c08c4f6c5 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -589,6 +589,32 @@ def compute_or_load_stat( """Compute or load the statistics.""" return self.atomic_model.compute_or_load_stat(sampled_func, stat_file_path) + @torch.jit.export + def get_observed_type_list(self) -> list[str]: + """Get observed types (elements) of the model during data statistics. + + Returns + ------- + observed_type_list: a list of the observed types in this model. + """ + type_map = self.get_type_map() + out_bias = self.atomic_model.get_out_bias()[0] + + assert out_bias is not None, "No out_bias found in the model." + assert out_bias.dim() == 2, "The supported out_bias should be a 2D tensor." + assert out_bias.size(0) == len(type_map), ( + "The out_bias shape does not match the type_map length." + ) + bias_mask = ( + torch.gt(torch.abs(out_bias), 1e-6).any(dim=-1).detach().cpu() + ) # 1e-6 for stability + + observed_type_list: list[str] = [] + for i in range(len(type_map)): + if bias_mask[i]: + observed_type_list.append(type_map[i]) + return observed_type_list + def get_sel(self) -> list[int]: """Returns the number of selected atoms for each type.""" return self.atomic_model.get_sel() diff --git a/deepmd/pt_expt/atomic_model/__init__.py b/deepmd/pt_expt/atomic_model/__init__.py index 51ee9f4186..6ceb116d85 100644 --- a/deepmd/pt_expt/atomic_model/__init__.py +++ b/deepmd/pt_expt/atomic_model/__init__.py @@ -1,12 +1 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from .dp_atomic_model import ( - DPAtomicModel, -) -from .energy_atomic_model import ( - DPEnergyAtomicModel, -) - -__all__ = [ - "DPAtomicModel", - "DPEnergyAtomicModel", -] diff --git a/deepmd/pt_expt/atomic_model/dp_atomic_model.py b/deepmd/pt_expt/atomic_model/dp_atomic_model.py deleted file mode 100644 index b87935bd09..0000000000 --- a/deepmd/pt_expt/atomic_model/dp_atomic_model.py +++ /dev/null @@ -1,45 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later - -import torch - -from deepmd.dpmodel.atomic_model.dp_atomic_model import DPAtomicModel as DPAtomicModelDP -from deepmd.pt_expt.common import ( - register_dpmodel_mapping, - torch_module, -) -from deepmd.pt_expt.descriptor.base_descriptor import ( - BaseDescriptor, -) -from deepmd.pt_expt.fitting.base_fitting import ( - BaseFitting, -) - - -@torch_module -class DPAtomicModel(DPAtomicModelDP): - base_descriptor_cls = BaseDescriptor - base_fitting_cls = BaseFitting - - def forward( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - ) -> dict[str, torch.Tensor]: - return self.forward_atomic( - extended_coord, - extended_atype, - nlist, - mapping=mapping, - fparam=fparam, - aparam=aparam, - ) - - -register_dpmodel_mapping( - DPAtomicModelDP, - lambda v: DPAtomicModel.deserialize(v.serialize()), -) diff --git a/deepmd/pt_expt/atomic_model/energy_atomic_model.py b/deepmd/pt_expt/atomic_model/energy_atomic_model.py deleted file mode 100644 index 5f34d215cf..0000000000 --- a/deepmd/pt_expt/atomic_model/energy_atomic_model.py +++ /dev/null @@ -1,27 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -from deepmd.dpmodel.atomic_model.energy_atomic_model import ( - DPEnergyAtomicModel as DPEnergyAtomicModelDP, -) -from deepmd.pt_expt.common import ( - register_dpmodel_mapping, -) - -from .dp_atomic_model import ( - DPAtomicModel, -) - - -class DPEnergyAtomicModel(DPAtomicModel): - """Energy atomic model for pt_expt backend. - - This is a thin wrapper around DPAtomicModel that validates - the fitting is an EnergyFittingNet or InvarFitting. - """ - - pass - - -register_dpmodel_mapping( - DPEnergyAtomicModelDP, - lambda v: DPEnergyAtomicModel.deserialize(v.serialize()), -) diff --git a/deepmd/pt_expt/common.py b/deepmd/pt_expt/common.py index d00e016f6c..d46da28ba8 100644 --- a/deepmd/pt_expt/common.py +++ b/deepmd/pt_expt/common.py @@ -297,6 +297,11 @@ def dpmodel_setattr(obj: torch.nn.Module, name: str, value: Any) -> tuple[bool, if name in obj._buffers: obj._buffers[name] = tensor return True, tensor + # If the attribute already exists as a regular attribute (e.g. set to + # None during __init__ and later reassigned as an ndarray in + # deserialize), remove it first so register_buffer doesn't conflict. + if hasattr(obj, name) and name not in obj._buffers: + delattr(obj, name) obj.register_buffer(name, tensor) return True, tensor diff --git a/deepmd/pt_expt/model/__init__.py b/deepmd/pt_expt/model/__init__.py index 5d1c5ffb5d..da120091e0 100644 --- a/deepmd/pt_expt/model/__init__.py +++ b/deepmd/pt_expt/model/__init__.py @@ -1,8 +1,32 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from .dipole_model import ( + DipoleModel, +) +from .dos_model import ( + DOSModel, +) +from .dp_zbl_model import ( + DPZBLModel, +) from .ener_model import ( EnergyModel, ) +from .model import ( + BaseModel, +) +from .polar_model import ( + PolarModel, +) +from .property_model import ( + PropertyModel, +) __all__ = [ + "BaseModel", + "DOSModel", + "DPZBLModel", + "DipoleModel", "EnergyModel", + "PolarModel", + "PropertyModel", ] diff --git a/deepmd/pt_expt/model/dipole_model.py b/deepmd/pt_expt/model/dipole_model.py new file mode 100644 index 0000000000..73ebba6bac --- /dev/null +++ b/deepmd/pt_expt/model/dipole_model.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.dpmodel.atomic_model import ( + DPDipoleAtomicModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) + +from .make_model import ( + make_model, +) +from .model import ( + BaseModel, +) + +DPDipoleModel_ = make_model(DPDipoleAtomicModel, T_Bases=(BaseModel,)) + + +@BaseModel.register("dipole") +class DipoleModel(DPModelCommon, DPDipoleModel_): + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPDipoleModel_.__init__(self, *args, **kwargs) + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["dipole"] = model_ret["dipole"] + model_predict["global_dipole"] = model_ret["dipole_redu"] + if self.do_grad_r("dipole") and model_ret["dipole_derv_r"] is not None: + model_predict["force"] = model_ret["dipole_derv_r"] + if self.do_grad_c("dipole") and model_ret["dipole_derv_c_redu"] is not None: + model_predict["virial"] = model_ret["dipole_derv_c_redu"] + if do_atomic_virial and model_ret["dipole_derv_c"] is not None: + model_predict["atom_virial"] = model_ret["dipole_derv_c"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["dipole"] = model_ret["dipole"] + model_predict["global_dipole"] = model_ret["dipole_redu"] + if self.do_grad_r("dipole") and model_ret.get("dipole_derv_r") is not None: + model_predict["extended_force"] = model_ret["dipole_derv_r"] + if self.do_grad_c("dipole") and model_ret.get("dipole_derv_c_redu") is not None: + model_predict["virial"] = model_ret["dipole_derv_c_redu"] + if do_atomic_virial and model_ret.get("dipole_derv_c") is not None: + model_predict["extended_virial"] = model_ret["dipole_derv_c"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + output_def = { + "dipole": out_def_data["dipole"], + "global_dipole": out_def_data["dipole_redu"], + } + if self.do_grad_r("dipole"): + output_def["force"] = out_def_data["dipole_derv_r"] + output_def["force"].squeeze(-2) + if self.do_grad_c("dipole"): + output_def["virial"] = out_def_data["dipole_derv_c_redu"] + output_def["virial"].squeeze(-2) + output_def["atom_virial"] = out_def_data["dipole_derv_c"] + output_def["atom_virial"].squeeze(-2) + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_def + + def forward_lower_exportable( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> torch.nn.Module: + model = self + + def fn( + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + extended_coord = extended_coord.detach().requires_grad_(True) + return model.forward_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + + return make_fx(fn)( + extended_coord, extended_atype, nlist, mapping, fparam, aparam + ) diff --git a/deepmd/pt_expt/model/dos_model.py b/deepmd/pt_expt/model/dos_model.py new file mode 100644 index 0000000000..137c2b2901 --- /dev/null +++ b/deepmd/pt_expt/model/dos_model.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.dpmodel.atomic_model import ( + DPDOSAtomicModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) + +from .make_model import ( + make_model, +) +from .model import ( + BaseModel, +) + +DPDOSModel_ = make_model(DPDOSAtomicModel, T_Bases=(BaseModel,)) + + +@BaseModel.register("dos") +class DOSModel(DPModelCommon, DPDOSModel_): + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPDOSModel_.__init__(self, *args, **kwargs) + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["atom_dos"] = model_ret["dos"] + model_predict["dos"] = model_ret["dos_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["atom_dos"] = model_ret["dos"] + model_predict["dos"] = model_ret["dos_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + output_def = { + "atom_dos": out_def_data["dos"], + "dos": out_def_data["dos_redu"], + } + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_def + + def forward_lower_exportable( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> torch.nn.Module: + model = self + + def fn( + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + extended_coord = extended_coord.detach().requires_grad_(True) + return model.forward_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + + return make_fx(fn)( + extended_coord, extended_atype, nlist, mapping, fparam, aparam + ) diff --git a/deepmd/pt_expt/model/dp_zbl_model.py b/deepmd/pt_expt/model/dp_zbl_model.py new file mode 100644 index 0000000000..c4bb668353 --- /dev/null +++ b/deepmd/pt_expt/model/dp_zbl_model.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.dpmodel.atomic_model.linear_atomic_model import ( + DPZBLLinearEnergyAtomicModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) + +from .make_model import ( + make_model, +) +from .model import ( + BaseModel, +) + +DPZBLModel_ = make_model(DPZBLLinearEnergyAtomicModel, T_Bases=(BaseModel,)) + + +@BaseModel.register("zbl") +class DPZBLModel(DPModelCommon, DPZBLModel_): + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPZBLModel_.__init__(self, *args, **kwargs) + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["atom_energy"] = model_ret["energy"] + model_predict["energy"] = model_ret["energy_redu"] + if self.do_grad_r("energy") and model_ret["energy_derv_r"] is not None: + model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2) + if self.do_grad_c("energy") and model_ret["energy_derv_c_redu"] is not None: + model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) + if do_atomic_virial and model_ret["energy_derv_c"] is not None: + model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-2) + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["atom_energy"] = model_ret["energy"] + model_predict["energy"] = model_ret["energy_redu"] + if self.do_grad_r("energy") and model_ret.get("energy_derv_r") is not None: + model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) + if self.do_grad_c("energy") and model_ret.get("energy_derv_c_redu") is not None: + model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) + if do_atomic_virial and model_ret.get("energy_derv_c") is not None: + model_predict["extended_virial"] = model_ret["energy_derv_c"].squeeze( + -2 + ) + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + output_def = { + "atom_energy": out_def_data["energy"], + "energy": out_def_data["energy_redu"], + } + if self.do_grad_r("energy"): + output_def["force"] = out_def_data["energy_derv_r"] + output_def["force"].squeeze(-2) + if self.do_grad_c("energy"): + output_def["virial"] = out_def_data["energy_derv_c_redu"] + output_def["virial"].squeeze(-2) + output_def["atom_virial"] = out_def_data["energy_derv_c"] + output_def["atom_virial"].squeeze(-2) + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_def + + def forward_lower_exportable( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> torch.nn.Module: + model = self + + def fn( + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + extended_coord = extended_coord.detach().requires_grad_(True) + return model.forward_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + + return make_fx(fn)( + extended_coord, extended_atype, nlist, mapping, fparam, aparam + ) diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index 5f30f3a227..271028d2ff 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -8,23 +8,25 @@ make_fx, ) +from deepmd.dpmodel.atomic_model import ( + DPEnergyAtomicModel, +) from deepmd.dpmodel.model.dp_model import ( DPModelCommon, ) -from deepmd.pt_expt.atomic_model import ( - DPEnergyAtomicModel, -) from .make_model import ( make_model, ) +from .model import ( + BaseModel, +) -DPEnergyModel_ = make_model(DPEnergyAtomicModel) +DPEnergyModel_ = make_model(DPEnergyAtomicModel, T_Bases=(BaseModel,)) +@BaseModel.register("ener") class EnergyModel(DPModelCommon, DPEnergyModel_): - model_type = "ener" - def __init__( self, *args: Any, @@ -63,7 +65,7 @@ def forward( model_predict["mask"] = model_ret["mask"] return model_predict - def _forward_lower( + def forward_lower( self, extended_coord: torch.Tensor, extended_atype: torch.Tensor, @@ -97,25 +99,23 @@ def _forward_lower( model_predict["mask"] = model_ret["mask"] return model_predict - def forward_lower( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - return self._forward_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + output_def = { + "atom_energy": out_def_data["energy"], + "energy": out_def_data["energy_redu"], + } + if self.do_grad_r("energy"): + output_def["force"] = out_def_data["energy_derv_r"] + output_def["force"].squeeze(-2) + if self.do_grad_c("energy"): + output_def["virial"] = out_def_data["energy_derv_c_redu"] + output_def["virial"].squeeze(-2) + output_def["atom_virial"] = out_def_data["energy_derv_c"] + output_def["atom_virial"].squeeze(-2) + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_def def forward_lower_exportable( self, @@ -127,7 +127,7 @@ def forward_lower_exportable( aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> torch.nn.Module: - """Trace ``_forward_lower`` into an exportable module. + """Trace ``forward_lower`` into an exportable module. Uses ``make_fx`` to trace through ``torch.autograd.grad``, decomposing the backward pass into primitive ops. The returned @@ -156,7 +156,7 @@ def fn( aparam: torch.Tensor | None, ) -> dict[str, torch.Tensor]: extended_coord = extended_coord.detach().requires_grad_(True) - return model._forward_lower( + return model.forward_lower( extended_coord, extended_atype, nlist, diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 2785fade14..56cabafe81 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -18,7 +18,10 @@ ) -def make_model(T_AtomicModel: type[BaseAtomicModel]) -> type: +def make_model( + T_AtomicModel: type[BaseAtomicModel], + T_Bases: tuple[type, ...] = (), +) -> type: """Make a model as a derived class of an atomic model. Wraps dpmodel's make_model with torch.nn.Module and overrides @@ -28,6 +31,10 @@ def make_model(T_AtomicModel: type[BaseAtomicModel]) -> type: ---------- T_AtomicModel The atomic model. + T_Bases + Additional base classes for the returned model class. + For example, pass ``(BaseModel,)`` so that the concrete model + inherits the pt_expt ``BaseModel`` plugin registry. Returns ------- @@ -38,7 +45,7 @@ def make_model(T_AtomicModel: type[BaseAtomicModel]) -> type: DPModel = make_model_dp(T_AtomicModel) @torch_module - class CM(DPModel): + class CM(DPModel, *T_Bases): def forward(self, *args: Any, **kwargs: Any) -> dict[str, torch.Tensor]: """Default forward delegates to call(). diff --git a/deepmd/pt_expt/model/model.py b/deepmd/pt_expt/model/model.py new file mode 100644 index 0000000000..83842eaabd --- /dev/null +++ b/deepmd/pt_expt/model/model.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.model.base_model import ( + make_base_model, +) + + +class BaseModel(make_base_model()): + """Base class for pt_expt models. + + Provides the plugin registry so that model classes can be + registered with ``@BaseModel.register("ener")`` etc. + + See Also + -------- + deepmd.dpmodel.model.base_model.BaseBaseModel + Backend-independent BaseModel class. + """ + + pass diff --git a/deepmd/pt_expt/model/polar_model.py b/deepmd/pt_expt/model/polar_model.py new file mode 100644 index 0000000000..2bec72d4f7 --- /dev/null +++ b/deepmd/pt_expt/model/polar_model.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.dpmodel.atomic_model import ( + DPPolarAtomicModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) + +from .make_model import ( + make_model, +) +from .model import ( + BaseModel, +) + +DPPolarModel_ = make_model(DPPolarAtomicModel, T_Bases=(BaseModel,)) + + +@BaseModel.register("polar") +class PolarModel(DPModelCommon, DPPolarModel_): + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPPolarModel_.__init__(self, *args, **kwargs) + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["polar"] = model_ret["polarizability"] + model_predict["global_polar"] = model_ret["polarizability_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["polar"] = model_ret["polarizability"] + model_predict["global_polar"] = model_ret["polarizability_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + output_def = { + "polar": out_def_data["polarizability"], + "global_polar": out_def_data["polarizability_redu"], + } + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_def + + def forward_lower_exportable( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> torch.nn.Module: + model = self + + def fn( + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + extended_coord = extended_coord.detach().requires_grad_(True) + return model.forward_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + + return make_fx(fn)( + extended_coord, extended_atype, nlist, mapping, fparam, aparam + ) diff --git a/deepmd/pt_expt/model/property_model.py b/deepmd/pt_expt/model/property_model.py new file mode 100644 index 0000000000..50f8f0eeb4 --- /dev/null +++ b/deepmd/pt_expt/model/property_model.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.dpmodel.atomic_model import ( + DPPropertyAtomicModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) + +from .make_model import ( + make_model, +) +from .model import ( + BaseModel, +) + +DPPropertyModel_ = make_model(DPPropertyAtomicModel, T_Bases=(BaseModel,)) + + +@BaseModel.register("property") +class PropertyModel(DPModelCommon, DPPropertyModel_): + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPPropertyModel_.__init__(self, *args, **kwargs) + + def get_var_name(self) -> str: + """Get the name of the property.""" + return self.get_fitting_net().var_name + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + var_name = self.get_var_name() + model_predict = {} + model_predict[f"atom_{var_name}"] = model_ret[var_name] + model_predict[var_name] = model_ret[f"{var_name}_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + var_name = self.get_var_name() + model_predict = {} + model_predict[f"atom_{var_name}"] = model_ret[var_name] + model_predict[var_name] = model_ret[f"{var_name}_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + var_name = self.get_var_name() + output_def = { + f"atom_{var_name}": out_def_data[var_name], + var_name: out_def_data[f"{var_name}_redu"], + } + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_def + + def forward_lower_exportable( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> torch.nn.Module: + model = self + + def fn( + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + extended_coord = extended_coord.detach().requires_grad_(True) + return model.forward_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + + return make_fx(fn)( + extended_coord, extended_atype, nlist, mapping, fparam, aparam + ) diff --git a/source/tests/pt_expt/atomic_model/test_atomic_model_atomic_stat.py b/source/tests/common/dpmodel/test_atomic_model_atomic_stat.py similarity index 56% rename from source/tests/pt_expt/atomic_model/test_atomic_model_atomic_stat.py rename to source/tests/common/dpmodel/test_atomic_model_atomic_stat.py index c393ad4b3b..572c4aa696 100644 --- a/source/tests/pt_expt/atomic_model/test_atomic_model_atomic_stat.py +++ b/source/tests/common/dpmodel/test_atomic_model_atomic_stat.py @@ -1,471 +1,381 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import tempfile -import unittest -from pathlib import ( - Path, -) -from typing import ( - NoReturn, -) - -import h5py -import numpy as np -import torch - -from deepmd.dpmodel.output_def import ( - FittingOutputDef, - OutputVariableDef, -) -from deepmd.pt_expt.atomic_model import ( - DPAtomicModel, -) -from deepmd.pt_expt.descriptor.se_e2_a import ( - DescrptSeA, -) -from deepmd.pt_expt.fitting.base_fitting import ( - BaseFitting, -) -from deepmd.pt_expt.utils import ( - env, -) -from deepmd.utils.path import ( - DPPath, -) - -from ...pt.model.test_env_mat import ( - TestCaseSingleFrameWithNlist, -) - - -class FooFitting(BaseFitting, torch.nn.Module): - """Test fitting that returns fixed values for testing bias computation.""" - - def __init__(self): - torch.nn.Module.__init__(self) - BaseFitting.__init__(self) - - def output_def(self): - return FittingOutputDef( - [ - OutputVariableDef( - "foo", - [1], - reducible=True, - r_differentiable=True, - c_differentiable=True, - ), - OutputVariableDef( - "bar", - [1, 2], - reducible=True, - r_differentiable=True, - c_differentiable=True, - ), - ] - ) - - def serialize(self) -> dict: - return { - "@class": "Fitting", - "type": "foo", - "@version": 1, - } - - @classmethod - def deserialize(cls, data: dict): - return cls() - - def get_dim_fparam(self) -> int: - return 0 - - def get_dim_aparam(self) -> int: - return 0 - - def get_sel_type(self) -> list[int]: - return [] - - def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None - ) -> None: - pass - - def get_type_map(self) -> list[str]: - return [] - - def forward( - self, - descriptor: torch.Tensor, - atype: torch.Tensor, - gr: torch.Tensor | None = None, - g2: torch.Tensor | None = None, - h2: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - ): - nf, nloc, _ = descriptor.shape - ret = {} - ret["foo"] = ( - torch.Tensor( - [ - [1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - ] - ) - .view([nf, nloc, *self.output_def()["foo"].shape]) - .to(dtype=torch.float64, device=env.DEVICE) - ) - ret["bar"] = ( - torch.Tensor( - [ - [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], - [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], - ] - ) - .view([nf, nloc, *self.output_def()["bar"].shape]) - .to(dtype=torch.float64, device=env.DEVICE) - ) - return ret - - -class TestAtomicModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist): - def tearDown(self) -> None: - self.tempdir.cleanup() - - def setUp(self) -> None: - TestCaseSingleFrameWithNlist.setUp(self) - self.device = env.DEVICE - self.merged_output_stat = [ - { - "coord": torch.tensor(np.zeros([2, 3, 3]), device=self.device), - "atype": torch.tensor( - np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), device=self.device - ), - "atype_ext": torch.tensor( - np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), - device=self.device, - ), - "box": torch.tensor(np.zeros([2, 3, 3]), device=self.device), - "natoms": torch.tensor( - np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), - device=self.device, - ), - # bias of foo: 5, 6 - "atom_foo": torch.tensor( - np.array([[5.0, 5.0, 5.0], [5.0, 6.0, 7.0]]).reshape(2, 3, 1), - device=self.device, - ), - # bias of bar: [1, 5], [3, 2] - "bar": torch.tensor( - np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2), device=self.device - ), - "find_atom_foo": np.float32(1.0), - "find_bar": np.float32(1.0), - }, - { - "coord": torch.tensor(np.zeros([2, 3, 3]), device=self.device), - "atype": torch.tensor( - np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), device=self.device - ), - "atype_ext": torch.tensor( - np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), - device=self.device, - ), - "box": torch.tensor(np.zeros([2, 3, 3]), device=self.device), - "natoms": torch.tensor( - np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), - device=self.device, - ), - # bias of foo: 5, 6 from atomic label. - "foo": torch.tensor( - np.array([5.0, 7.0]).reshape(2, 1), device=self.device - ), - # bias of bar: [1, 5], [3, 2] - "bar": torch.tensor( - np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2), device=self.device - ), - "find_foo": np.float32(1.0), - "find_bar": np.float32(1.0), - }, - ] - self.tempdir = tempfile.TemporaryDirectory() - h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) - with h5py.File(h5file, "w") as f: - pass - self.stat_file_path = DPPath(h5file, "a") - - def test_output_stat(self) -> None: - """Test output statistics computation for pt_expt atomic model.""" - nf, nloc, nnei = self.nlist.shape - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ).to(self.device) - ft = FooFitting().to(self.device) - type_map = ["foo", "bar"] - md0 = DPAtomicModel( - ds, - ft, - type_map=type_map, - ).to(self.device) - args = [ - torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device), - torch.tensor(self.atype_ext, dtype=torch.int64, device=self.device), - torch.tensor(self.nlist, dtype=torch.int64, device=self.device), - ] - # nf x nloc - at = self.atype_ext[:, :nloc] - - def cvt_ret(x): - return {kk: vv.detach().cpu().numpy() for kk, vv in x.items()} - - # 1. test run without bias - # nf x na x odim - ret0 = md0.forward_common_atomic(*args) - ret0 = cvt_ret(ret0) - expected_ret0 = {} - expected_ret0["foo"] = np.array( - [ - [1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["foo"].shape]) - expected_ret0["bar"] = np.array( - [ - [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], - [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["bar"].shape]) - for kk in ["foo", "bar"]: - np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) - - # 2. test bias is applied - md0.compute_or_load_out_stat( - self.merged_output_stat, stat_file_path=self.stat_file_path - ) - ret1 = md0.forward_common_atomic(*args) - expected_std = np.ones( - (2, 2, 2), dtype=np.float64 - ) # 2 keys, 2 atypes, 2 max dims. - expected_std[0, :, :1] = np.array([0.0, 0.816496]).reshape( - 2, 1 - ) # updating std for foo based on [5.0, 5.0, 5.0], [5.0, 6.0, 7.0]] - np.testing.assert_almost_equal( - md0.out_std.detach().cpu().numpy(), expected_std, decimal=4 - ) - ret1 = cvt_ret(ret1) - # nt x odim - foo_bias = np.array([5.0, 6.0]).reshape(2, 1) - bar_bias = np.array([1.0, 5.0, 3.0, 2.0]).reshape(2, 1, 2) - expected_ret1 = {} - expected_ret1["foo"] = ret0["foo"] + foo_bias[at] - expected_ret1["bar"] = ret0["bar"] + bar_bias[at] - for kk in ["foo", "bar"]: - np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) - - # 3. test bias load from file - def raise_error() -> NoReturn: - raise RuntimeError - - md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) - ret2 = md0.forward_common_atomic(*args) - ret2 = cvt_ret(ret2) - for kk in ["foo", "bar"]: - np.testing.assert_almost_equal(ret1[kk], ret2[kk]) - np.testing.assert_almost_equal( - md0.out_std.detach().cpu().numpy(), expected_std, decimal=4 - ) - - # 4. test change bias - md0.change_out_bias( - self.merged_output_stat, bias_adjust_mode="change-by-statistic" - ) - # use atype_ext from merged_output_stat for inference (matching pt backend test) - args = [ - torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device), - self.merged_output_stat[0]["atype_ext"].to( - dtype=torch.int64, device=self.device - ), - torch.tensor(self.nlist, dtype=torch.int64, device=self.device), - ] - ret3 = md0.forward_common_atomic(*args) - ret3 = cvt_ret(ret3) - expected_std[0, :, :1] = np.array([1.24722, 0.47140]).reshape( - 2, 1 - ) # updating std for foo based on [4.0, 3.0, 2.0], [1.0, 1.0, 1.0]] - expected_ret3 = {} - # new bias [2.666, 1.333] - expected_ret3["foo"] = np.array( - [[3.6667, 4.6667, 4.3333], [6.6667, 6.3333, 7.3333]] - ).reshape(2, 3, 1) - for kk in ["foo"]: - np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk], decimal=4) - np.testing.assert_almost_equal( - md0.out_std.detach().cpu().numpy(), expected_std, decimal=4 - ) - - -class TestAtomicModelStatMergeGlobalAtomic( - unittest.TestCase, TestCaseSingleFrameWithNlist -): - """Test merging atomic and global stat when atomic label only covers some types.""" - - def tearDown(self) -> None: - self.tempdir.cleanup() - - def setUp(self) -> None: - TestCaseSingleFrameWithNlist.setUp(self) - self.device = env.DEVICE - self.merged_output_stat = [ - { - "coord": torch.tensor(np.zeros([2, 3, 3]), device=self.device), - "atype": torch.tensor( - np.array([[0, 0, 0], [0, 0, 0]], dtype=np.int32), - device=self.device, - ), - "atype_ext": torch.tensor( - np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), - device=self.device, - ), - "box": torch.tensor(np.zeros([2, 3, 3]), device=self.device), - "natoms": torch.tensor( - np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), - device=self.device, - ), - # bias of foo: 5.5, nan (only type 0 atoms) - "atom_foo": torch.tensor( - np.array([[5.0, 5.0, 5.0], [5.0, 6.0, 7.0]]).reshape(2, 3, 1), - device=self.device, - ), - # bias of bar: [1, 5], [3, 2] - "bar": torch.tensor( - np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2), - device=self.device, - ), - "find_atom_foo": np.float32(1.0), - "find_bar": np.float32(1.0), - }, - { - "coord": torch.tensor(np.zeros([2, 3, 3]), device=self.device), - "atype": torch.tensor( - np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), - device=self.device, - ), - "atype_ext": torch.tensor( - np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), - device=self.device, - ), - "box": torch.tensor(np.zeros([2, 3, 3]), device=self.device), - "natoms": torch.tensor( - np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), - device=self.device, - ), - # bias of foo: 5.5, 3 from global label. - "foo": torch.tensor( - np.array([5.0, 7.0]).reshape(2, 1), device=self.device - ), - # bias of bar: [1, 5], [3, 2] - "bar": torch.tensor( - np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2), - device=self.device, - ), - "find_foo": np.float32(1.0), - "find_bar": np.float32(1.0), - }, - ] - self.tempdir = tempfile.TemporaryDirectory() - h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) - with h5py.File(h5file, "w") as f: - pass - self.stat_file_path = DPPath(h5file, "a") - - def test_output_stat(self) -> None: - """Test merging atomic (type 0 only) and global stat for type 1.""" - nf, nloc, nnei = self.nlist.shape - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ).to(self.device) - ft = FooFitting().to(self.device) - type_map = ["foo", "bar"] - md0 = DPAtomicModel( - ds, - ft, - type_map=type_map, - ).to(self.device) - args = [ - torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device), - torch.tensor(self.atype_ext, dtype=torch.int64, device=self.device), - torch.tensor(self.nlist, dtype=torch.int64, device=self.device), - ] - # nf x nloc - at = self.atype_ext[:, :nloc] - - def cvt_ret(x): - return {kk: vv.detach().cpu().numpy() for kk, vv in x.items()} - - # 1. test run without bias - ret0 = md0.forward_common_atomic(*args) - ret0 = cvt_ret(ret0) - expected_ret0 = {} - expected_ret0["foo"] = np.array( - [ - [1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["foo"].shape]) - expected_ret0["bar"] = np.array( - [ - [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], - [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["bar"].shape]) - for kk in ["foo", "bar"]: - np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) - - # 2. test bias is applied - # foo: type 0 from atomic (mean=5.5), type 1 from global (lstsq=3.0) - md0.compute_or_load_out_stat( - self.merged_output_stat, stat_file_path=self.stat_file_path - ) - ret1 = md0.forward_common_atomic(*args) - ret1 = cvt_ret(ret1) - # nt x odim - foo_bias = np.array([5.5, 3.0]).reshape(2, 1) - bar_bias = np.array([1.0, 5.0, 3.0, 2.0]).reshape(2, 1, 2) - expected_ret1 = {} - expected_ret1["foo"] = ret0["foo"] + foo_bias[at] - expected_ret1["bar"] = ret0["bar"] + bar_bias[at] - for kk in ["foo", "bar"]: - np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) - - # 3. test bias load from file - def raise_error() -> NoReturn: - raise RuntimeError - - md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) - ret2 = md0.forward_common_atomic(*args) - ret2 = cvt_ret(ret2) - for kk in ["foo", "bar"]: - np.testing.assert_almost_equal(ret1[kk], ret2[kk]) - - # 4. test change bias - md0.change_out_bias( - self.merged_output_stat, bias_adjust_mode="change-by-statistic" - ) - # use atype_ext from merged_output_stat for inference - args = [ - torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device), - self.merged_output_stat[0]["atype_ext"].to( - dtype=torch.int64, device=self.device - ), - torch.tensor(self.nlist, dtype=torch.int64, device=self.device), - ] - ret3 = md0.forward_common_atomic(*args) - ret3 = cvt_ret(ret3) - expected_ret3 = {} - # new bias [2, -5] - expected_ret3["foo"] = np.array([[3, 4, -2], [6, 0, 1]]).reshape(2, 3, 1) - for kk in ["foo"]: - np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk], decimal=4) +# SPDX-License-Identifier: LGPL-3.0-or-later +import tempfile +import unittest +from pathlib import ( + Path, +) +from typing import ( + NoReturn, +) + +import h5py +import numpy as np + +from deepmd.dpmodel.atomic_model import ( + DPAtomicModel, +) +from deepmd.dpmodel.common import ( + NativeOP, +) +from deepmd.dpmodel.descriptor import ( + DescrptSeA, +) +from deepmd.dpmodel.fitting.base_fitting import ( + BaseFitting, +) +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + OutputVariableDef, +) +from deepmd.utils.path import ( + DPPath, +) + +from .case_single_frame_with_nlist import ( + TestCaseSingleFrameWithNlist, +) + + +class FooFitting(NativeOP, BaseFitting): + """Test fitting that returns fixed values for testing bias computation.""" + + def __init__(self): + pass + + def output_def(self): + return FittingOutputDef( + [ + OutputVariableDef( + "foo", + [1], + reducible=True, + r_differentiable=True, + c_differentiable=True, + ), + OutputVariableDef( + "bar", + [1, 2], + reducible=True, + r_differentiable=True, + c_differentiable=True, + ), + ] + ) + + def serialize(self) -> dict: + return { + "@class": "Fitting", + "type": "foo", + "@version": 1, + } + + @classmethod + def deserialize(cls, data: dict): + return cls() + + def get_dim_fparam(self) -> int: + return 0 + + def get_dim_aparam(self) -> int: + return 0 + + def get_sel_type(self) -> list[int]: + return [] + + def change_type_map( + self, type_map: list[str], model_with_new_type_stat=None + ) -> None: + pass + + def get_type_map(self) -> list[str]: + return [] + + def call( + self, + descriptor, + atype, + gr=None, + g2=None, + h2=None, + fparam=None, + aparam=None, + ): + nf, nloc, _ = descriptor.shape + ret = {} + ret["foo"] = np.array( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + ] + ).reshape([nf, nloc, *self.output_def()["foo"].shape]) + ret["bar"] = np.array( + [ + [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], + [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], + ] + ).reshape([nf, nloc, *self.output_def()["bar"].shape]) + return ret + + +class TestAtomicModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist): + def tearDown(self) -> None: + self.tempdir.cleanup() + + def setUp(self) -> None: + TestCaseSingleFrameWithNlist.setUp(self) + self.merged_output_stat = [ + { + "coord": np.zeros([2, 3, 3]), + "atype": np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), + "atype_ext": np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), + "box": np.zeros([2, 3, 3]), + "natoms": np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), + # bias of foo: 5, 6 + "atom_foo": np.array([[5.0, 5.0, 5.0], [5.0, 6.0, 7.0]]).reshape( + 2, 3, 1 + ), + # bias of bar: [1, 5], [3, 2] + "bar": np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2), + "find_atom_foo": np.float32(1.0), + "find_bar": np.float32(1.0), + }, + { + "coord": np.zeros([2, 3, 3]), + "atype": np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), + "atype_ext": np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), + "box": np.zeros([2, 3, 3]), + "natoms": np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), + # bias of foo: 5.5, 3 from global label. + "foo": np.array([5.0, 7.0]).reshape(2, 1), + # bias of bar: [1, 5], [3, 2] + "bar": np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2), + "find_foo": np.float32(1.0), + "find_bar": np.float32(1.0), + }, + ] + self.tempdir = tempfile.TemporaryDirectory() + h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) + with h5py.File(h5file, "w") as f: + pass + self.stat_file_path = DPPath(h5file, "a") + + def test_output_stat(self) -> None: + """Test output statistics computation for dpmodel atomic model.""" + nf, nloc, nnei = self.nlist.shape + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = FooFitting() + type_map = ["foo", "bar"] + md0 = DPAtomicModel( + ds, + ft, + type_map=type_map, + ) + args = [self.coord_ext, self.atype_ext, self.nlist] + # nf x nloc + at = self.atype_ext[:, :nloc] + + # 1. test run without bias + # nf x na x odim + ret0 = md0.forward_common_atomic(*args) + expected_ret0 = {} + expected_ret0["foo"] = np.array( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["foo"].shape]) + expected_ret0["bar"] = np.array( + [ + [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], + [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["bar"].shape]) + for kk in ["foo", "bar"]: + np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) + + # 2. test bias is applied + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + ret1 = md0.forward_common_atomic(*args) + expected_std = np.ones( + (2, 2, 2), dtype=np.float64 + ) # 2 keys, 2 atypes, 2 max dims. + expected_std[0, :, :1] = np.array([0.0, 0.816496]).reshape( + 2, 1 + ) # updating std for foo based on [5.0, 5.0, 5.0], [5.0, 6.0, 7.0]] + np.testing.assert_almost_equal(md0.out_std, expected_std, decimal=4) + # nt x odim + foo_bias = np.array([5.0, 6.0]).reshape(2, 1) + bar_bias = np.array([1.0, 5.0, 3.0, 2.0]).reshape(2, 1, 2) + expected_ret1 = {} + expected_ret1["foo"] = ret0["foo"] + foo_bias[at] + expected_ret1["bar"] = ret0["bar"] + bar_bias[at] + for kk in ["foo", "bar"]: + np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) + + # 3. test bias load from file + def raise_error() -> NoReturn: + raise RuntimeError + + md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) + ret2 = md0.forward_common_atomic(*args) + for kk in ["foo", "bar"]: + np.testing.assert_almost_equal(ret1[kk], ret2[kk]) + np.testing.assert_almost_equal(md0.out_std, expected_std, decimal=4) + + # 4. test change bias + md0.change_out_bias( + self.merged_output_stat, bias_adjust_mode="change-by-statistic" + ) + # use atype_ext from merged_output_stat for inference (matching pt backend test) + args = [ + self.coord_ext, + np.array(self.merged_output_stat[0]["atype_ext"], dtype=np.int64), + self.nlist, + ] + ret3 = md0.forward_common_atomic(*args) + expected_std[0, :, :1] = np.array([1.24722, 0.47140]).reshape( + 2, 1 + ) # updating std for foo based on [4.0, 3.0, 2.0], [1.0, 1.0, 1.0]] + expected_ret3 = {} + # new bias [2.666, 1.333] + expected_ret3["foo"] = np.array( + [[3.6667, 4.6667, 4.3333], [6.6667, 6.3333, 7.3333]] + ).reshape(2, 3, 1) + for kk in ["foo"]: + np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk], decimal=4) + np.testing.assert_almost_equal(md0.out_std, expected_std, decimal=4) + + +class TestAtomicModelStatMergeGlobalAtomic( + unittest.TestCase, TestCaseSingleFrameWithNlist +): + """Test merging atomic and global stat when atomic label only covers some types.""" + + def tearDown(self) -> None: + self.tempdir.cleanup() + + def setUp(self) -> None: + TestCaseSingleFrameWithNlist.setUp(self) + self.merged_output_stat = [ + { + "coord": np.zeros([2, 3, 3]), + "atype": np.array([[0, 0, 0], [0, 0, 0]], dtype=np.int32), + "atype_ext": np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), + "box": np.zeros([2, 3, 3]), + "natoms": np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), + # bias of foo: 5.5, nan (only type 0 atoms) + "atom_foo": np.array([[5.0, 5.0, 5.0], [5.0, 6.0, 7.0]]).reshape( + 2, 3, 1 + ), + # bias of bar: [1, 5], [3, 2] + "bar": np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2), + "find_atom_foo": np.float32(1.0), + "find_bar": np.float32(1.0), + }, + { + "coord": np.zeros([2, 3, 3]), + "atype": np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), + "atype_ext": np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), + "box": np.zeros([2, 3, 3]), + "natoms": np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), + # bias of foo: 5.5, 3 from global label. + "foo": np.array([5.0, 7.0]).reshape(2, 1), + # bias of bar: [1, 5], [3, 2] + "bar": np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2), + "find_foo": np.float32(1.0), + "find_bar": np.float32(1.0), + }, + ] + self.tempdir = tempfile.TemporaryDirectory() + h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) + with h5py.File(h5file, "w") as f: + pass + self.stat_file_path = DPPath(h5file, "a") + + def test_output_stat(self) -> None: + """Test merging atomic (type 0 only) and global stat for type 1.""" + nf, nloc, nnei = self.nlist.shape + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = FooFitting() + type_map = ["foo", "bar"] + md0 = DPAtomicModel( + ds, + ft, + type_map=type_map, + ) + args = [self.coord_ext, self.atype_ext, self.nlist] + # nf x nloc + at = self.atype_ext[:, :nloc] + + # 1. test run without bias + ret0 = md0.forward_common_atomic(*args) + expected_ret0 = {} + expected_ret0["foo"] = np.array( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["foo"].shape]) + expected_ret0["bar"] = np.array( + [ + [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], + [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["bar"].shape]) + for kk in ["foo", "bar"]: + np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) + + # 2. test bias is applied + # foo: type 0 from atomic (mean=5.5), type 1 from global (lstsq=3.0) + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + ret1 = md0.forward_common_atomic(*args) + # nt x odim + foo_bias = np.array([5.5, 3.0]).reshape(2, 1) + bar_bias = np.array([1.0, 5.0, 3.0, 2.0]).reshape(2, 1, 2) + expected_ret1 = {} + expected_ret1["foo"] = ret0["foo"] + foo_bias[at] + expected_ret1["bar"] = ret0["bar"] + bar_bias[at] + for kk in ["foo", "bar"]: + np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) + + # 3. test bias load from file + def raise_error() -> NoReturn: + raise RuntimeError + + md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) + ret2 = md0.forward_common_atomic(*args) + for kk in ["foo", "bar"]: + np.testing.assert_almost_equal(ret1[kk], ret2[kk]) + + # 4. test change bias + md0.change_out_bias( + self.merged_output_stat, bias_adjust_mode="change-by-statistic" + ) + # use atype_ext from merged_output_stat for inference + args = [ + self.coord_ext, + np.array(self.merged_output_stat[0]["atype_ext"], dtype=np.int64), + self.nlist, + ] + ret3 = md0.forward_common_atomic(*args) + expected_ret3 = {} + # new bias [2, -5] + expected_ret3["foo"] = np.array([[3, 4, -2], [6, 0, 1]]).reshape(2, 3, 1) + for kk in ["foo"]: + np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk], decimal=4) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt_expt/atomic_model/test_atomic_model_global_stat.py b/source/tests/common/dpmodel/test_atomic_model_global_stat.py similarity index 64% rename from source/tests/pt_expt/atomic_model/test_atomic_model_global_stat.py rename to source/tests/common/dpmodel/test_atomic_model_global_stat.py index e09e7c0c91..88470a57c9 100644 --- a/source/tests/pt_expt/atomic_model/test_atomic_model_global_stat.py +++ b/source/tests/common/dpmodel/test_atomic_model_global_stat.py @@ -1,759 +1,629 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import tempfile -import unittest -from pathlib import ( - Path, -) -from typing import ( - NoReturn, -) - -import h5py -import numpy as np -import torch - -from deepmd.dpmodel.atomic_model import DPAtomicModel as DPDPAtomicModel -from deepmd.dpmodel.output_def import ( - FittingOutputDef, - OutputVariableDef, -) -from deepmd.pt_expt.atomic_model import ( - DPAtomicModel, -) -from deepmd.pt_expt.descriptor.se_e2_a import ( - DescrptSeA, -) -from deepmd.pt_expt.fitting import ( - InvarFitting, -) -from deepmd.pt_expt.fitting.base_fitting import ( - BaseFitting, -) -from deepmd.pt_expt.utils import ( - env, -) -from deepmd.utils.path import ( - DPPath, -) - -from ...pt.model.test_env_mat import ( - TestCaseSingleFrameWithNlist, -) -from ...seed import ( - GLOBAL_SEED, -) - - -class FooFitting(BaseFitting, torch.nn.Module): - """Test fitting with multiple outputs for testing global statistics.""" - - def __init__(self): - torch.nn.Module.__init__(self) - BaseFitting.__init__(self) - - def output_def(self): - return FittingOutputDef( - [ - OutputVariableDef( - "foo", - [1], - reducible=True, - r_differentiable=True, - c_differentiable=True, - ), - OutputVariableDef( - "pix", - [1], - reducible=True, - r_differentiable=True, - c_differentiable=True, - ), - OutputVariableDef( - "bar", - [1, 2], - reducible=True, - r_differentiable=True, - c_differentiable=True, - ), - ] - ) - - def serialize(self) -> dict: - return { - "@class": "Fitting", - "type": "foo", - "@version": 1, - } - - @classmethod - def deserialize(cls, data: dict): - return cls() - - def get_dim_fparam(self) -> int: - return 0 - - def get_dim_aparam(self) -> int: - return 0 - - def get_sel_type(self) -> list[int]: - return [] - - def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None - ) -> None: - pass - - def get_type_map(self) -> list[str]: - return [] - - def forward( - self, - descriptor: torch.Tensor, - atype: torch.Tensor, - gr: torch.Tensor | None = None, - g2: torch.Tensor | None = None, - h2: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - ): - nf, nloc, _ = descriptor.shape - ret = {} - ret["foo"] = ( - torch.Tensor( - [ - [1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - ] - ) - .view([nf, nloc, *self.output_def()["foo"].shape]) - .to(dtype=torch.float64, device=env.DEVICE) - ) - ret["pix"] = ( - torch.Tensor( - [ - [3.0, 2.0, 1.0], - [6.0, 5.0, 4.0], - ] - ) - .view([nf, nloc, *self.output_def()["pix"].shape]) - .to(dtype=torch.float64, device=env.DEVICE) - ) - ret["bar"] = ( - torch.Tensor( - [ - [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], - [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], - ] - ) - .view([nf, nloc, *self.output_def()["bar"].shape]) - .to(dtype=torch.float64, device=env.DEVICE) - ) - return ret - - -def _to_numpy(x): - return x.detach().cpu().numpy() - - -class TestAtomicModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist): - def tearDown(self) -> None: - self.tempdir.cleanup() - - def setUp(self) -> None: - TestCaseSingleFrameWithNlist.setUp(self) - self.device = env.DEVICE - self.merged_output_stat = [ - { - "coord": torch.tensor(np.zeros([2, 3, 3]), device=self.device), - "atype": torch.tensor( - np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), - device=self.device, - ), - "atype_ext": torch.tensor( - np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), - device=self.device, - ), - "box": torch.tensor(np.zeros([2, 3, 3]), device=self.device), - "natoms": torch.tensor( - np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), - device=self.device, - ), - # bias of foo: 1, 3 - "foo": torch.tensor( - np.array([5.0, 7.0]).reshape(2, 1), device=self.device - ), - # no bias of pix - # bias of bar: [1, 5], [3, 2] - "bar": torch.tensor( - np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2), - device=self.device, - ), - "find_foo": np.float32(1.0), - "find_bar": np.float32(1.0), - } - ] - self.tempdir = tempfile.TemporaryDirectory() - h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) - with h5py.File(h5file, "w") as f: - pass - self.stat_file_path = DPPath(h5file, "a") - - def test_output_stat(self) -> None: - nf, nloc, nnei = self.nlist.shape - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ).to(self.device) - ft = FooFitting().to(self.device) - type_map = ["foo", "bar"] - md0 = DPAtomicModel( - ds, - ft, - type_map=type_map, - ).to(self.device) - args = [ - torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device), - torch.tensor(self.atype_ext, dtype=torch.int64, device=self.device), - torch.tensor(self.nlist, dtype=torch.int64, device=self.device), - ] - # nf x nloc - at = self.atype_ext[:, :nloc] - - def cvt_ret(x): - return {kk: _to_numpy(vv) for kk, vv in x.items()} - - # 1. test run without bias - # nf x na x odim - ret0 = md0.forward_common_atomic(*args) - ret0 = cvt_ret(ret0) - - expected_ret0 = {} - expected_ret0["foo"] = np.array( - [ - [1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["foo"].shape]) - expected_ret0["pix"] = np.array( - [ - [3.0, 2.0, 1.0], - [6.0, 5.0, 4.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["pix"].shape]) - expected_ret0["bar"] = np.array( - [ - [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], - [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["bar"].shape]) - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) - - # 2. test bias is applied - md0.compute_or_load_out_stat( - self.merged_output_stat, stat_file_path=self.stat_file_path - ) - ret1 = md0.forward_common_atomic(*args) - ret1 = cvt_ret(ret1) - expected_std = np.ones((3, 2, 2)) # 3 keys, 2 atypes, 2 max dims. - # nt x odim - foo_bias = np.array([1.0, 3.0]).reshape(2, 1) - bar_bias = np.array([1.0, 5.0, 3.0, 2.0]).reshape(2, 1, 2) - expected_ret1 = {} - expected_ret1["foo"] = ret0["foo"] + foo_bias[at] - expected_ret1["pix"] = ret0["pix"] - expected_ret1["bar"] = ret0["bar"] + bar_bias[at] - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) - np.testing.assert_almost_equal(_to_numpy(md0.out_std), expected_std) - - # 3. test bias load from file - def raise_error() -> NoReturn: - raise RuntimeError - - md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) - ret2 = md0.forward_common_atomic(*args) - ret2 = cvt_ret(ret2) - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret1[kk], ret2[kk]) - np.testing.assert_almost_equal(_to_numpy(md0.out_std), expected_std) - - # 4. test change bias - md0.change_out_bias( - self.merged_output_stat, bias_adjust_mode="change-by-statistic" - ) - # use atype_ext from merged_output_stat for inference (matching pt backend test) - args = [ - torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device), - self.merged_output_stat[0]["atype_ext"].to( - dtype=torch.int64, device=self.device - ), - torch.tensor(self.nlist, dtype=torch.int64, device=self.device), - ] - ret3 = md0.forward_common_atomic(*args) - ret3 = cvt_ret(ret3) - ## model output on foo: [[2, 3, 6], [5, 8, 9]] given bias [1, 3] - ## foo sumed: [11, 22] compared with [5, 7], fit target is [-6, -15] - ## fit bias is [1, -8] - ## old bias + fit bias [2, -5] - ## new model output is [[3, 4, -2], [6, 0, 1]], which sumed to [5, 7] - expected_ret3 = {} - expected_ret3["foo"] = np.array([[3, 4, -2], [6, 0, 1]]).reshape(2, 3, 1) - expected_ret3["pix"] = ret0["pix"] - for kk in ["foo", "pix"]: - np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk]) - # bar is too complicated to be manually computed. - np.testing.assert_almost_equal(_to_numpy(md0.out_std), expected_std) - - def test_preset_bias(self) -> None: - nf, nloc, nnei = self.nlist.shape - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ).to(self.device) - ft = FooFitting().to(self.device) - type_map = ["foo", "bar"] - preset_out_bias = { - "foo": [None, 2], - "bar": np.array([7.0, 5.0, 13.0, 11.0]).reshape(2, 1, 2), - } - md0 = DPAtomicModel( - ds, - ft, - type_map=type_map, - preset_out_bias=preset_out_bias, - ).to(self.device) - args = [ - torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device), - torch.tensor(self.atype_ext, dtype=torch.int64, device=self.device), - torch.tensor(self.nlist, dtype=torch.int64, device=self.device), - ] - # nf x nloc - at = self.atype_ext[:, :nloc] - - def cvt_ret(x): - return {kk: _to_numpy(vv) for kk, vv in x.items()} - - # 1. test run without bias - # nf x na x odim - ret0 = md0.forward_common_atomic(*args) - ret0 = cvt_ret(ret0) - expected_ret0 = {} - expected_ret0["foo"] = np.array( - [ - [1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["foo"].shape]) - expected_ret0["pix"] = np.array( - [ - [3.0, 2.0, 1.0], - [6.0, 5.0, 4.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["pix"].shape]) - expected_ret0["bar"] = np.array( - [ - [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], - [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["bar"].shape]) - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) - - # 2. test bias is applied - md0.compute_or_load_out_stat( - self.merged_output_stat, stat_file_path=self.stat_file_path - ) - ret1 = md0.forward_common_atomic(*args) - ret1 = cvt_ret(ret1) - # foo sums: [5, 7], - # given bias of type 1 being 2, the bias left for type 0 is [5-2*1, 7-2*2] = [3,3] - # the solution of type 0 is 1.8 - foo_bias = np.array([1.8, preset_out_bias["foo"][1]]).reshape(2, 1) - bar_bias = preset_out_bias["bar"] - expected_ret1 = {} - expected_ret1["foo"] = ret0["foo"] + foo_bias[at] - expected_ret1["pix"] = ret0["pix"] - expected_ret1["bar"] = ret0["bar"] + bar_bias[at] - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) - - # 3. test bias load from file - def raise_error() -> NoReturn: - raise RuntimeError - - md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) - ret2 = md0.forward_common_atomic(*args) - ret2 = cvt_ret(ret2) - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret1[kk], ret2[kk]) - - # 4. test change bias - md0.change_out_bias( - self.merged_output_stat, bias_adjust_mode="change-by-statistic" - ) - # use atype_ext from merged_output_stat for inference - args = [ - torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device), - self.merged_output_stat[0]["atype_ext"].to( - dtype=torch.int64, device=self.device - ), - torch.tensor(self.nlist, dtype=torch.int64, device=self.device), - ] - ret3 = md0.forward_common_atomic(*args) - ret3 = cvt_ret(ret3) - ## model output on foo: [[2.8, 3.8, 5], [5.8, 7., 8.]] given bias [1.8, 2] - ## foo sumed: [11.6, 20.8] compared with [5, 7], fit target is [-6.6, -13.8] - ## fit bias is [-7, 2] (2 is assigned. -7 is fit to [-8.6, -17.8]) - ## old bias[1.8,2] + fit bias[-7, 2] = [-5.2, 4] - ## new model output is [[-4.2, -3.2, 7], [-1.2, 9, 10]] - expected_ret3 = {} - expected_ret3["foo"] = np.array([[-4.2, -3.2, 7.0], [-1.2, 9.0, 10.0]]).reshape( - 2, 3, 1 - ) - expected_ret3["pix"] = ret0["pix"] - for kk in ["foo", "pix"]: - np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk]) - # bar is too complicated to be manually computed. - - def test_preset_bias_all_none(self) -> None: - nf, nloc, nnei = self.nlist.shape - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ).to(self.device) - ft = FooFitting().to(self.device) - type_map = ["foo", "bar"] - preset_out_bias = { - "foo": [None, None], - } - md0 = DPAtomicModel( - ds, - ft, - type_map=type_map, - preset_out_bias=preset_out_bias, - ).to(self.device) - args = [ - torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device), - torch.tensor(self.atype_ext, dtype=torch.int64, device=self.device), - torch.tensor(self.nlist, dtype=torch.int64, device=self.device), - ] - # nf x nloc - at = self.atype_ext[:, :nloc] - - def cvt_ret(x): - return {kk: _to_numpy(vv) for kk, vv in x.items()} - - # 1. test run without bias - # nf x na x odim - ret0 = md0.forward_common_atomic(*args) - ret0 = cvt_ret(ret0) - expected_ret0 = {} - expected_ret0["foo"] = np.array( - [ - [1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["foo"].shape]) - expected_ret0["pix"] = np.array( - [ - [3.0, 2.0, 1.0], - [6.0, 5.0, 4.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["pix"].shape]) - expected_ret0["bar"] = np.array( - [ - [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], - [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["bar"].shape]) - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) - - # 2. test bias is applied (all None preset = same as no preset) - md0.compute_or_load_out_stat( - self.merged_output_stat, stat_file_path=self.stat_file_path - ) - ret1 = md0.forward_common_atomic(*args) - ret1 = cvt_ret(ret1) - # nt x odim - foo_bias = np.array([1.0, 3.0]).reshape(2, 1) - bar_bias = np.array([1.0, 5.0, 3.0, 2.0]).reshape(2, 1, 2) - expected_ret1 = {} - expected_ret1["foo"] = ret0["foo"] + foo_bias[at] - expected_ret1["pix"] = ret0["pix"] - expected_ret1["bar"] = ret0["bar"] + bar_bias[at] - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) - - def test_serialize(self) -> None: - nf, nloc, nnei = self.nlist.shape - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ).to(self.device) - ft = InvarFitting( - "foo", - self.nt, - ds.get_dim_out(), - 1, - mixed_types=ds.mixed_types(), - seed=GLOBAL_SEED, - ).to(self.device) - type_map = ["A", "B"] - md0 = DPAtomicModel( - ds, - ft, - type_map=type_map, - ).to(self.device) - args = [ - torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device), - torch.tensor(self.atype_ext, dtype=torch.int64, device=self.device), - torch.tensor(self.nlist, dtype=torch.int64, device=self.device), - ] - - def cvt_ret(x): - return {kk: _to_numpy(vv) for kk, vv in x.items()} - - md0.compute_or_load_out_stat( - self.merged_output_stat, stat_file_path=self.stat_file_path - ) - ret0 = md0.forward_common_atomic(*args) - ret0 = cvt_ret(ret0) - md1 = DPAtomicModel.deserialize(md0.serialize()) - ret1 = md1.forward_common_atomic(*args) - ret1 = cvt_ret(ret1) - - for kk in ["foo"]: - np.testing.assert_almost_equal(ret0[kk], ret1[kk]) - - md2 = DPDPAtomicModel.deserialize(md0.serialize()) - args_np = [self.coord_ext, self.atype_ext, self.nlist] - ret2 = md2.forward_common_atomic(*args_np) - for kk in ["foo"]: - np.testing.assert_almost_equal(ret0[kk], ret2[kk]) - - -class TestChangeByStatMixedLabels(unittest.TestCase, TestCaseSingleFrameWithNlist): - """Test change-by-statistic with mixed atomic and global labels.""" - - def tearDown(self) -> None: - self.tempdir.cleanup() - - def setUp(self) -> None: - TestCaseSingleFrameWithNlist.setUp(self) - self.device = env.DEVICE - self.merged_output_stat = [ - { - "coord": torch.tensor(np.zeros([2, 3, 3]), device=self.device), - "atype": torch.tensor( - np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), - device=self.device, - ), - "atype_ext": torch.tensor( - np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), - device=self.device, - ), - "box": torch.tensor(np.zeros([2, 3, 3]), device=self.device), - "natoms": torch.tensor( - np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), - device=self.device, - ), - # foo: atomic label - "atom_foo": torch.tensor( - np.array([[5.0, 5.0, 5.0], [5.0, 6.0, 7.0]]).reshape(2, 3, 1), - device=self.device, - ), - # pix: global label - "pix": torch.tensor( - np.array([5.0, 12.0]).reshape(2, 1), device=self.device - ), - # bar: global label - "bar": torch.tensor( - np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2), - device=self.device, - ), - "find_atom_foo": np.float32(1.0), - "find_pix": np.float32(1.0), - "find_bar": np.float32(1.0), - }, - ] - self.tempdir = tempfile.TemporaryDirectory() - h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) - with h5py.File(h5file, "w") as f: - pass - self.stat_file_path = DPPath(h5file, "a") - - def test_change_by_statistic(self) -> None: - """Test change-by-statistic with atomic foo + global pix + global bar.""" - nf, nloc, nnei = self.nlist.shape - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ).to(self.device) - ft = FooFitting().to(self.device) - type_map = ["foo", "bar"] - md0 = DPAtomicModel( - ds, - ft, - type_map=type_map, - ).to(self.device) - args = [ - torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device), - torch.tensor(self.atype_ext, dtype=torch.int64, device=self.device), - torch.tensor(self.nlist, dtype=torch.int64, device=self.device), - ] - - def cvt_ret(x): - return {kk: _to_numpy(vv) for kk, vv in x.items()} - - ret0 = md0.forward_common_atomic(*args) - ret0 = cvt_ret(ret0) - - # set initial bias - md0.compute_or_load_out_stat( - self.merged_output_stat, stat_file_path=self.stat_file_path - ) - - # change bias - md0.change_out_bias( - self.merged_output_stat, bias_adjust_mode="change-by-statistic" - ) - # use atype_ext from merged_output_stat for inference - args = [ - torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device), - self.merged_output_stat[0]["atype_ext"].to( - dtype=torch.int64, device=self.device - ), - torch.tensor(self.nlist, dtype=torch.int64, device=self.device), - ] - ret3 = md0.forward_common_atomic(*args) - ret3 = cvt_ret(ret3) - # foo: atomic label, bias after set-by-stat: [5, 6] - # model output with bias [5,6], atype [[0,0,1],[0,1,1]]: - # [[6, 7, 9], [9, 11, 12]] - # atom_foo labels: [[5, 5, 5], [5, 6, 7]] - # per-atom delta: [[-1, -2, -4], [-4, -5, -5]] - # delta bias (mean per type): type0=-7/3, type1=-14/3 - # new bias = [5-7/3, 6-14/3] = [8/3, 4/3] - # new output: [[11/3, 14/3, 13/3], [20/3, 19/3, 22/3]] - expected_ret3 = {} - expected_ret3["foo"] = np.array( - [[3.6667, 4.6667, 4.3333], [6.6667, 6.3333, 7.3333]] - ).reshape(2, 3, 1) - # pix: global label, bias after set-by-stat: [-2/3, 19/3] - # model pix with bias, atype [[0,0,1],[0,1,1]]: - # [[7/3, 4/3, 22/3], [16/3, 34/3, 31/3]], sums [11, 27] - # labels [5, 12], delta [-6, -15] - # lstsq: delta bias [1, -8], new bias [1/3, -5/3] - # new output: [[10/3, 7/3, -2/3], [19/3, 10/3, 7/3]] - expected_ret3["pix"] = np.array( - [[3.3333, 2.3333, -0.6667], [6.3333, 3.3333, 2.3333]] - ).reshape(2, 3, 1) - for kk in ["foo", "pix"]: - np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk], decimal=4) - # bar is too complicated to be manually computed. - - -class TestEnergyModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist): - """Test statistics computation with real energy fitting net.""" - - def tearDown(self) -> None: - self.tempdir.cleanup() - - def setUp(self) -> None: - TestCaseSingleFrameWithNlist.setUp(self) - self.device = env.DEVICE - self.merged_output_stat = [ - { - "coord": torch.tensor(np.zeros([2, 3, 3]), device=self.device), - "atype": torch.tensor( - np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), - device=self.device, - ), - "atype_ext": torch.tensor( - np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), - device=self.device, - ), - "box": torch.tensor(np.zeros([2, 3, 3]), device=self.device), - "natoms": torch.tensor( - np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), - device=self.device, - ), - # energy data - "energy": torch.tensor( - np.array([10.0, 20.0]).reshape(2, 1), device=self.device - ), - "find_energy": np.float32(1.0), - }, - ] - self.tempdir = tempfile.TemporaryDirectory() - h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) - with h5py.File(h5file, "w") as f: - pass - self.stat_file_path = DPPath(h5file, "a") - - def test_energy_stat(self) -> None: - """Test energy statistics computation with real energy fitting net.""" - nf, nloc, nnei = self.nlist.shape - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ).to(self.device) - ft = InvarFitting( - "energy", - self.nt, - ds.get_dim_out(), - 1, - mixed_types=ds.mixed_types(), - seed=GLOBAL_SEED, - ).to(self.device) - type_map = ["foo", "bar"] - md0 = DPAtomicModel( - ds, - ft, - type_map=type_map, - ).to(self.device) - args = [ - torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device), - torch.tensor(self.atype_ext, dtype=torch.int64, device=self.device), - torch.tensor(self.nlist, dtype=torch.int64, device=self.device), - ] - - # test run without bias - ret0 = md0.forward_common_atomic(*args) - self.assertIn("energy", ret0) - - # compute statistics - md0.compute_or_load_out_stat( - self.merged_output_stat, stat_file_path=self.stat_file_path - ) - ret1 = md0.forward_common_atomic(*args) - self.assertIn("energy", ret1) - - # Check that bias was computed (out_bias should be non-zero) - self.assertFalse(torch.all(md0.out_bias == 0)) - - # test bias load from file - def raise_error() -> NoReturn: - raise RuntimeError - - md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) - ret2 = md0.forward_common_atomic(*args) - np.testing.assert_allclose( - ret1["energy"].detach().cpu().numpy(), - ret2["energy"].detach().cpu().numpy(), - ) - - # test change bias - md0.change_out_bias( - self.merged_output_stat, bias_adjust_mode="change-by-statistic" - ) - ret3 = md0.forward_common_atomic(*args) - self.assertIn("energy", ret3) +# SPDX-License-Identifier: LGPL-3.0-or-later +import tempfile +import unittest +from pathlib import ( + Path, +) +from typing import ( + NoReturn, +) + +import h5py +import numpy as np + +from deepmd.dpmodel.atomic_model import DPAtomicModel as DPDPAtomicModel +from deepmd.dpmodel.common import ( + NativeOP, +) +from deepmd.dpmodel.descriptor import ( + DescrptSeA, +) +from deepmd.dpmodel.fitting import ( + InvarFitting, +) +from deepmd.dpmodel.fitting.base_fitting import ( + BaseFitting, +) +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + OutputVariableDef, +) +from deepmd.utils.path import ( + DPPath, +) + +from .case_single_frame_with_nlist import ( + TestCaseSingleFrameWithNlist, +) + + +class FooFitting(NativeOP, BaseFitting): + """Test fitting with multiple outputs for testing global statistics.""" + + def __init__(self): + pass + + def output_def(self): + return FittingOutputDef( + [ + OutputVariableDef( + "foo", + [1], + reducible=True, + r_differentiable=True, + c_differentiable=True, + ), + OutputVariableDef( + "pix", + [1], + reducible=True, + r_differentiable=True, + c_differentiable=True, + ), + OutputVariableDef( + "bar", + [1, 2], + reducible=True, + r_differentiable=True, + c_differentiable=True, + ), + ] + ) + + def serialize(self) -> dict: + return { + "@class": "Fitting", + "type": "foo", + "@version": 1, + } + + @classmethod + def deserialize(cls, data: dict): + return cls() + + def get_dim_fparam(self) -> int: + return 0 + + def get_dim_aparam(self) -> int: + return 0 + + def get_sel_type(self) -> list[int]: + return [] + + def change_type_map( + self, type_map: list[str], model_with_new_type_stat=None + ) -> None: + pass + + def get_type_map(self) -> list[str]: + return [] + + def call( + self, + descriptor, + atype, + gr=None, + g2=None, + h2=None, + fparam=None, + aparam=None, + ): + nf, nloc, _ = descriptor.shape + ret = {} + ret["foo"] = np.array( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + ] + ).reshape([nf, nloc, *self.output_def()["foo"].shape]) + ret["pix"] = np.array( + [ + [3.0, 2.0, 1.0], + [6.0, 5.0, 4.0], + ] + ).reshape([nf, nloc, *self.output_def()["pix"].shape]) + ret["bar"] = np.array( + [ + [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], + [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], + ] + ).reshape([nf, nloc, *self.output_def()["bar"].shape]) + return ret + + +class TestAtomicModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist): + def tearDown(self) -> None: + self.tempdir.cleanup() + + def setUp(self) -> None: + TestCaseSingleFrameWithNlist.setUp(self) + self.merged_output_stat = [ + { + "coord": np.zeros([2, 3, 3]), + "atype": np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), + "atype_ext": np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), + "box": np.zeros([2, 3, 3]), + "natoms": np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), + # bias of foo: 1, 3 + "foo": np.array([5.0, 7.0]).reshape(2, 1), + # no bias of pix + # bias of bar: [1, 5], [3, 2] + "bar": np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2), + "find_foo": np.float32(1.0), + "find_bar": np.float32(1.0), + } + ] + self.tempdir = tempfile.TemporaryDirectory() + h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) + with h5py.File(h5file, "w") as f: + pass + self.stat_file_path = DPPath(h5file, "a") + + def test_output_stat(self) -> None: + nf, nloc, nnei = self.nlist.shape + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = FooFitting() + type_map = ["foo", "bar"] + md0 = DPDPAtomicModel( + ds, + ft, + type_map=type_map, + ) + args = [self.coord_ext, self.atype_ext, self.nlist] + # nf x nloc + at = self.atype_ext[:, :nloc] + + # 1. test run without bias + # nf x na x odim + ret0 = md0.forward_common_atomic(*args) + + expected_ret0 = {} + expected_ret0["foo"] = np.array( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["foo"].shape]) + expected_ret0["pix"] = np.array( + [ + [3.0, 2.0, 1.0], + [6.0, 5.0, 4.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["pix"].shape]) + expected_ret0["bar"] = np.array( + [ + [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], + [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["bar"].shape]) + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) + + # 2. test bias is applied + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + ret1 = md0.forward_common_atomic(*args) + expected_std = np.ones((3, 2, 2)) # 3 keys, 2 atypes, 2 max dims. + # nt x odim + foo_bias = np.array([1.0, 3.0]).reshape(2, 1) + bar_bias = np.array([1.0, 5.0, 3.0, 2.0]).reshape(2, 1, 2) + expected_ret1 = {} + expected_ret1["foo"] = ret0["foo"] + foo_bias[at] + expected_ret1["pix"] = ret0["pix"] + expected_ret1["bar"] = ret0["bar"] + bar_bias[at] + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) + np.testing.assert_almost_equal(md0.out_std, expected_std) + + # 3. test bias load from file + def raise_error() -> NoReturn: + raise RuntimeError + + md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) + ret2 = md0.forward_common_atomic(*args) + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret1[kk], ret2[kk]) + np.testing.assert_almost_equal(md0.out_std, expected_std) + + # 4. test change bias + md0.change_out_bias( + self.merged_output_stat, bias_adjust_mode="change-by-statistic" + ) + # use atype_ext from merged_output_stat for inference + args = [ + self.coord_ext, + np.array(self.merged_output_stat[0]["atype_ext"], dtype=np.int64), + self.nlist, + ] + ret3 = md0.forward_common_atomic(*args) + ## model output on foo: [[2, 3, 6], [5, 8, 9]] given bias [1, 3] + ## foo sumed: [11, 22] compared with [5, 7], fit target is [-6, -15] + ## fit bias is [1, -8] + ## old bias + fit bias [2, -5] + ## new model output is [[3, 4, -2], [6, 0, 1]], which sumed to [5, 7] + expected_ret3 = {} + expected_ret3["foo"] = np.array([[3, 4, -2], [6, 0, 1]]).reshape(2, 3, 1) + expected_ret3["pix"] = ret0["pix"] + for kk in ["foo", "pix"]: + np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk]) + # bar is too complicated to be manually computed. + np.testing.assert_almost_equal(md0.out_std, expected_std) + + def test_preset_bias(self) -> None: + nf, nloc, nnei = self.nlist.shape + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = FooFitting() + type_map = ["foo", "bar"] + preset_out_bias = { + "foo": [None, 2], + "bar": np.array([7.0, 5.0, 13.0, 11.0]).reshape(2, 1, 2), + } + md0 = DPDPAtomicModel( + ds, + ft, + type_map=type_map, + preset_out_bias=preset_out_bias, + ) + args = [self.coord_ext, self.atype_ext, self.nlist] + # nf x nloc + at = self.atype_ext[:, :nloc] + + # 1. test run without bias + # nf x na x odim + ret0 = md0.forward_common_atomic(*args) + expected_ret0 = {} + expected_ret0["foo"] = np.array( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["foo"].shape]) + expected_ret0["pix"] = np.array( + [ + [3.0, 2.0, 1.0], + [6.0, 5.0, 4.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["pix"].shape]) + expected_ret0["bar"] = np.array( + [ + [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], + [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["bar"].shape]) + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) + + # 2. test bias is applied + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + ret1 = md0.forward_common_atomic(*args) + # foo sums: [5, 7], + # given bias of type 1 being 2, the bias left for type 0 is [5-2*1, 7-2*2] = [3,3] + # the solution of type 0 is 1.8 + foo_bias = np.array([1.8, preset_out_bias["foo"][1]]).reshape(2, 1) + bar_bias = preset_out_bias["bar"] + expected_ret1 = {} + expected_ret1["foo"] = ret0["foo"] + foo_bias[at] + expected_ret1["pix"] = ret0["pix"] + expected_ret1["bar"] = ret0["bar"] + bar_bias[at] + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) + + # 3. test bias load from file + def raise_error() -> NoReturn: + raise RuntimeError + + md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) + ret2 = md0.forward_common_atomic(*args) + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret1[kk], ret2[kk]) + + # 4. test change bias + md0.change_out_bias( + self.merged_output_stat, bias_adjust_mode="change-by-statistic" + ) + # use atype_ext from merged_output_stat for inference + args = [ + self.coord_ext, + np.array(self.merged_output_stat[0]["atype_ext"], dtype=np.int64), + self.nlist, + ] + ret3 = md0.forward_common_atomic(*args) + ## model output on foo: [[2.8, 3.8, 5], [5.8, 7., 8.]] given bias [1.8, 2] + ## foo sumed: [11.6, 20.8] compared with [5, 7], fit target is [-6.6, -13.8] + ## fit bias is [-7, 2] (2 is assigned. -7 is fit to [-8.6, -17.8]) + ## old bias[1.8,2] + fit bias[-7, 2] = [-5.2, 4] + ## new model output is [[-4.2, -3.2, 7], [-1.2, 9, 10]] + expected_ret3 = {} + expected_ret3["foo"] = np.array([[-4.2, -3.2, 7.0], [-1.2, 9.0, 10.0]]).reshape( + 2, 3, 1 + ) + expected_ret3["pix"] = ret0["pix"] + for kk in ["foo", "pix"]: + np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk]) + # bar is too complicated to be manually computed. + + def test_preset_bias_all_none(self) -> None: + nf, nloc, nnei = self.nlist.shape + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = FooFitting() + type_map = ["foo", "bar"] + preset_out_bias = { + "foo": [None, None], + } + md0 = DPDPAtomicModel( + ds, + ft, + type_map=type_map, + preset_out_bias=preset_out_bias, + ) + args = [self.coord_ext, self.atype_ext, self.nlist] + # nf x nloc + at = self.atype_ext[:, :nloc] + + # 1. test run without bias + # nf x na x odim + ret0 = md0.forward_common_atomic(*args) + expected_ret0 = {} + expected_ret0["foo"] = np.array( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["foo"].shape]) + expected_ret0["pix"] = np.array( + [ + [3.0, 2.0, 1.0], + [6.0, 5.0, 4.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["pix"].shape]) + expected_ret0["bar"] = np.array( + [ + [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], + [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["bar"].shape]) + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) + + # 2. test bias is applied (all None preset = same as no preset) + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + ret1 = md0.forward_common_atomic(*args) + # nt x odim + foo_bias = np.array([1.0, 3.0]).reshape(2, 1) + bar_bias = np.array([1.0, 5.0, 3.0, 2.0]).reshape(2, 1, 2) + expected_ret1 = {} + expected_ret1["foo"] = ret0["foo"] + foo_bias[at] + expected_ret1["pix"] = ret0["pix"] + expected_ret1["bar"] = ret0["bar"] + bar_bias[at] + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) + + def test_serialize(self) -> None: + nf, nloc, nnei = self.nlist.shape + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = InvarFitting( + "foo", + self.nt, + ds.get_dim_out(), + 1, + mixed_types=ds.mixed_types(), + ) + type_map = ["A", "B"] + md0 = DPDPAtomicModel( + ds, + ft, + type_map=type_map, + ) + args = [self.coord_ext, self.atype_ext, self.nlist] + + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + ret0 = md0.forward_common_atomic(*args) + md1 = DPDPAtomicModel.deserialize(md0.serialize()) + ret1 = md1.forward_common_atomic(*args) + + for kk in ["foo"]: + np.testing.assert_almost_equal(ret0[kk], ret1[kk]) + + +class TestChangeByStatMixedLabels(unittest.TestCase, TestCaseSingleFrameWithNlist): + """Test change-by-statistic with mixed atomic and global labels.""" + + def tearDown(self) -> None: + self.tempdir.cleanup() + + def setUp(self) -> None: + TestCaseSingleFrameWithNlist.setUp(self) + self.merged_output_stat = [ + { + "coord": np.zeros([2, 3, 3]), + "atype": np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), + "atype_ext": np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), + "box": np.zeros([2, 3, 3]), + "natoms": np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), + # foo: atomic label + "atom_foo": np.array([[5.0, 5.0, 5.0], [5.0, 6.0, 7.0]]).reshape( + 2, 3, 1 + ), + # pix: global label + "pix": np.array([5.0, 12.0]).reshape(2, 1), + # bar: global label + "bar": np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2), + "find_atom_foo": np.float32(1.0), + "find_pix": np.float32(1.0), + "find_bar": np.float32(1.0), + }, + ] + self.tempdir = tempfile.TemporaryDirectory() + h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) + with h5py.File(h5file, "w") as f: + pass + self.stat_file_path = DPPath(h5file, "a") + + def test_change_by_statistic(self) -> None: + """Test change-by-statistic with atomic foo + global pix + global bar.""" + nf, nloc, nnei = self.nlist.shape + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = FooFitting() + type_map = ["foo", "bar"] + md0 = DPDPAtomicModel( + ds, + ft, + type_map=type_map, + ) + args = [self.coord_ext, self.atype_ext, self.nlist] + + ret0 = md0.forward_common_atomic(*args) + + # set initial bias + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + + # change bias + md0.change_out_bias( + self.merged_output_stat, bias_adjust_mode="change-by-statistic" + ) + # use atype_ext from merged_output_stat for inference + args = [ + self.coord_ext, + np.array(self.merged_output_stat[0]["atype_ext"], dtype=np.int64), + self.nlist, + ] + ret3 = md0.forward_common_atomic(*args) + # foo: atomic label, bias after set-by-stat: [5, 6] + # model output with bias [5,6], atype [[0,0,1],[0,1,1]]: + # [[6, 7, 9], [9, 11, 12]] + # atom_foo labels: [[5, 5, 5], [5, 6, 7]] + # per-atom delta: [[-1, -2, -4], [-4, -5, -5]] + # delta bias (mean per type): type0=-7/3, type1=-14/3 + # new bias = [5-7/3, 6-14/3] = [8/3, 4/3] + # new output: [[11/3, 14/3, 13/3], [20/3, 19/3, 22/3]] + expected_ret3 = {} + expected_ret3["foo"] = np.array( + [[3.6667, 4.6667, 4.3333], [6.6667, 6.3333, 7.3333]] + ).reshape(2, 3, 1) + # pix: global label, bias after set-by-stat: [-2/3, 19/3] + # model pix with bias, atype [[0,0,1],[0,1,1]]: + # [[7/3, 4/3, 22/3], [16/3, 34/3, 31/3]], sums [11, 27] + # labels [5, 12], delta [-6, -15] + # lstsq: delta bias [1, -8], new bias [1/3, -5/3] + # new output: [[10/3, 7/3, -2/3], [19/3, 10/3, 7/3]] + expected_ret3["pix"] = np.array( + [[3.3333, 2.3333, -0.6667], [6.3333, 3.3333, 2.3333]] + ).reshape(2, 3, 1) + for kk in ["foo", "pix"]: + np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk], decimal=4) + # bar is too complicated to be manually computed. + + +class TestEnergyModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist): + """Test statistics computation with real energy fitting net.""" + + def tearDown(self) -> None: + self.tempdir.cleanup() + + def setUp(self) -> None: + TestCaseSingleFrameWithNlist.setUp(self) + self.merged_output_stat = [ + { + "coord": np.zeros([2, 3, 3]), + "atype": np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), + "atype_ext": np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), + "box": np.zeros([2, 3, 3]), + "natoms": np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), + # energy data + "energy": np.array([10.0, 20.0]).reshape(2, 1), + "find_energy": np.float32(1.0), + }, + ] + self.tempdir = tempfile.TemporaryDirectory() + h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) + with h5py.File(h5file, "w") as f: + pass + self.stat_file_path = DPPath(h5file, "a") + + def test_energy_stat(self) -> None: + """Test energy statistics computation with real energy fitting net.""" + nf, nloc, nnei = self.nlist.shape + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = InvarFitting( + "energy", + self.nt, + ds.get_dim_out(), + 1, + mixed_types=ds.mixed_types(), + ) + type_map = ["foo", "bar"] + md0 = DPDPAtomicModel( + ds, + ft, + type_map=type_map, + ) + args = [self.coord_ext, self.atype_ext, self.nlist] + + # test run without bias + ret0 = md0.forward_common_atomic(*args) + self.assertIn("energy", ret0) + + # compute statistics + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + ret1 = md0.forward_common_atomic(*args) + self.assertIn("energy", ret1) + + # Check that bias was computed (out_bias should be non-zero) + self.assertFalse(np.all(md0.out_bias == 0)) + + # test bias load from file + def raise_error() -> NoReturn: + raise RuntimeError + + md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) + ret2 = md0.forward_common_atomic(*args) + np.testing.assert_allclose( + ret1["energy"], + ret2["energy"], + ) + + # test change bias + md0.change_out_bias( + self.merged_output_stat, bias_adjust_mode="change-by-statistic" + ) + ret3 = md0.forward_common_atomic(*args) + self.assertIn("energy", ret3) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/common/dpmodel/test_dp_atomic_model.py b/source/tests/common/dpmodel/test_dp_atomic_model.py index 8ebd214865..ae5a9c610d 100644 --- a/source/tests/common/dpmodel/test_dp_atomic_model.py +++ b/source/tests/common/dpmodel/test_dp_atomic_model.py @@ -109,7 +109,7 @@ def test_excl_consistency(self) -> None: md0.reinit_pair_exclude(pair_excl) # hacking! md1.descriptor.reinit_exclude(pair_excl) - md1.fitting.reinit_exclude(atom_excl) + md1.fitting_net.reinit_exclude(atom_excl) # check energy consistency args = [self.coord_ext, self.atype_ext, self.nlist] diff --git a/source/tests/consistent/model/common.py b/source/tests/consistent/model/common.py index 04966c02a1..3dff24dcba 100644 --- a/source/tests/consistent/model/common.py +++ b/source/tests/consistent/model/common.py @@ -3,6 +3,8 @@ Any, ) +import numpy as np + from deepmd.common import ( make_default_mesh, ) @@ -145,3 +147,31 @@ def eval_pd_model(self, pd_obj: Any, natoms, coords, atype, box) -> Any: do_atomic_virial=True, ).items() } + + +def compare_variables_recursive( + d1: dict, d2: dict, path: str = "", rtol: float = 1e-10, atol: float = 1e-10 +) -> None: + """Recursively compare ``@variables`` sections in two serialized dicts.""" + for key in d1: + if key not in d2: + continue + child_path = f"{path}/{key}" if path else key + v1, v2 = d1[key], d2[key] + if key == "@variables" and isinstance(v1, dict) and isinstance(v2, dict): + for vk in v1: + if vk not in v2: + continue + a1 = np.asarray(v1[vk]) if v1[vk] is not None else None + a2 = np.asarray(v2[vk]) if v2[vk] is not None else None + if a1 is None and a2 is None: + continue + np.testing.assert_allclose( + a1, + a2, + rtol=rtol, + atol=atol, + err_msg=f"@variables mismatch at {child_path}/{vk}", + ) + elif isinstance(v1, dict) and isinstance(v2, dict): + compare_variables_recursive(v1, v2, child_path, rtol, atol) diff --git a/source/tests/consistent/model/test_dipole.py b/source/tests/consistent/model/test_dipole.py index bcd199a633..76251af6d2 100644 --- a/source/tests/consistent/model/test_dipole.py +++ b/source/tests/consistent/model/test_dipole.py @@ -6,8 +6,18 @@ import numpy as np +from deepmd.dpmodel.common import ( + to_numpy_array, +) from deepmd.dpmodel.model.dipole_model import DipoleModel as DipoleModelDP from deepmd.dpmodel.model.model import get_model as get_model_dp +from deepmd.dpmodel.utils.nlist import ( + build_neighbor_list, + extend_coord_with_ghosts, +) +from deepmd.dpmodel.utils.region import ( + normalize_coord, +) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) @@ -15,16 +25,21 @@ from ..common import ( INSTALLED_JAX, INSTALLED_PT, + INSTALLED_PT_EXPT, INSTALLED_TF, CommonTest, + parameterized, ) from .common import ( ModelTest, + compare_variables_recursive, ) if INSTALLED_PT: from deepmd.pt.model.model import get_model as get_model_pt from deepmd.pt.model.model.dipole_model import DipoleModel as DipoleModelPT + from deepmd.pt.utils.utils import to_numpy_array as torch_to_numpy + from deepmd.pt.utils.utils import to_torch_tensor as numpy_to_torch else: DipoleModelPT = None if INSTALLED_TF: @@ -36,6 +51,11 @@ from deepmd.jax.model.model import get_model as get_model_jax else: DipoleModelJAX = None +if INSTALLED_PT_EXPT: + from deepmd.pt_expt.common import to_torch_array as pt_expt_numpy_to_torch + from deepmd.pt_expt.model import DipoleModel as DipoleModelPTExpt +else: + DipoleModelPTExpt = None from deepmd.utils.argcheck import ( model_args, ) @@ -71,6 +91,7 @@ def data(self) -> dict: tf_class = DipoleModelTF dp_class = DipoleModelDP pt_class = DipoleModelPT + pt_expt_class = DipoleModelPTExpt jax_class = DipoleModelJAX args = model_args() atol = 1e-8 @@ -84,6 +105,8 @@ def get_reference_backend(self): return self.RefBackend.PT if not self.skip_tf: return self.RefBackend.TF + if not self.skip_pt_expt and self.pt_expt_class is not None: + return self.RefBackend.PT_EXPT if not self.skip_dp: return self.RefBackend.DP raise ValueError("No available reference") @@ -105,6 +128,9 @@ def pass_data_to_cls(self, cls, data) -> Any: model = get_model_pt(data) model.atomic_model.out_bias.uniform_() return model + elif cls is DipoleModelPTExpt: + dp_model = get_model_dp(data) + return DipoleModelPTExpt.deserialize(dp_model.serialize()) elif cls is DipoleModelJAX: return get_model_jax(data) return cls(**data, **self.additional_data) @@ -177,6 +203,15 @@ def eval_pt(self, pt_obj: Any) -> Any: self.box, ) + def eval_pt_expt(self, pt_expt_obj: Any) -> Any: + return self.eval_pt_expt_model( + pt_expt_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + def eval_jax(self, jax_obj: Any) -> Any: return self.eval_jax_model( jax_obj, @@ -196,6 +231,7 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: elif backend in { self.RefBackend.DP, self.RefBackend.PT, + self.RefBackend.PT_EXPT, self.RefBackend.JAX, }: return ( @@ -215,3 +251,1305 @@ def test_atom_exclude_types(self): tf_obj = self.tf_class.deserialize(data, suffix=self.unique_id) pt_obj = self.pt_class.deserialize(data) self.assertEqual(tf_obj.get_sel_type(), pt_obj.get_sel_type()) + + +@unittest.skipUnless(INSTALLED_PT and INSTALLED_PT_EXPT, "PyTorch is not installed") +class TestDipoleModelAPIs(unittest.TestCase): + """Test consistency of model-level APIs between pt and dpmodel backends. + + Both models are constructed from the same serialized weights + (dpmodel -> serialize -> pt deserialize) so that numerical outputs + can be compared directly. + """ + + def setUp(self) -> None: + data = model_args().normalize_value( + { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20], + "rcut_smth": 1.8, + "rcut": 6.0, + "neuron": [2, 4, 8], + "resnet_dt": False, + "axis_neuron": 8, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "type": "dipole", + "neuron": [4, 4, 4], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + "numb_fparam": 2, + "numb_aparam": 3, + "default_fparam": [0.5, -0.3], + }, + }, + trim_pattern="_*", + ) + # Build dpmodel first, then deserialize into pt/pt_expt to share weights + self.dp_model = get_model_dp(data) + serialized = self.dp_model.serialize() + self.pt_model = DipoleModelPT.deserialize(serialized) + self.pt_expt_model = DipoleModelPTExpt.deserialize(serialized) + + # Coords / atype / box + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, -1, 3) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32).reshape(1, -1) + self.box = np.array( + [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, 9) + + # Build extended coords + nlist for lower-level calls + rcut = 6.0 + nframes, nloc = self.atype.shape[:2] + coord_normalized = normalize_coord( + self.coords.reshape(nframes, nloc, 3), + self.box.reshape(nframes, 3, 3), + ) + extended_coord, extended_atype, mapping = extend_coord_with_ghosts( + coord_normalized, self.atype, self.box, rcut + ) + nlist = build_neighbor_list( + extended_coord, + extended_atype, + nloc, + rcut, + [20, 20], + distinguish_types=True, + ) + self.extended_coord = extended_coord.reshape(nframes, -1, 3) + self.extended_atype = extended_atype + self.mapping = mapping + self.nlist = nlist + + # aparam for forward evaluation (1 frame, 6 atoms, 3 aparam) + rng = np.random.default_rng(42) + self.eval_aparam = rng.normal(size=(1, nloc, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + def test_translated_output_def(self) -> None: + """translated_output_def should return the same keys on dp, pt, and pt_expt.""" + dp_def = self.dp_model.translated_output_def() + pt_def = self.pt_model.translated_output_def() + pt_expt_def = self.pt_expt_model.translated_output_def() + self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + self.assertEqual(set(dp_def.keys()), set(pt_expt_def.keys())) + for key in dp_def: + self.assertEqual(dp_def[key].shape, pt_def[key].shape) + self.assertEqual(dp_def[key].shape, pt_expt_def[key].shape) + + def test_get_descriptor(self) -> None: + """get_descriptor should return a non-None object on both backends.""" + self.assertIsNotNone(self.dp_model.get_descriptor()) + self.assertIsNotNone(self.pt_model.get_descriptor()) + + def test_get_fitting_net(self) -> None: + """get_fitting_net should return a non-None object on both backends.""" + self.assertIsNotNone(self.dp_model.get_fitting_net()) + self.assertIsNotNone(self.pt_model.get_fitting_net()) + + def test_get_out_bias(self) -> None: + """get_out_bias should return numerically equal values on dp and pt. + + DipoleModel's apply_out_stat is a no-op, so bias won't affect output, + but the bias storage should still be consistent across backends. + """ + dp_bias = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias = torch_to_numpy(self.pt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias, pt_bias, rtol=1e-10, atol=1e-10) + # Verify shape: (n_output_keys x ntypes x 3) for dipole + self.assertEqual(dp_bias.shape[1], 2) # ntypes + self.assertGreater(dp_bias.shape[0], 0) # at least one output key + + def test_set_out_bias(self) -> None: + """set_out_bias should update the bias on both backends.""" + dp_bias = to_numpy_array(self.dp_model.get_out_bias()) + new_bias = dp_bias + 1.0 + # dp + self.dp_model.set_out_bias(new_bias) + np.testing.assert_allclose( + to_numpy_array(self.dp_model.get_out_bias()), + new_bias, + rtol=1e-10, + atol=1e-10, + ) + # pt + self.pt_model.set_out_bias(numpy_to_torch(new_bias)) + np.testing.assert_allclose( + torch_to_numpy(self.pt_model.get_out_bias()), + new_bias, + rtol=1e-10, + atol=1e-10, + ) + + def test_model_output_def(self) -> None: + """model_output_def should return the same keys and shapes on dp and pt.""" + dp_def = self.dp_model.model_output_def().get_data() + pt_def = self.pt_model.model_output_def().get_data() + self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + for key in dp_def: + self.assertEqual(dp_def[key].shape, pt_def[key].shape) + + def test_model_output_type(self) -> None: + """model_output_type should return the same list on dp and pt.""" + self.assertEqual( + self.dp_model.model_output_type(), + self.pt_model.model_output_type(), + ) + + def test_do_grad_r(self) -> None: + """do_grad_r should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.do_grad_r("dipole"), + self.pt_model.do_grad_r("dipole"), + ) + self.assertTrue(self.dp_model.do_grad_r("dipole")) + + def test_do_grad_c(self) -> None: + """do_grad_c should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.do_grad_c("dipole"), + self.pt_model.do_grad_c("dipole"), + ) + self.assertTrue(self.dp_model.do_grad_c("dipole")) + + def test_get_rcut(self) -> None: + """get_rcut should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_rcut(), self.pt_model.get_rcut()) + self.assertAlmostEqual(self.dp_model.get_rcut(), 6.0) + + def test_get_type_map(self) -> None: + """get_type_map should return the same list on dp and pt.""" + self.assertEqual(self.dp_model.get_type_map(), self.pt_model.get_type_map()) + self.assertEqual(self.dp_model.get_type_map(), ["O", "H"]) + + def test_get_sel(self) -> None: + """get_sel should return the same list on dp and pt.""" + self.assertEqual(self.dp_model.get_sel(), self.pt_model.get_sel()) + self.assertEqual(self.dp_model.get_sel(), [20, 20]) + + def test_get_nsel(self) -> None: + """get_nsel should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_nsel(), self.pt_model.get_nsel()) + self.assertEqual(self.dp_model.get_nsel(), 40) + + def test_get_nnei(self) -> None: + """get_nnei should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_nnei(), self.pt_model.get_nnei()) + self.assertEqual(self.dp_model.get_nnei(), 40) + + def test_mixed_types(self) -> None: + """mixed_types should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.mixed_types(), self.pt_model.mixed_types()) + # se_e2_a is not mixed-types + self.assertFalse(self.dp_model.mixed_types()) + + def test_has_message_passing(self) -> None: + """has_message_passing should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.has_message_passing(), + self.pt_model.has_message_passing(), + ) + self.assertFalse(self.dp_model.has_message_passing()) + + def test_need_sorted_nlist_for_lower(self) -> None: + """need_sorted_nlist_for_lower should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.need_sorted_nlist_for_lower(), + self.pt_model.need_sorted_nlist_for_lower(), + ) + self.assertFalse(self.dp_model.need_sorted_nlist_for_lower()) + + def test_get_dim_fparam(self) -> None: + """get_dim_fparam should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_dim_fparam(), self.pt_model.get_dim_fparam()) + self.assertEqual(self.dp_model.get_dim_fparam(), 2) + + def test_get_dim_aparam(self) -> None: + """get_dim_aparam should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_dim_aparam(), self.pt_model.get_dim_aparam()) + self.assertEqual(self.dp_model.get_dim_aparam(), 3) + + def test_get_sel_type(self) -> None: + """get_sel_type should return the same list on dp and pt.""" + self.assertEqual(self.dp_model.get_sel_type(), self.pt_model.get_sel_type()) + self.assertEqual(self.dp_model.get_sel_type(), [0, 1]) + + def test_is_aparam_nall(self) -> None: + """is_aparam_nall should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.is_aparam_nall(), self.pt_model.is_aparam_nall()) + self.assertFalse(self.dp_model.is_aparam_nall()) + + def test_get_model_def_script(self) -> None: + """get_model_def_script should return the same value on dp, pt, and pt_expt.""" + dp_val = self.dp_model.get_model_def_script() + pt_val = self.pt_model.get_model_def_script() + pe_val = self.pt_expt_model.get_model_def_script() + self.assertEqual(dp_val, pt_val) + self.assertEqual(dp_val, pe_val) + + def test_get_min_nbor_dist(self) -> None: + """get_min_nbor_dist should return the same value on dp, pt, and pt_expt.""" + dp_val = self.dp_model.get_min_nbor_dist() + pt_val = self.pt_model.get_min_nbor_dist() + pe_val = self.pt_expt_model.get_min_nbor_dist() + self.assertEqual(dp_val, pt_val) + self.assertEqual(dp_val, pe_val) + + def test_set_case_embd(self) -> None: + """set_case_embd should produce consistent results across backends. + + Also verifies that different case indices produce different outputs, + confirming the embedding is actually used. + """ + from deepmd.utils.argcheck import ( + model_args, + ) + + # Build a model with dim_case_embd > 0 + data = model_args().normalize_value( + { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "type": "dipole", + "neuron": [5, 5], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + "dim_case_embd": 3, + }, + }, + trim_pattern="_*", + ) + dp_model = get_model_dp(data) + serialized = dp_model.serialize() + pt_model = DipoleModelPT.deserialize(serialized) + pe_model = DipoleModelPTExpt.deserialize(serialized) + + def _eval(case_idx): + dp_model.set_case_embd(case_idx) + pt_model.set_case_embd(case_idx) + pe_model.set_case_embd(case_idx) + dp_ret = dp_model(self.coords, self.atype, box=self.box) + pt_ret = { + k: torch_to_numpy(v) + for k, v in pt_model( + numpy_to_torch(self.coords), + numpy_to_torch(self.atype), + box=numpy_to_torch(self.box), + ).items() + } + coord_t = pt_expt_numpy_to_torch(self.coords) + coord_t.requires_grad_(True) + pe_ret = { + k: v.detach().cpu().numpy() + for k, v in pe_model( + coord_t, + pt_expt_numpy_to_torch(self.atype), + box=pt_expt_numpy_to_torch(self.box), + ).items() + } + return dp_ret, pt_ret, pe_ret + + dp0, pt0, pe0 = _eval(0) + dp1, pt1, pe1 = _eval(1) + + # Cross-backend consistency for each case index + for key in ("dipole", "global_dipole"): + np.testing.assert_allclose( + dp0[key], + pt0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 0: dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp0[key], + pe0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 0: dp vs pt_expt mismatch in {key}", + ) + np.testing.assert_allclose( + dp1[key], + pt1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 1: dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp1[key], + pe1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 1: dp vs pt_expt mismatch in {key}", + ) + # Different case indices should produce different outputs + self.assertFalse( + np.allclose(dp0["global_dipole"], dp1["global_dipole"]), + "set_case_embd(0) and set_case_embd(1) produced the same global_dipole", + ) + + def test_atomic_output_def(self) -> None: + """atomic_output_def should return the same keys and shapes on dp and pt.""" + dp_def = self.dp_model.atomic_output_def() + pt_def = self.pt_model.atomic_output_def() + self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + for key in dp_def.keys(): + self.assertEqual(dp_def[key].shape, pt_def[key].shape) + + def test_format_nlist(self) -> None: + """format_nlist should produce the same result on dp and pt.""" + dp_nlist = self.dp_model.format_nlist( + self.extended_coord, + self.extended_atype, + self.nlist, + ) + pt_nlist = torch_to_numpy( + self.pt_model.format_nlist( + numpy_to_torch(self.extended_coord), + numpy_to_torch(self.extended_atype), + numpy_to_torch(self.nlist), + ) + ) + np.testing.assert_equal(dp_nlist, pt_nlist) + + def test_forward_common_atomic(self) -> None: + """forward_common_atomic should produce consistent results on dp and pt. + + Compares at the atomic_model level, where both backends define this method. + """ + dp_ret = self.dp_model.atomic_model.forward_common_atomic( + self.extended_coord, + self.extended_atype, + self.nlist, + mapping=self.mapping, + aparam=self.eval_aparam, + ) + pt_ret = self.pt_model.atomic_model.forward_common_atomic( + numpy_to_torch(self.extended_coord), + numpy_to_torch(self.extended_atype), + numpy_to_torch(self.nlist), + mapping=numpy_to_torch(self.mapping), + aparam=numpy_to_torch(self.eval_aparam), + ) + # Compare the common keys + common_keys = set(dp_ret.keys()) & set(pt_ret.keys()) + self.assertTrue(len(common_keys) > 0) + for key in common_keys: + if dp_ret[key] is not None and pt_ret[key] is not None: + np.testing.assert_allclose( + dp_ret[key], + torch_to_numpy(pt_ret[key]), + rtol=1e-10, + atol=1e-10, + err_msg=f"Mismatch in forward_common_atomic key '{key}'", + ) + + def test_has_default_fparam(self) -> None: + """has_default_fparam should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.has_default_fparam(), + self.pt_model.has_default_fparam(), + ) + self.assertTrue(self.dp_model.has_default_fparam()) + + def test_get_default_fparam(self) -> None: + """get_default_fparam should return consistent values on dp and pt.""" + dp_val = self.dp_model.get_default_fparam() + pt_val = self.pt_model.get_default_fparam() + np.testing.assert_allclose(dp_val, pt_val, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dp_val, [0.5, -0.3], rtol=1e-10, atol=1e-10) + + def test_change_out_bias(self) -> None: + """change_out_bias should produce consistent bias on dp, pt, and pt_expt. + + Tests both set-by-statistic and change-by-statistic modes. + Note: change_out_bias only updates the output bias, not fitting input + stats (fparam/aparam). Fitting stats are updated by compute_or_load_stat. + """ + nframes = 2 + nloc = 6 + numb_fparam = 2 + numb_aparam = 3 + rng = np.random.default_rng(123) + + # Use realistic coords (from setUp, tiled for 2 frames) + coords_2f = np.tile(self.coords, (nframes, 1, 1)) # (2, 6, 3) + atype_2f = np.array([[0, 0, 1, 1, 1, 1], [0, 1, 1, 0, 1, 1]], dtype=np.int32) + box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) + natoms_data = np.array([[6, 6, 2, 4], [6, 6, 2, 4]], dtype=np.int32) + dipole_data = rng.normal(size=(nframes, 3)).astype(GLOBAL_NP_FLOAT_PRECISION) + fparam_data = rng.normal(size=(nframes, numb_fparam)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + aparam_data = rng.normal(size=(nframes, nloc, numb_aparam)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + # dpmodel stat data (numpy) + dp_merged = [ + { + "coord": coords_2f, + "atype": atype_2f, + "atype_ext": atype_2f, + "box": box_2f, + "natoms": natoms_data, + "dipole": dipole_data, + "find_dipole": np.float32(1.0), + "fparam": fparam_data, + "aparam": aparam_data, + } + ] + # pt stat data (torch tensors) + pt_merged = [ + { + "coord": numpy_to_torch(coords_2f), + "atype": numpy_to_torch(atype_2f), + "atype_ext": numpy_to_torch(atype_2f), + "box": numpy_to_torch(box_2f), + "natoms": numpy_to_torch(natoms_data), + "dipole": numpy_to_torch(dipole_data), + "find_dipole": np.float32(1.0), + "fparam": numpy_to_torch(fparam_data), + "aparam": numpy_to_torch(aparam_data), + } + ] + # pt_expt stat data (numpy, same as dp) + pe_merged = dp_merged + + # Save initial (zero) bias + dp_bias_init = to_numpy_array(self.dp_model.get_out_bias()).copy() + + # --- Test "set-by-statistic" mode --- + self.dp_model.change_out_bias(dp_merged, bias_adjust_mode="set-by-statistic") + self.pt_model.change_out_bias(pt_merged, bias_adjust_mode="set-by-statistic") + self.pt_expt_model.change_out_bias( + pe_merged, bias_adjust_mode="set-by-statistic" + ) + + # Verify out bias consistency + dp_bias = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias = torch_to_numpy(self.pt_model.get_out_bias()) + pe_bias = to_numpy_array(self.pt_expt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias, pt_bias, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dp_bias, pe_bias, rtol=1e-10, atol=1e-10) + self.assertFalse( + np.allclose(dp_bias, dp_bias_init), + "set-by-statistic did not change the bias from initial values", + ) + + # --- Test "change-by-statistic" mode --- + dp_bias_before = dp_bias.copy() + self.dp_model.change_out_bias(dp_merged, bias_adjust_mode="change-by-statistic") + self.pt_model.change_out_bias(pt_merged, bias_adjust_mode="change-by-statistic") + self.pt_expt_model.change_out_bias( + pe_merged, bias_adjust_mode="change-by-statistic" + ) + + # Verify out bias consistency + dp_bias2 = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias2 = torch_to_numpy(self.pt_model.get_out_bias()) + pe_bias2 = to_numpy_array(self.pt_expt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias2, pt_bias2, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dp_bias2, pe_bias2, rtol=1e-10, atol=1e-10) + self.assertFalse( + np.allclose(dp_bias2, dp_bias_before), + "change-by-statistic did not further change the bias", + ) + + def test_change_type_map(self) -> None: + """change_type_map should produce consistent results on dp and pt. + + Uses a DPA1 (se_atten) descriptor since se_e2_a does not support + change_type_map (non-mixed-types descriptors raise NotImplementedError). + """ + from deepmd.utils.argcheck import model_args as model_args_fn + + data = model_args_fn().normalize_value( + { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_atten", + "sel": 20, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "seed": 1, + "attn": 6, + "attn_layer": 0, + }, + "fitting_net": { + "type": "dipole", + "neuron": [4, 4, 4], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + }, + }, + trim_pattern="_*", + ) + dp_model = get_model_dp(data) + pt_model = DipoleModelPT.deserialize(dp_model.serialize()) + + # Set non-zero out_bias so the swap is non-trivial + dp_bias_orig = to_numpy_array(dp_model.get_out_bias()).copy() + new_bias = dp_bias_orig.copy() + new_bias[:, 0, :] = 1.5 # type 0 ("O") + new_bias[:, 1, :] = -3.7 # type 1 ("H") + dp_model.set_out_bias(new_bias) + pt_model.set_out_bias(numpy_to_torch(new_bias)) + + new_type_map = ["H", "O"] + dp_model.change_type_map(new_type_map) + pt_model.change_type_map(new_type_map) + + # Both should have the new type_map + self.assertEqual(dp_model.get_type_map(), new_type_map) + self.assertEqual(pt_model.get_type_map(), new_type_map) + + # Out_bias should be reordered consistently between backends + dp_bias_new = to_numpy_array(dp_model.get_out_bias()) + pt_bias_new = torch_to_numpy(pt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias_new, pt_bias_new, rtol=1e-10, atol=1e-10) + + # Verify the reorder is correct: old type 0 -> new type 1, old type 1 -> new type 0 + np.testing.assert_allclose( + dp_bias_new[:, 0, :], + new_bias[:, 1, :], + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + dp_bias_new[:, 1, :], + new_bias[:, 0, :], + rtol=1e-10, + atol=1e-10, + ) + + def test_change_type_map_extend_stat(self) -> None: + """change_type_map with model_with_new_type_stat should propagate stats consistently across dp, pt, and pt_expt. + + Verifies that the model-level change_type_map correctly unwraps + model_with_new_type_stat.atomic_model before forwarding to the + atomic model. + """ + from deepmd.utils.argcheck import model_args as model_args_fn + + small_tm = ["O", "H"] + large_tm = ["O", "H", "Li"] + + small_data = model_args_fn().normalize_value( + { + "type_map": small_tm, + "descriptor": { + "type": "se_atten", + "sel": 20, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "seed": 1, + "attn": 6, + "attn_layer": 0, + }, + "fitting_net": { + "type": "dipole", + "neuron": [4, 4, 4], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + }, + }, + trim_pattern="_*", + ) + large_data = model_args_fn().normalize_value( + { + "type_map": large_tm, + "descriptor": { + "type": "se_atten", + "sel": 20, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "seed": 2, + "attn": 6, + "attn_layer": 0, + }, + "fitting_net": { + "type": "dipole", + "neuron": [4, 4, 4], + "resnet_dt": True, + "precision": "float64", + "seed": 2, + }, + }, + trim_pattern="_*", + ) + + dp_small = get_model_dp(small_data) + dp_large = get_model_dp(large_data) + + # Set distinguishable random stats on the large model's descriptor + rng = np.random.default_rng(42) + desc_large = dp_large.get_descriptor() + mean_large, std_large = desc_large.get_stat_mean_and_stddev() + mean_rand = rng.random(size=to_numpy_array(mean_large).shape) + std_rand = rng.random(size=to_numpy_array(std_large).shape) + desc_large.set_stat_mean_and_stddev(mean_rand, std_rand) + + # Build pt and pt_expt models from dp serialization + pt_small = DipoleModelPT.deserialize(dp_small.serialize()) + pt_large = DipoleModelPT.deserialize(dp_large.serialize()) + pt_expt_small = DipoleModelPTExpt.deserialize(dp_small.serialize()) + pt_expt_large = DipoleModelPTExpt.deserialize(dp_large.serialize()) + + # Extend type map with model_with_new_type_stat at the model level + dp_small.change_type_map(large_tm, model_with_new_type_stat=dp_large) + pt_small.change_type_map(large_tm, model_with_new_type_stat=pt_large) + pt_expt_small.change_type_map(large_tm, model_with_new_type_stat=pt_expt_large) + + # Descriptor stats should be consistent across backends + dp_mean, dp_std = dp_small.get_descriptor().get_stat_mean_and_stddev() + pt_mean, pt_std = pt_small.get_descriptor().get_stat_mean_and_stddev() + pt_expt_mean, pt_expt_std = ( + pt_expt_small.get_descriptor().get_stat_mean_and_stddev() + ) + np.testing.assert_allclose( + to_numpy_array(dp_mean), + torch_to_numpy(pt_mean), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + to_numpy_array(dp_std), + torch_to_numpy(pt_std), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + to_numpy_array(dp_mean), + to_numpy_array(pt_expt_mean), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + to_numpy_array(dp_std), + to_numpy_array(pt_expt_std), + rtol=1e-10, + atol=1e-10, + ) + + def test_update_sel(self) -> None: + """update_sel should return the same result on dp and pt.""" + from unittest.mock import ( + patch, + ) + + from deepmd.dpmodel.model.dp_model import DPModelCommon as DPModelCommonDP + from deepmd.pt.model.model.dp_model import DPModelCommon as DPModelCommonPT + + mock_min_nbor_dist = 0.5 + mock_sel = [10, 20] + local_jdata = { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": "auto", + "rcut_smth": 0.50, + "rcut": 6.00, + }, + "fitting_net": { + "type": "dipole", + "neuron": [4, 4, 4], + }, + } + type_map = ["O", "H"] + + with patch( + "deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat", + return_value=(mock_min_nbor_dist, mock_sel), + ): + dp_result, dp_min_dist = DPModelCommonDP.update_sel( + None, type_map, local_jdata + ) + + with patch( + "deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat", + return_value=(mock_min_nbor_dist, mock_sel), + ): + pt_result, pt_min_dist = DPModelCommonPT.update_sel( + None, type_map, local_jdata + ) + + self.assertEqual(dp_result, pt_result) + self.assertEqual(dp_min_dist, pt_min_dist) + # Verify sel was actually updated (not still "auto") + self.assertIsInstance(dp_result["descriptor"]["sel"], list) + self.assertNotEqual(dp_result["descriptor"]["sel"], "auto") + + def test_get_ntypes(self) -> None: + """get_ntypes should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_ntypes(), self.pt_model.get_ntypes()) + self.assertEqual(self.dp_model.get_ntypes(), 2) + + def test_compute_or_load_out_stat(self) -> None: + """compute_or_load_out_stat should produce consistent bias on dp and pt. + + Tests both the compute path (from data) and the load path (from file). + Note: DipoleModel's apply_out_stat is a no-op, so bias doesn't affect + output, but the stored bias should still be consistent. + """ + import tempfile + from pathlib import ( + Path, + ) + + import h5py + + from deepmd.utils.path import ( + DPPath, + ) + + nframes = 2 + coords_2f = np.tile(self.coords, (nframes, 1, 1)) + atype_2f = np.array([[0, 0, 1, 1, 1, 1], [0, 1, 1, 0, 1, 1]], dtype=np.int32) + box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) + natoms_data = np.array([[6, 6, 2, 4], [6, 6, 2, 4]], dtype=np.int32) + dipole_data = ( + np.random.default_rng(42) + .normal(size=(nframes, 3)) + .astype(GLOBAL_NP_FLOAT_PRECISION) + ) + + dp_merged = [ + { + "coord": coords_2f, + "atype": atype_2f, + "atype_ext": atype_2f, + "box": box_2f, + "natoms": natoms_data, + "dipole": dipole_data, + "find_dipole": np.float32(1.0), + } + ] + pt_merged = [ + { + "coord": numpy_to_torch(coords_2f), + "atype": numpy_to_torch(atype_2f), + "atype_ext": numpy_to_torch(atype_2f), + "box": numpy_to_torch(box_2f), + "natoms": numpy_to_torch(natoms_data), + "dipole": numpy_to_torch(dipole_data), + "find_dipole": np.float32(1.0), + } + ] + + # Verify bias is initially identical + dp_bias_before = to_numpy_array(self.dp_model.get_out_bias()).copy() + pt_bias_before = torch_to_numpy(self.pt_model.get_out_bias()).copy() + np.testing.assert_allclose( + dp_bias_before, pt_bias_before, rtol=1e-10, atol=1e-10 + ) + + with tempfile.TemporaryDirectory() as tmpdir: + # Create separate h5 files for dp and pt + dp_h5 = str((Path(tmpdir) / "dp_stat.h5").resolve()) + pt_h5 = str((Path(tmpdir) / "pt_stat.h5").resolve()) + with h5py.File(dp_h5, "w"): + pass + with h5py.File(pt_h5, "w"): + pass + dp_stat_path = DPPath(dp_h5, "a") + pt_stat_path = DPPath(pt_h5, "a") + + # 1. Compute stats and save to file + self.dp_model.atomic_model.compute_or_load_out_stat( + dp_merged, stat_file_path=dp_stat_path + ) + self.pt_model.atomic_model.compute_or_load_out_stat( + pt_merged, stat_file_path=pt_stat_path + ) + + dp_bias_after = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias_after = torch_to_numpy(self.pt_model.get_out_bias()) + np.testing.assert_allclose( + dp_bias_after, pt_bias_after, rtol=1e-10, atol=1e-10 + ) + + # 2. Verify both backends saved the same file content + with h5py.File(dp_h5, "r") as dp_f, h5py.File(pt_h5, "r") as pt_f: + dp_keys = sorted(dp_f.keys()) + pt_keys = sorted(pt_f.keys()) + self.assertEqual(dp_keys, pt_keys) + for key in dp_keys: + np.testing.assert_allclose( + np.array(dp_f[key]), + np.array(pt_f[key]), + rtol=1e-10, + atol=1e-10, + err_msg=f"Stat file content mismatch for key {key}", + ) + + # 3. Reset biases to zero, then load from file + zero_bias = np.zeros_like(dp_bias_after) + self.dp_model.set_out_bias(zero_bias) + self.pt_model.set_out_bias(numpy_to_torch(zero_bias)) + + # Use a callable that raises to ensure it loads from file, not recomputes + def raise_error(): + raise RuntimeError("Should not recompute — should load from file") + + self.dp_model.atomic_model.compute_or_load_out_stat( + raise_error, stat_file_path=dp_stat_path + ) + self.pt_model.atomic_model.compute_or_load_out_stat( + raise_error, stat_file_path=pt_stat_path + ) + + dp_bias_loaded = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias_loaded = torch_to_numpy(self.pt_model.get_out_bias()) + + # Loaded biases should match between backends + np.testing.assert_allclose( + dp_bias_loaded, pt_bias_loaded, rtol=1e-10, atol=1e-10 + ) + # Loaded biases should match the originally computed biases + np.testing.assert_allclose( + dp_bias_loaded, dp_bias_after, rtol=1e-10, atol=1e-10 + ) + + def test_get_observed_type_list(self) -> None: + """get_observed_type_list should be consistent across dp, pt, pt_expt. + + Uses mock data containing only type 0 ("O") so that type 1 ("H") is + unobserved and should be absent from the returned list. + """ + nframes = 2 + natoms = 6 + # All atoms are type 0 — type 1 is unobserved + atype_2f = np.zeros((nframes, natoms), dtype=np.int32) + coords_2f = np.tile(self.coords, (nframes, 1, 1)) + box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) + natoms_data = np.array([[natoms, natoms, natoms, 0]] * nframes, dtype=np.int32) + dipole_data = ( + np.random.default_rng(42) + .normal(size=(nframes, 3)) + .astype(GLOBAL_NP_FLOAT_PRECISION) + ) + + dp_merged = [ + { + "coord": coords_2f, + "atype": atype_2f, + "atype_ext": atype_2f, + "box": box_2f, + "natoms": natoms_data, + "dipole": dipole_data, + "find_dipole": np.float32(1.0), + } + ] + pt_merged = [ + { + "coord": numpy_to_torch(coords_2f), + "atype": numpy_to_torch(atype_2f), + "atype_ext": numpy_to_torch(atype_2f), + "box": numpy_to_torch(box_2f), + "natoms": numpy_to_torch(natoms_data), + "dipole": numpy_to_torch(dipole_data), + "find_dipole": np.float32(1.0), + } + ] + + self.dp_model.atomic_model.compute_or_load_out_stat(dp_merged) + self.pt_model.atomic_model.compute_or_load_out_stat(pt_merged) + self.pt_expt_model.atomic_model.compute_or_load_out_stat(dp_merged) + + dp_observed = self.dp_model.get_observed_type_list() + pt_observed = self.pt_model.get_observed_type_list() + pe_observed = self.pt_expt_model.get_observed_type_list() + + self.assertEqual(dp_observed, pt_observed) + self.assertEqual(dp_observed, pe_observed) + # Only type 0 ("O") should be observed + self.assertEqual(dp_observed, ["O"]) + + +@parameterized( + (([], []), ([[0, 1]], [1])), # (pair_exclude_types, atom_exclude_types) + (False, True), # fparam_in_data +) +@unittest.skipUnless(INSTALLED_PT and INSTALLED_PT_EXPT, "PT and PT_EXPT are required") +class TestDipoleComputeOrLoadStat(unittest.TestCase): + """Test that compute_or_load_stat produces identical statistics on dp, pt, and pt_expt. + + Covers descriptor stats (dstd), fitting stats (fparam, aparam), and output bias. + Parameterized over exclusion types and whether fparam is explicitly provided or + injected via default_fparam. + """ + + def setUp(self) -> None: + (pair_exclude_types, atom_exclude_types), self.fparam_in_data = self.param + data = model_args().normalize_value( + { + "type_map": ["O", "H"], + "pair_exclude_types": pair_exclude_types, + "atom_exclude_types": atom_exclude_types, + "descriptor": { + "type": "dpa3", + "repflow": { + "n_dim": 20, + "e_dim": 10, + "a_dim": 8, + "nlayers": 3, + "e_rcut": 6.0, + "e_rcut_smth": 5.0, + "e_sel": 10, + "a_rcut": 4.0, + "a_rcut_smth": 3.5, + "a_sel": 8, + "axis_neuron": 4, + "update_angle": True, + "update_style": "res_residual", + "update_residual": 0.1, + "update_residual_init": "const", + }, + "precision": "float64", + "seed": 1, + }, + "fitting_net": { + "type": "dipole", + "neuron": [10, 10], + "precision": "float64", + "seed": 1, + "numb_fparam": 2, + "default_fparam": [0.5, -0.3], + "numb_aparam": 3, + }, + }, + trim_pattern="_*", + ) + + # Save data for reuse in load-from-file test + self._model_data = data + + # Build dp model, then deserialize into pt and pt_expt to share weights + self.dp_model = get_model_dp(data) + serialized = self.dp_model.serialize() + self.pt_model = DipoleModelPT.deserialize(serialized) + self.pt_expt_model = DipoleModelPTExpt.deserialize(serialized) + + # Test coords / atype / box for forward evaluation + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 0.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, -1, 3) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32).reshape(1, -1) + self.box = np.array( + [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, 9) + + # Mock training data for compute_or_load_stat + natoms = 6 + nframes = 3 + rng = np.random.default_rng(42) + coords_stat = rng.normal(size=(nframes, natoms, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + atype_stat = np.array([[0, 0, 1, 1, 1, 1]] * nframes, dtype=np.int32) + box_stat = np.tile( + np.eye(3, dtype=GLOBAL_NP_FLOAT_PRECISION).reshape(1, 3, 3) * 13.0, + (nframes, 1, 1), + ) + natoms_stat = np.array([[natoms, natoms, 2, 4]] * nframes, dtype=np.int32) + dipole_stat = rng.normal(size=(nframes, 3)).astype(GLOBAL_NP_FLOAT_PRECISION) + aparam_stat = rng.normal(size=(nframes, natoms, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + # dp / pt_expt sample (numpy) + np_sample = { + "coord": coords_stat, + "atype": atype_stat, + "atype_ext": atype_stat, + "box": box_stat, + "natoms": natoms_stat, + "dipole": dipole_stat, + "find_dipole": np.float32(1.0), + "aparam": aparam_stat, + } + # pt sample (torch tensors) + pt_sample = { + "coord": numpy_to_torch(coords_stat), + "atype": numpy_to_torch(atype_stat), + "atype_ext": numpy_to_torch(atype_stat), + "box": numpy_to_torch(box_stat), + "natoms": numpy_to_torch(natoms_stat), + "dipole": numpy_to_torch(dipole_stat), + "find_dipole": np.float32(1.0), + "aparam": numpy_to_torch(aparam_stat), + } + + if self.fparam_in_data: + fparam_stat = rng.normal(size=(nframes, 2)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + np_sample["fparam"] = fparam_stat + pt_sample["fparam"] = numpy_to_torch(fparam_stat) + self.expected_fparam_avg = np.mean(fparam_stat, axis=0) + else: + # No fparam → _make_wrapped_sampler injects default_fparam + self.expected_fparam_avg = np.array([0.5, -0.3]) + + self.np_sampled = [np_sample] + self.pt_sampled = [pt_sample] + + # aparam for forward evaluation (1 frame, 6 atoms, 3 aparam) + self.eval_aparam = rng.normal(size=(1, natoms, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + def _eval_dp(self) -> dict: + return self.dp_model( + self.coords, self.atype, box=self.box, aparam=self.eval_aparam + ) + + def _eval_pt(self) -> dict: + return { + kk: torch_to_numpy(vv) + for kk, vv in self.pt_model( + numpy_to_torch(self.coords), + numpy_to_torch(self.atype), + box=numpy_to_torch(self.box), + aparam=numpy_to_torch(self.eval_aparam), + do_atomic_virial=True, + ).items() + } + + def _eval_pt_expt(self) -> dict: + coord_t = pt_expt_numpy_to_torch(self.coords) + coord_t.requires_grad_(True) + return { + k: v.detach().cpu().numpy() + for k, v in self.pt_expt_model( + coord_t, + pt_expt_numpy_to_torch(self.atype), + box=pt_expt_numpy_to_torch(self.box), + aparam=pt_expt_numpy_to_torch(self.eval_aparam), + do_atomic_virial=True, + ).items() + } + + def test_compute_stat(self) -> None: + # 1. Pre-stat forward consistency + dp_ret0 = self._eval_dp() + pt_ret0 = self._eval_pt() + pe_ret0 = self._eval_pt_expt() + for key in ("global_dipole", "dipole"): + np.testing.assert_allclose( + dp_ret0[key], + pt_ret0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Pre-stat dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp_ret0[key], + pe_ret0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Pre-stat dp vs pt_expt mismatch in {key}", + ) + + # 2. Run compute_or_load_stat on all three backends + self.dp_model.compute_or_load_stat(lambda: self.np_sampled) + self.pt_model.compute_or_load_stat(lambda: self.pt_sampled) + self.pt_expt_model.compute_or_load_stat(lambda: self.np_sampled) + + # 3. Serialize all three and compare @variables + dp_ser = self.dp_model.serialize() + pt_ser = self.pt_model.serialize() + pe_ser = self.pt_expt_model.serialize() + compare_variables_recursive(dp_ser, pt_ser) + compare_variables_recursive(dp_ser, pe_ser) + + # 4. Post-stat forward consistency + # Note: DipoleModel's apply_out_stat is a no-op, so output won't change + # after stat computation, but we still verify cross-backend consistency. + dp_ret1 = self._eval_dp() + pt_ret1 = self._eval_pt() + pe_ret1 = self._eval_pt_expt() + for key in ("global_dipole", "dipole"): + np.testing.assert_allclose( + dp_ret1[key], + pt_ret1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Post-stat dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp_ret1[key], + pe_ret1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Post-stat dp vs pt_expt mismatch in {key}", + ) + + # 5. Non-triviality checks + fit_vars = dp_ser["fitting"]["@variables"] + # fparam stats were computed + fparam_avg = np.asarray(fit_vars["fparam_avg"]) + self.assertFalse( + np.allclose(fparam_avg, 0.0), + "fparam_avg is still zero — fparam stats were not computed", + ) + np.testing.assert_allclose( + fparam_avg, + self.expected_fparam_avg, + rtol=1e-10, + atol=1e-10, + err_msg="fparam_avg does not match expected values", + ) + # aparam stats were computed + aparam_avg = np.asarray(fit_vars["aparam_avg"]) + self.assertFalse( + np.allclose(aparam_avg, 0.0), + "aparam_avg is still zero — aparam stats were not computed", + ) + + def test_load_stat_from_file(self) -> None: + import tempfile + from pathlib import ( + Path, + ) + + import h5py + + from deepmd.utils.path import ( + DPPath, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + # Create separate stat files for each backend + dp_h5 = str((Path(tmpdir) / "dp_stat.h5").resolve()) + pt_h5 = str((Path(tmpdir) / "pt_stat.h5").resolve()) + pe_h5 = str((Path(tmpdir) / "pe_stat.h5").resolve()) + for p in (dp_h5, pt_h5, pe_h5): + with h5py.File(p, "w"): + pass + + # 1. Compute stats and save to file + self.dp_model.compute_or_load_stat( + lambda: self.np_sampled, stat_file_path=DPPath(dp_h5, "a") + ) + self.pt_model.compute_or_load_stat( + lambda: self.pt_sampled, stat_file_path=DPPath(pt_h5, "a") + ) + self.pt_expt_model.compute_or_load_stat( + lambda: self.np_sampled, stat_file_path=DPPath(pe_h5, "a") + ) + + # Save the computed serializations as reference + dp_ser_computed = self.dp_model.serialize() + pt_ser_computed = self.pt_model.serialize() + pe_ser_computed = self.pt_expt_model.serialize() + + # 2. Build fresh models from the same initial weights + dp_model2 = get_model_dp(self._model_data) + pt_model2 = DipoleModelPT.deserialize(dp_model2.serialize()) + pe_model2 = DipoleModelPTExpt.deserialize(dp_model2.serialize()) + + # 3. Load stats from file (should NOT call the sampled func) + def raise_error(): + raise RuntimeError("Should load from file, not recompute") + + dp_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(dp_h5, "a") + ) + pt_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(pt_h5, "a") + ) + pe_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(pe_h5, "a") + ) + + # 4. Loaded models should match the computed ones + dp_ser_loaded = dp_model2.serialize() + pt_ser_loaded = pt_model2.serialize() + pe_ser_loaded = pe_model2.serialize() + compare_variables_recursive(dp_ser_computed, dp_ser_loaded) + compare_variables_recursive(pt_ser_computed, pt_ser_loaded) + compare_variables_recursive(pe_ser_computed, pe_ser_loaded) + + # 5. Cross-backend consistency after loading + compare_variables_recursive(dp_ser_loaded, pt_ser_loaded) + compare_variables_recursive(dp_ser_loaded, pe_ser_loaded) diff --git a/source/tests/consistent/model/test_dos.py b/source/tests/consistent/model/test_dos.py index f967973913..5f801129d4 100644 --- a/source/tests/consistent/model/test_dos.py +++ b/source/tests/consistent/model/test_dos.py @@ -6,8 +6,18 @@ import numpy as np +from deepmd.dpmodel.common import ( + to_numpy_array, +) from deepmd.dpmodel.model.dos_model import DOSModel as DOSModelDP from deepmd.dpmodel.model.model import get_model as get_model_dp +from deepmd.dpmodel.utils.nlist import ( + build_neighbor_list, + extend_coord_with_ghosts, +) +from deepmd.dpmodel.utils.region import ( + normalize_coord, +) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) @@ -15,16 +25,21 @@ from ..common import ( INSTALLED_JAX, INSTALLED_PT, + INSTALLED_PT_EXPT, INSTALLED_TF, CommonTest, + parameterized, ) from .common import ( ModelTest, + compare_variables_recursive, ) if INSTALLED_PT: from deepmd.pt.model.model import get_model as get_model_pt from deepmd.pt.model.model.dos_model import DOSModel as DOSModelPT + from deepmd.pt.utils.utils import to_numpy_array as torch_to_numpy + from deepmd.pt.utils.utils import to_torch_tensor as numpy_to_torch else: DOSModelPT = None if INSTALLED_TF: @@ -36,6 +51,11 @@ from deepmd.jax.model.model import get_model as get_model_jax else: DOSModelJAX = None +if INSTALLED_PT_EXPT: + from deepmd.pt_expt.common import to_torch_array as pt_expt_numpy_to_torch + from deepmd.pt_expt.model import DOSModel as DOSModelPTExpt +else: + DOSModelPTExpt = None from deepmd.utils.argcheck import ( model_args, ) @@ -72,6 +92,7 @@ def data(self) -> dict: tf_class = DOSModelTF dp_class = DOSModelDP pt_class = DOSModelPT + pt_expt_class = DOSModelPTExpt jax_class = DOSModelJAX args = model_args() @@ -84,6 +105,8 @@ def get_reference_backend(self): return self.RefBackend.PT if not self.skip_tf: return self.RefBackend.TF + if not self.skip_pt_expt and self.pt_expt_class is not None: + return self.RefBackend.PT_EXPT if not self.skip_dp: return self.RefBackend.DP raise ValueError("No available reference") @@ -105,6 +128,9 @@ def pass_data_to_cls(self, cls, data) -> Any: model = get_model_pt(data) model.atomic_model.out_bias.uniform_() return model + elif cls is DOSModelPTExpt: + dp_model = get_model_dp(data) + return DOSModelPTExpt.deserialize(dp_model.serialize()) elif cls is DOSModelJAX: return get_model_jax(data) return cls(**data, **self.additional_data) @@ -171,6 +197,15 @@ def eval_pt(self, pt_obj: Any) -> Any: self.box, ) + def eval_pt_expt(self, pt_expt_obj: Any) -> Any: + return self.eval_pt_expt_model( + pt_expt_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + def eval_jax(self, jax_obj: Any) -> Any: return self.eval_jax_model( jax_obj, @@ -190,6 +225,7 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: elif backend in { self.RefBackend.DP, self.RefBackend.PT, + self.RefBackend.PT_EXPT, self.RefBackend.JAX, }: return ( @@ -197,3 +233,1313 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: ret["atom_dos"].ravel(), ) raise ValueError(f"Unknown backend: {backend}") + + +@unittest.skipUnless(INSTALLED_PT and INSTALLED_PT_EXPT, "PyTorch is not installed") +class TestDOSModelAPIs(unittest.TestCase): + """Test consistency of model-level APIs between pt and dpmodel backends. + + Both models are constructed from the same serialized weights + (dpmodel -> serialize -> pt deserialize) so that numerical outputs + can be compared directly. + """ + + def setUp(self) -> None: + data = model_args().normalize_value( + { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20], + "rcut_smth": 1.8, + "rcut": 6.0, + "neuron": [2, 4, 8], + "resnet_dt": False, + "axis_neuron": 8, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "type": "dos", + "numb_dos": 2, + "neuron": [4, 4, 4], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + "numb_fparam": 2, + "numb_aparam": 3, + "default_fparam": [0.5, -0.3], + }, + }, + trim_pattern="_*", + ) + # Build dpmodel first, then deserialize into pt/pt_expt to share weights + self.dp_model = get_model_dp(data) + serialized = self.dp_model.serialize() + self.pt_model = DOSModelPT.deserialize(serialized) + self.pt_expt_model = DOSModelPTExpt.deserialize(serialized) + + # Coords / atype / box + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, -1, 3) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32).reshape(1, -1) + self.box = np.array( + [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, 9) + + # Build extended coords + nlist for lower-level calls + rcut = 6.0 + nframes, nloc = self.atype.shape[:2] + coord_normalized = normalize_coord( + self.coords.reshape(nframes, nloc, 3), + self.box.reshape(nframes, 3, 3), + ) + extended_coord, extended_atype, mapping = extend_coord_with_ghosts( + coord_normalized, self.atype, self.box, rcut + ) + nlist = build_neighbor_list( + extended_coord, + extended_atype, + nloc, + rcut, + [20, 20], + distinguish_types=True, + ) + self.extended_coord = extended_coord.reshape(nframes, -1, 3) + self.extended_atype = extended_atype + self.mapping = mapping + self.nlist = nlist + + # aparam for forward evaluation (1 frame, 6 atoms, 3 aparam) + rng = np.random.default_rng(42) + self.eval_aparam = rng.normal(size=(1, nloc, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + def test_translated_output_def(self) -> None: + """translated_output_def should return the same keys on dp, pt, and pt_expt.""" + dp_def = self.dp_model.translated_output_def() + pt_def = self.pt_model.translated_output_def() + pt_expt_def = self.pt_expt_model.translated_output_def() + self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + self.assertEqual(set(dp_def.keys()), set(pt_expt_def.keys())) + for key in dp_def: + self.assertEqual(dp_def[key].shape, pt_def[key].shape) + self.assertEqual(dp_def[key].shape, pt_expt_def[key].shape) + + def test_get_descriptor(self) -> None: + """get_descriptor should return a non-None object on both backends.""" + self.assertIsNotNone(self.dp_model.get_descriptor()) + self.assertIsNotNone(self.pt_model.get_descriptor()) + + def test_get_fitting_net(self) -> None: + """get_fitting_net should return a non-None object on both backends.""" + self.assertIsNotNone(self.dp_model.get_fitting_net()) + self.assertIsNotNone(self.pt_model.get_fitting_net()) + + def test_get_out_bias(self) -> None: + """get_out_bias should return numerically equal values on dp and pt. + + DOSModel uses default apply_out_stat (per-atom type bias), + so the bias storage should be consistent across backends. + """ + dp_bias = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias = torch_to_numpy(self.pt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias, pt_bias, rtol=1e-10, atol=1e-10) + # Verify shape: (n_output_keys x ntypes x numb_dos) for dos + self.assertEqual(dp_bias.shape[1], 2) # ntypes + self.assertGreater(dp_bias.shape[0], 0) # at least one output key + + def test_set_out_bias(self) -> None: + """set_out_bias should update the bias on both backends.""" + dp_bias = to_numpy_array(self.dp_model.get_out_bias()) + new_bias = dp_bias + 1.0 + # dp + self.dp_model.set_out_bias(new_bias) + np.testing.assert_allclose( + to_numpy_array(self.dp_model.get_out_bias()), + new_bias, + rtol=1e-10, + atol=1e-10, + ) + # pt + self.pt_model.set_out_bias(numpy_to_torch(new_bias)) + np.testing.assert_allclose( + torch_to_numpy(self.pt_model.get_out_bias()), + new_bias, + rtol=1e-10, + atol=1e-10, + ) + + def test_model_output_def(self) -> None: + """model_output_def should return the same keys and shapes on dp and pt.""" + dp_def = self.dp_model.model_output_def().get_data() + pt_def = self.pt_model.model_output_def().get_data() + self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + for key in dp_def: + self.assertEqual(dp_def[key].shape, pt_def[key].shape) + + def test_model_output_type(self) -> None: + """model_output_type should return the same list on dp and pt.""" + self.assertEqual( + self.dp_model.model_output_type(), + self.pt_model.model_output_type(), + ) + + def test_do_grad_r(self) -> None: + """do_grad_r should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.do_grad_r("dos"), + self.pt_model.do_grad_r("dos"), + ) + self.assertFalse(self.dp_model.do_grad_r("dos")) + + def test_do_grad_c(self) -> None: + """do_grad_c should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.do_grad_c("dos"), + self.pt_model.do_grad_c("dos"), + ) + self.assertFalse(self.dp_model.do_grad_c("dos")) + + def test_get_rcut(self) -> None: + """get_rcut should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_rcut(), self.pt_model.get_rcut()) + self.assertAlmostEqual(self.dp_model.get_rcut(), 6.0) + + def test_get_type_map(self) -> None: + """get_type_map should return the same list on dp and pt.""" + self.assertEqual(self.dp_model.get_type_map(), self.pt_model.get_type_map()) + self.assertEqual(self.dp_model.get_type_map(), ["O", "H"]) + + def test_get_sel(self) -> None: + """get_sel should return the same list on dp and pt.""" + self.assertEqual(self.dp_model.get_sel(), self.pt_model.get_sel()) + self.assertEqual(self.dp_model.get_sel(), [20, 20]) + + def test_get_nsel(self) -> None: + """get_nsel should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_nsel(), self.pt_model.get_nsel()) + self.assertEqual(self.dp_model.get_nsel(), 40) + + def test_get_nnei(self) -> None: + """get_nnei should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_nnei(), self.pt_model.get_nnei()) + self.assertEqual(self.dp_model.get_nnei(), 40) + + def test_mixed_types(self) -> None: + """mixed_types should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.mixed_types(), self.pt_model.mixed_types()) + # se_e2_a is not mixed-types + self.assertFalse(self.dp_model.mixed_types()) + + def test_has_message_passing(self) -> None: + """has_message_passing should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.has_message_passing(), + self.pt_model.has_message_passing(), + ) + self.assertFalse(self.dp_model.has_message_passing()) + + def test_need_sorted_nlist_for_lower(self) -> None: + """need_sorted_nlist_for_lower should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.need_sorted_nlist_for_lower(), + self.pt_model.need_sorted_nlist_for_lower(), + ) + self.assertFalse(self.dp_model.need_sorted_nlist_for_lower()) + + def test_get_dim_fparam(self) -> None: + """get_dim_fparam should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_dim_fparam(), self.pt_model.get_dim_fparam()) + self.assertEqual(self.dp_model.get_dim_fparam(), 2) + + def test_get_dim_aparam(self) -> None: + """get_dim_aparam should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_dim_aparam(), self.pt_model.get_dim_aparam()) + self.assertEqual(self.dp_model.get_dim_aparam(), 3) + + def test_get_sel_type(self) -> None: + """get_sel_type should return the same list on dp and pt.""" + self.assertEqual(self.dp_model.get_sel_type(), self.pt_model.get_sel_type()) + self.assertEqual(self.dp_model.get_sel_type(), [0, 1]) + + def test_is_aparam_nall(self) -> None: + """is_aparam_nall should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.is_aparam_nall(), self.pt_model.is_aparam_nall()) + self.assertFalse(self.dp_model.is_aparam_nall()) + + def test_get_model_def_script(self) -> None: + """get_model_def_script should return the same value on dp, pt, and pt_expt.""" + dp_val = self.dp_model.get_model_def_script() + pt_val = self.pt_model.get_model_def_script() + pe_val = self.pt_expt_model.get_model_def_script() + self.assertEqual(dp_val, pt_val) + self.assertEqual(dp_val, pe_val) + + def test_get_min_nbor_dist(self) -> None: + """get_min_nbor_dist should return the same value on dp, pt, and pt_expt.""" + dp_val = self.dp_model.get_min_nbor_dist() + pt_val = self.pt_model.get_min_nbor_dist() + pe_val = self.pt_expt_model.get_min_nbor_dist() + self.assertEqual(dp_val, pt_val) + self.assertEqual(dp_val, pe_val) + + def test_set_case_embd(self) -> None: + """set_case_embd should produce consistent results across backends. + + Also verifies that different case indices produce different outputs, + confirming the embedding is actually used. + """ + from deepmd.utils.argcheck import ( + model_args, + ) + + # Build a model with dim_case_embd > 0 + data = model_args().normalize_value( + { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "type": "dos", + "numb_dos": 2, + "neuron": [5, 5], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + "dim_case_embd": 3, + }, + }, + trim_pattern="_*", + ) + dp_model = get_model_dp(data) + serialized = dp_model.serialize() + pt_model = DOSModelPT.deserialize(serialized) + pe_model = DOSModelPTExpt.deserialize(serialized) + + def _eval(case_idx): + dp_model.set_case_embd(case_idx) + pt_model.set_case_embd(case_idx) + pe_model.set_case_embd(case_idx) + dp_ret = dp_model(self.coords, self.atype, box=self.box) + pt_ret = { + k: torch_to_numpy(v) + for k, v in pt_model( + numpy_to_torch(self.coords), + numpy_to_torch(self.atype), + box=numpy_to_torch(self.box), + ).items() + } + coord_t = pt_expt_numpy_to_torch(self.coords) + coord_t.requires_grad_(True) + pe_ret = { + k: v.detach().cpu().numpy() + for k, v in pe_model( + coord_t, + pt_expt_numpy_to_torch(self.atype), + box=pt_expt_numpy_to_torch(self.box), + ).items() + } + return dp_ret, pt_ret, pe_ret + + dp0, pt0, pe0 = _eval(0) + dp1, pt1, pe1 = _eval(1) + + # Cross-backend consistency for each case index + for key in ("dos", "atom_dos"): + np.testing.assert_allclose( + dp0[key], + pt0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 0: dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp0[key], + pe0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 0: dp vs pt_expt mismatch in {key}", + ) + np.testing.assert_allclose( + dp1[key], + pt1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 1: dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp1[key], + pe1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 1: dp vs pt_expt mismatch in {key}", + ) + # Different case indices should produce different outputs + self.assertFalse( + np.allclose(dp0["dos"], dp1["dos"]), + "set_case_embd(0) and set_case_embd(1) produced the same dos", + ) + + def test_atomic_output_def(self) -> None: + """atomic_output_def should return the same keys and shapes on dp and pt.""" + dp_def = self.dp_model.atomic_output_def() + pt_def = self.pt_model.atomic_output_def() + self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + for key in dp_def.keys(): + self.assertEqual(dp_def[key].shape, pt_def[key].shape) + + def test_format_nlist(self) -> None: + """format_nlist should produce the same result on dp and pt.""" + dp_nlist = self.dp_model.format_nlist( + self.extended_coord, + self.extended_atype, + self.nlist, + ) + pt_nlist = torch_to_numpy( + self.pt_model.format_nlist( + numpy_to_torch(self.extended_coord), + numpy_to_torch(self.extended_atype), + numpy_to_torch(self.nlist), + ) + ) + np.testing.assert_equal(dp_nlist, pt_nlist) + + def test_forward_common_atomic(self) -> None: + """forward_common_atomic should produce consistent results on dp and pt. + + Compares at the atomic_model level, where both backends define this method. + """ + dp_ret = self.dp_model.atomic_model.forward_common_atomic( + self.extended_coord, + self.extended_atype, + self.nlist, + mapping=self.mapping, + aparam=self.eval_aparam, + ) + pt_ret = self.pt_model.atomic_model.forward_common_atomic( + numpy_to_torch(self.extended_coord), + numpy_to_torch(self.extended_atype), + numpy_to_torch(self.nlist), + mapping=numpy_to_torch(self.mapping), + aparam=numpy_to_torch(self.eval_aparam), + ) + # Compare the common keys + common_keys = set(dp_ret.keys()) & set(pt_ret.keys()) + self.assertTrue(len(common_keys) > 0) + for key in common_keys: + if dp_ret[key] is not None and pt_ret[key] is not None: + np.testing.assert_allclose( + dp_ret[key], + torch_to_numpy(pt_ret[key]), + rtol=1e-10, + atol=1e-10, + err_msg=f"Mismatch in forward_common_atomic key '{key}'", + ) + + def test_has_default_fparam(self) -> None: + """has_default_fparam should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.has_default_fparam(), + self.pt_model.has_default_fparam(), + ) + self.assertTrue(self.dp_model.has_default_fparam()) + + def test_get_default_fparam(self) -> None: + """get_default_fparam should return consistent values on dp and pt.""" + dp_val = self.dp_model.get_default_fparam() + pt_val = self.pt_model.get_default_fparam() + np.testing.assert_allclose(dp_val, pt_val, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dp_val, [0.5, -0.3], rtol=1e-10, atol=1e-10) + + def test_change_out_bias(self) -> None: + """change_out_bias should produce consistent bias on dp, pt, and pt_expt. + + Tests both set-by-statistic and change-by-statistic modes. + Note: change_out_bias only updates the output bias, not fitting input + stats (fparam/aparam). Fitting stats are updated by compute_or_load_stat. + """ + nframes = 2 + nloc = 6 + numb_fparam = 2 + numb_aparam = 3 + rng = np.random.default_rng(123) + + # Use realistic coords (from setUp, tiled for 2 frames) + coords_2f = np.tile(self.coords, (nframes, 1, 1)) # (2, 6, 3) + atype_2f = np.array([[0, 0, 1, 1, 1, 1], [0, 1, 1, 0, 1, 1]], dtype=np.int32) + box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) + natoms_data = np.array([[6, 6, 2, 4], [6, 6, 2, 4]], dtype=np.int32) + dos_data = rng.normal(size=(nframes, 2)).astype(GLOBAL_NP_FLOAT_PRECISION) + fparam_data = rng.normal(size=(nframes, numb_fparam)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + aparam_data = rng.normal(size=(nframes, nloc, numb_aparam)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + # dpmodel stat data (numpy) + dp_merged = [ + { + "coord": coords_2f, + "atype": atype_2f, + "atype_ext": atype_2f, + "box": box_2f, + "natoms": natoms_data, + "dos": dos_data, + "find_dos": np.float32(1.0), + "fparam": fparam_data, + "aparam": aparam_data, + } + ] + # pt stat data (torch tensors) + pt_merged = [ + { + "coord": numpy_to_torch(coords_2f), + "atype": numpy_to_torch(atype_2f), + "atype_ext": numpy_to_torch(atype_2f), + "box": numpy_to_torch(box_2f), + "natoms": numpy_to_torch(natoms_data), + "dos": numpy_to_torch(dos_data), + "find_dos": np.float32(1.0), + "fparam": numpy_to_torch(fparam_data), + "aparam": numpy_to_torch(aparam_data), + } + ] + # pt_expt stat data (numpy, same as dp) + pe_merged = dp_merged + + # Save initial (zero) bias + dp_bias_init = to_numpy_array(self.dp_model.get_out_bias()).copy() + + # --- Test "set-by-statistic" mode --- + self.dp_model.change_out_bias(dp_merged, bias_adjust_mode="set-by-statistic") + self.pt_model.change_out_bias(pt_merged, bias_adjust_mode="set-by-statistic") + self.pt_expt_model.change_out_bias( + pe_merged, bias_adjust_mode="set-by-statistic" + ) + + # Verify out bias consistency + dp_bias = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias = torch_to_numpy(self.pt_model.get_out_bias()) + pe_bias = to_numpy_array(self.pt_expt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias, pt_bias, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dp_bias, pe_bias, rtol=1e-10, atol=1e-10) + self.assertFalse( + np.allclose(dp_bias, dp_bias_init), + "set-by-statistic did not change the bias from initial values", + ) + + # --- Test "change-by-statistic" mode --- + dp_bias_before = dp_bias.copy() + self.dp_model.change_out_bias(dp_merged, bias_adjust_mode="change-by-statistic") + self.pt_model.change_out_bias(pt_merged, bias_adjust_mode="change-by-statistic") + self.pt_expt_model.change_out_bias( + pe_merged, bias_adjust_mode="change-by-statistic" + ) + + # Verify out bias consistency + dp_bias2 = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias2 = torch_to_numpy(self.pt_model.get_out_bias()) + pe_bias2 = to_numpy_array(self.pt_expt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias2, pt_bias2, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dp_bias2, pe_bias2, rtol=1e-10, atol=1e-10) + self.assertFalse( + np.allclose(dp_bias2, dp_bias_before), + "change-by-statistic did not further change the bias", + ) + + def test_change_type_map(self) -> None: + """change_type_map should produce consistent results on dp and pt. + + Uses a DPA1 (se_atten) descriptor since se_e2_a does not support + change_type_map (non-mixed-types descriptors raise NotImplementedError). + """ + from deepmd.utils.argcheck import model_args as model_args_fn + + data = model_args_fn().normalize_value( + { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_atten", + "sel": 20, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "seed": 1, + "attn": 6, + "attn_layer": 0, + }, + "fitting_net": { + "type": "dos", + "numb_dos": 2, + "neuron": [4, 4, 4], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + }, + }, + trim_pattern="_*", + ) + dp_model = get_model_dp(data) + pt_model = DOSModelPT.deserialize(dp_model.serialize()) + + # Set non-zero out_bias so the swap is non-trivial + dp_bias_orig = to_numpy_array(dp_model.get_out_bias()).copy() + new_bias = dp_bias_orig.copy() + new_bias[:, 0, :] = 1.5 # type 0 ("O") + new_bias[:, 1, :] = -3.7 # type 1 ("H") + dp_model.set_out_bias(new_bias) + pt_model.set_out_bias(numpy_to_torch(new_bias)) + + new_type_map = ["H", "O"] + dp_model.change_type_map(new_type_map) + pt_model.change_type_map(new_type_map) + + # Both should have the new type_map + self.assertEqual(dp_model.get_type_map(), new_type_map) + self.assertEqual(pt_model.get_type_map(), new_type_map) + + # Out_bias should be reordered consistently between backends + dp_bias_new = to_numpy_array(dp_model.get_out_bias()) + pt_bias_new = torch_to_numpy(pt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias_new, pt_bias_new, rtol=1e-10, atol=1e-10) + + # Verify the reorder is correct: old type 0 -> new type 1, old type 1 -> new type 0 + np.testing.assert_allclose( + dp_bias_new[:, 0, :], + new_bias[:, 1, :], + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + dp_bias_new[:, 1, :], + new_bias[:, 0, :], + rtol=1e-10, + atol=1e-10, + ) + + def test_change_type_map_extend_stat(self) -> None: + """change_type_map with model_with_new_type_stat should propagate stats consistently across dp, pt, and pt_expt. + + Verifies that the model-level change_type_map correctly unwraps + model_with_new_type_stat.atomic_model before forwarding to the + atomic model. + """ + from deepmd.utils.argcheck import model_args as model_args_fn + + small_tm = ["O", "H"] + large_tm = ["O", "H", "Li"] + + small_data = model_args_fn().normalize_value( + { + "type_map": small_tm, + "descriptor": { + "type": "se_atten", + "sel": 20, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "seed": 1, + "attn": 6, + "attn_layer": 0, + }, + "fitting_net": { + "type": "dos", + "numb_dos": 2, + "neuron": [4, 4, 4], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + }, + }, + trim_pattern="_*", + ) + large_data = model_args_fn().normalize_value( + { + "type_map": large_tm, + "descriptor": { + "type": "se_atten", + "sel": 20, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "seed": 2, + "attn": 6, + "attn_layer": 0, + }, + "fitting_net": { + "type": "dos", + "numb_dos": 2, + "neuron": [4, 4, 4], + "resnet_dt": True, + "precision": "float64", + "seed": 2, + }, + }, + trim_pattern="_*", + ) + + dp_small = get_model_dp(small_data) + dp_large = get_model_dp(large_data) + + # Set distinguishable random stats on the large model's descriptor + rng = np.random.default_rng(42) + desc_large = dp_large.get_descriptor() + mean_large, std_large = desc_large.get_stat_mean_and_stddev() + mean_rand = rng.random(size=to_numpy_array(mean_large).shape) + std_rand = rng.random(size=to_numpy_array(std_large).shape) + desc_large.set_stat_mean_and_stddev(mean_rand, std_rand) + + # Build pt and pt_expt models from dp serialization + pt_small = DOSModelPT.deserialize(dp_small.serialize()) + pt_large = DOSModelPT.deserialize(dp_large.serialize()) + pt_expt_small = DOSModelPTExpt.deserialize(dp_small.serialize()) + pt_expt_large = DOSModelPTExpt.deserialize(dp_large.serialize()) + + # Extend type map with model_with_new_type_stat at the model level + dp_small.change_type_map(large_tm, model_with_new_type_stat=dp_large) + pt_small.change_type_map(large_tm, model_with_new_type_stat=pt_large) + pt_expt_small.change_type_map(large_tm, model_with_new_type_stat=pt_expt_large) + + # Descriptor stats should be consistent across backends + dp_mean, dp_std = dp_small.get_descriptor().get_stat_mean_and_stddev() + pt_mean, pt_std = pt_small.get_descriptor().get_stat_mean_and_stddev() + pt_expt_mean, pt_expt_std = ( + pt_expt_small.get_descriptor().get_stat_mean_and_stddev() + ) + np.testing.assert_allclose( + to_numpy_array(dp_mean), + torch_to_numpy(pt_mean), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + to_numpy_array(dp_std), + torch_to_numpy(pt_std), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + to_numpy_array(dp_mean), + to_numpy_array(pt_expt_mean), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + to_numpy_array(dp_std), + to_numpy_array(pt_expt_std), + rtol=1e-10, + atol=1e-10, + ) + + def test_update_sel(self) -> None: + """update_sel should return the same result on dp and pt.""" + from unittest.mock import ( + patch, + ) + + from deepmd.dpmodel.model.dp_model import DPModelCommon as DPModelCommonDP + from deepmd.pt.model.model.dp_model import DPModelCommon as DPModelCommonPT + + mock_min_nbor_dist = 0.5 + mock_sel = [10, 20] + local_jdata = { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": "auto", + "rcut_smth": 0.50, + "rcut": 6.00, + }, + "fitting_net": { + "type": "dos", + "numb_dos": 2, + "neuron": [4, 4, 4], + }, + } + type_map = ["O", "H"] + + with patch( + "deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat", + return_value=(mock_min_nbor_dist, mock_sel), + ): + dp_result, dp_min_dist = DPModelCommonDP.update_sel( + None, type_map, local_jdata + ) + + with patch( + "deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat", + return_value=(mock_min_nbor_dist, mock_sel), + ): + pt_result, pt_min_dist = DPModelCommonPT.update_sel( + None, type_map, local_jdata + ) + + self.assertEqual(dp_result, pt_result) + self.assertEqual(dp_min_dist, pt_min_dist) + # Verify sel was actually updated (not still "auto") + self.assertIsInstance(dp_result["descriptor"]["sel"], list) + self.assertNotEqual(dp_result["descriptor"]["sel"], "auto") + + def test_get_ntypes(self) -> None: + """get_ntypes should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_ntypes(), self.pt_model.get_ntypes()) + self.assertEqual(self.dp_model.get_ntypes(), 2) + + def test_compute_or_load_out_stat(self) -> None: + """compute_or_load_out_stat should produce consistent bias on dp and pt. + + Tests both the compute path (from data) and the load path (from file). + DOSModel uses default apply_out_stat (per-atom type bias), + so the stored bias should be consistent across backends. + """ + import tempfile + from pathlib import ( + Path, + ) + + import h5py + + from deepmd.utils.path import ( + DPPath, + ) + + nframes = 2 + coords_2f = np.tile(self.coords, (nframes, 1, 1)) + atype_2f = np.array([[0, 0, 1, 1, 1, 1], [0, 1, 1, 0, 1, 1]], dtype=np.int32) + box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) + natoms_data = np.array([[6, 6, 2, 4], [6, 6, 2, 4]], dtype=np.int32) + dos_data = ( + np.random.default_rng(42) + .normal(size=(nframes, 2)) + .astype(GLOBAL_NP_FLOAT_PRECISION) + ) + + dp_merged = [ + { + "coord": coords_2f, + "atype": atype_2f, + "atype_ext": atype_2f, + "box": box_2f, + "natoms": natoms_data, + "dos": dos_data, + "find_dos": np.float32(1.0), + } + ] + pt_merged = [ + { + "coord": numpy_to_torch(coords_2f), + "atype": numpy_to_torch(atype_2f), + "atype_ext": numpy_to_torch(atype_2f), + "box": numpy_to_torch(box_2f), + "natoms": numpy_to_torch(natoms_data), + "dos": numpy_to_torch(dos_data), + "find_dos": np.float32(1.0), + } + ] + + # Verify bias is initially identical + dp_bias_before = to_numpy_array(self.dp_model.get_out_bias()).copy() + pt_bias_before = torch_to_numpy(self.pt_model.get_out_bias()).copy() + np.testing.assert_allclose( + dp_bias_before, pt_bias_before, rtol=1e-10, atol=1e-10 + ) + + with tempfile.TemporaryDirectory() as tmpdir: + # Create separate h5 files for dp and pt + dp_h5 = str((Path(tmpdir) / "dp_stat.h5").resolve()) + pt_h5 = str((Path(tmpdir) / "pt_stat.h5").resolve()) + with h5py.File(dp_h5, "w"): + pass + with h5py.File(pt_h5, "w"): + pass + dp_stat_path = DPPath(dp_h5, "a") + pt_stat_path = DPPath(pt_h5, "a") + + # 1. Compute stats and save to file + self.dp_model.atomic_model.compute_or_load_out_stat( + dp_merged, stat_file_path=dp_stat_path + ) + self.pt_model.atomic_model.compute_or_load_out_stat( + pt_merged, stat_file_path=pt_stat_path + ) + + dp_bias_after = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias_after = torch_to_numpy(self.pt_model.get_out_bias()) + np.testing.assert_allclose( + dp_bias_after, pt_bias_after, rtol=1e-10, atol=1e-10 + ) + + # 2. Verify both backends saved the same file content + with h5py.File(dp_h5, "r") as dp_f, h5py.File(pt_h5, "r") as pt_f: + dp_keys = sorted(dp_f.keys()) + pt_keys = sorted(pt_f.keys()) + self.assertEqual(dp_keys, pt_keys) + for key in dp_keys: + np.testing.assert_allclose( + np.array(dp_f[key]), + np.array(pt_f[key]), + rtol=1e-10, + atol=1e-10, + err_msg=f"Stat file content mismatch for key {key}", + ) + + # 3. Reset biases to zero, then load from file + zero_bias = np.zeros_like(dp_bias_after) + self.dp_model.set_out_bias(zero_bias) + self.pt_model.set_out_bias(numpy_to_torch(zero_bias)) + + # Use a callable that raises to ensure it loads from file, not recomputes + def raise_error(): + raise RuntimeError("Should not recompute — should load from file") + + self.dp_model.atomic_model.compute_or_load_out_stat( + raise_error, stat_file_path=dp_stat_path + ) + self.pt_model.atomic_model.compute_or_load_out_stat( + raise_error, stat_file_path=pt_stat_path + ) + + dp_bias_loaded = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias_loaded = torch_to_numpy(self.pt_model.get_out_bias()) + + # Loaded biases should match between backends + np.testing.assert_allclose( + dp_bias_loaded, pt_bias_loaded, rtol=1e-10, atol=1e-10 + ) + # Loaded biases should match the originally computed biases + np.testing.assert_allclose( + dp_bias_loaded, dp_bias_after, rtol=1e-10, atol=1e-10 + ) + + def test_get_observed_type_list(self) -> None: + """get_observed_type_list should be consistent across dp, pt, pt_expt. + + Uses mock data containing only type 0 ("O") so that type 1 ("H") is + unobserved and should be absent from the returned list. + """ + nframes = 2 + natoms = 6 + # All atoms are type 0 — type 1 is unobserved + atype_2f = np.zeros((nframes, natoms), dtype=np.int32) + coords_2f = np.tile(self.coords, (nframes, 1, 1)) + box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) + natoms_data = np.array([[natoms, natoms, natoms, 0]] * nframes, dtype=np.int32) + dos_data = ( + np.random.default_rng(42) + .normal(size=(nframes, 2)) + .astype(GLOBAL_NP_FLOAT_PRECISION) + ) + + dp_merged = [ + { + "coord": coords_2f, + "atype": atype_2f, + "atype_ext": atype_2f, + "box": box_2f, + "natoms": natoms_data, + "dos": dos_data, + "find_dos": np.float32(1.0), + } + ] + pt_merged = [ + { + "coord": numpy_to_torch(coords_2f), + "atype": numpy_to_torch(atype_2f), + "atype_ext": numpy_to_torch(atype_2f), + "box": numpy_to_torch(box_2f), + "natoms": numpy_to_torch(natoms_data), + "dos": numpy_to_torch(dos_data), + "find_dos": np.float32(1.0), + } + ] + + self.dp_model.atomic_model.compute_or_load_out_stat(dp_merged) + self.pt_model.atomic_model.compute_or_load_out_stat(pt_merged) + self.pt_expt_model.atomic_model.compute_or_load_out_stat(dp_merged) + + dp_observed = self.dp_model.get_observed_type_list() + pt_observed = self.pt_model.get_observed_type_list() + pe_observed = self.pt_expt_model.get_observed_type_list() + + self.assertEqual(dp_observed, pt_observed) + self.assertEqual(dp_observed, pe_observed) + # DOS uses get_compute_stats_distinguish_types()=True (default), + # so only observed types get non-zero bias — only type 0 ("O"). + self.assertEqual(dp_observed, ["O"]) + + +@parameterized( + (([], []), ([[0, 1]], [1])), # (pair_exclude_types, atom_exclude_types) + (False, True), # fparam_in_data +) +@unittest.skipUnless(INSTALLED_PT and INSTALLED_PT_EXPT, "PT and PT_EXPT are required") +class TestDOSComputeOrLoadStat(unittest.TestCase): + """Test that compute_or_load_stat produces identical statistics on dp, pt, and pt_expt. + + Covers descriptor stats (dstd), fitting stats (fparam, aparam), and output bias. + Parameterized over exclusion types and whether fparam is explicitly provided or + injected via default_fparam. + """ + + def setUp(self) -> None: + (pair_exclude_types, atom_exclude_types), self.fparam_in_data = self.param + data = model_args().normalize_value( + { + "type_map": ["O", "H"], + "pair_exclude_types": pair_exclude_types, + "atom_exclude_types": atom_exclude_types, + "descriptor": { + "type": "dpa3", + "repflow": { + "n_dim": 20, + "e_dim": 10, + "a_dim": 8, + "nlayers": 3, + "e_rcut": 6.0, + "e_rcut_smth": 5.0, + "e_sel": 10, + "a_rcut": 4.0, + "a_rcut_smth": 3.5, + "a_sel": 8, + "axis_neuron": 4, + "update_angle": True, + "update_style": "res_residual", + "update_residual": 0.1, + "update_residual_init": "const", + }, + "precision": "float64", + "seed": 1, + }, + "fitting_net": { + "type": "dos", + "numb_dos": 2, + "neuron": [10, 10], + "precision": "float64", + "seed": 1, + "numb_fparam": 2, + "default_fparam": [0.5, -0.3], + "numb_aparam": 3, + }, + }, + trim_pattern="_*", + ) + + # Save data for reuse in load-from-file test + self._model_data = data + + # Build dp model, then deserialize into pt and pt_expt to share weights + self.dp_model = get_model_dp(data) + serialized = self.dp_model.serialize() + self.pt_model = DOSModelPT.deserialize(serialized) + self.pt_expt_model = DOSModelPTExpt.deserialize(serialized) + + # Test coords / atype / box for forward evaluation + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 0.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, -1, 3) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32).reshape(1, -1) + self.box = np.array( + [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, 9) + + # Mock training data for compute_or_load_stat + natoms = 6 + nframes = 3 + rng = np.random.default_rng(42) + coords_stat = rng.normal(size=(nframes, natoms, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + atype_stat = np.array([[0, 0, 1, 1, 1, 1]] * nframes, dtype=np.int32) + box_stat = np.tile( + np.eye(3, dtype=GLOBAL_NP_FLOAT_PRECISION).reshape(1, 3, 3) * 13.0, + (nframes, 1, 1), + ) + natoms_stat = np.array([[natoms, natoms, 2, 4]] * nframes, dtype=np.int32) + dos_stat = rng.normal(size=(nframes, 2)).astype(GLOBAL_NP_FLOAT_PRECISION) + aparam_stat = rng.normal(size=(nframes, natoms, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + # dp / pt_expt sample (numpy) + np_sample = { + "coord": coords_stat, + "atype": atype_stat, + "atype_ext": atype_stat, + "box": box_stat, + "natoms": natoms_stat, + "dos": dos_stat, + "find_dos": np.float32(1.0), + "aparam": aparam_stat, + } + # pt sample (torch tensors) + pt_sample = { + "coord": numpy_to_torch(coords_stat), + "atype": numpy_to_torch(atype_stat), + "atype_ext": numpy_to_torch(atype_stat), + "box": numpy_to_torch(box_stat), + "natoms": numpy_to_torch(natoms_stat), + "dos": numpy_to_torch(dos_stat), + "find_dos": np.float32(1.0), + "aparam": numpy_to_torch(aparam_stat), + } + + if self.fparam_in_data: + fparam_stat = rng.normal(size=(nframes, 2)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + np_sample["fparam"] = fparam_stat + pt_sample["fparam"] = numpy_to_torch(fparam_stat) + self.expected_fparam_avg = np.mean(fparam_stat, axis=0) + else: + # No fparam -> _make_wrapped_sampler injects default_fparam + self.expected_fparam_avg = np.array([0.5, -0.3]) + + self.np_sampled = [np_sample] + self.pt_sampled = [pt_sample] + + # aparam for forward evaluation (1 frame, 6 atoms, 3 aparam) + self.eval_aparam = rng.normal(size=(1, natoms, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + def _eval_dp(self) -> dict: + return self.dp_model( + self.coords, self.atype, box=self.box, aparam=self.eval_aparam + ) + + def _eval_pt(self) -> dict: + return { + kk: torch_to_numpy(vv) + for kk, vv in self.pt_model( + numpy_to_torch(self.coords), + numpy_to_torch(self.atype), + box=numpy_to_torch(self.box), + aparam=numpy_to_torch(self.eval_aparam), + do_atomic_virial=True, + ).items() + } + + def _eval_pt_expt(self) -> dict: + coord_t = pt_expt_numpy_to_torch(self.coords) + coord_t.requires_grad_(True) + return { + k: v.detach().cpu().numpy() + for k, v in self.pt_expt_model( + coord_t, + pt_expt_numpy_to_torch(self.atype), + box=pt_expt_numpy_to_torch(self.box), + aparam=pt_expt_numpy_to_torch(self.eval_aparam), + do_atomic_virial=True, + ).items() + } + + def test_compute_stat(self) -> None: + # 1. Pre-stat forward consistency + dp_ret0 = self._eval_dp() + pt_ret0 = self._eval_pt() + pe_ret0 = self._eval_pt_expt() + for key in ("dos", "atom_dos"): + np.testing.assert_allclose( + dp_ret0[key], + pt_ret0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Pre-stat dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp_ret0[key], + pe_ret0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Pre-stat dp vs pt_expt mismatch in {key}", + ) + + # 2. Run compute_or_load_stat on all three backends + self.dp_model.compute_or_load_stat(lambda: self.np_sampled) + self.pt_model.compute_or_load_stat(lambda: self.pt_sampled) + self.pt_expt_model.compute_or_load_stat(lambda: self.np_sampled) + + # 3. Serialize all three and compare @variables + dp_ser = self.dp_model.serialize() + pt_ser = self.pt_model.serialize() + pe_ser = self.pt_expt_model.serialize() + compare_variables_recursive(dp_ser, pt_ser) + compare_variables_recursive(dp_ser, pe_ser) + + # 4. Post-stat forward consistency + # DOSModel uses default apply_out_stat (per-atom type bias), so + # output values WILL differ from pre-stat after stat computation. + dp_ret1 = self._eval_dp() + pt_ret1 = self._eval_pt() + pe_ret1 = self._eval_pt_expt() + for key in ("dos", "atom_dos"): + np.testing.assert_allclose( + dp_ret1[key], + pt_ret1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Post-stat dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp_ret1[key], + pe_ret1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Post-stat dp vs pt_expt mismatch in {key}", + ) + + # 5. Non-triviality checks + fit_vars = dp_ser["fitting"]["@variables"] + # fparam stats were computed + fparam_avg = np.asarray(fit_vars["fparam_avg"]) + self.assertFalse( + np.allclose(fparam_avg, 0.0), + "fparam_avg is still zero — fparam stats were not computed", + ) + np.testing.assert_allclose( + fparam_avg, + self.expected_fparam_avg, + rtol=1e-10, + atol=1e-10, + err_msg="fparam_avg does not match expected values", + ) + # aparam stats were computed + aparam_avg = np.asarray(fit_vars["aparam_avg"]) + self.assertFalse( + np.allclose(aparam_avg, 0.0), + "aparam_avg is still zero — aparam stats were not computed", + ) + + def test_load_stat_from_file(self) -> None: + import tempfile + from pathlib import ( + Path, + ) + + import h5py + + from deepmd.utils.path import ( + DPPath, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + # Create separate stat files for each backend + dp_h5 = str((Path(tmpdir) / "dp_stat.h5").resolve()) + pt_h5 = str((Path(tmpdir) / "pt_stat.h5").resolve()) + pe_h5 = str((Path(tmpdir) / "pe_stat.h5").resolve()) + for p in (dp_h5, pt_h5, pe_h5): + with h5py.File(p, "w"): + pass + + # 1. Compute stats and save to file + self.dp_model.compute_or_load_stat( + lambda: self.np_sampled, stat_file_path=DPPath(dp_h5, "a") + ) + self.pt_model.compute_or_load_stat( + lambda: self.pt_sampled, stat_file_path=DPPath(pt_h5, "a") + ) + self.pt_expt_model.compute_or_load_stat( + lambda: self.np_sampled, stat_file_path=DPPath(pe_h5, "a") + ) + + # Save the computed serializations as reference + dp_ser_computed = self.dp_model.serialize() + pt_ser_computed = self.pt_model.serialize() + pe_ser_computed = self.pt_expt_model.serialize() + + # 2. Build fresh models from the same initial weights + dp_model2 = get_model_dp(self._model_data) + pt_model2 = DOSModelPT.deserialize(dp_model2.serialize()) + pe_model2 = DOSModelPTExpt.deserialize(dp_model2.serialize()) + + # 3. Load stats from file (should NOT call the sampled func) + def raise_error(): + raise RuntimeError("Should load from file, not recompute") + + dp_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(dp_h5, "a") + ) + pt_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(pt_h5, "a") + ) + pe_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(pe_h5, "a") + ) + + # 4. Loaded models should match the computed ones + dp_ser_loaded = dp_model2.serialize() + pt_ser_loaded = pt_model2.serialize() + pe_ser_loaded = pe_model2.serialize() + compare_variables_recursive(dp_ser_computed, dp_ser_loaded) + compare_variables_recursive(pt_ser_computed, pt_ser_loaded) + compare_variables_recursive(pe_ser_computed, pe_ser_loaded) + + # 5. Cross-backend consistency after loading + compare_variables_recursive(dp_ser_loaded, pt_ser_loaded) + compare_variables_recursive(dp_ser_loaded, pe_ser_loaded) diff --git a/source/tests/consistent/model/test_dpa1.py b/source/tests/consistent/model/test_dpa1.py index bacca12413..b32570d024 100644 --- a/source/tests/consistent/model/test_dpa1.py +++ b/source/tests/consistent/model/test_dpa1.py @@ -16,6 +16,7 @@ INSTALLED_JAX, INSTALLED_PD, INSTALLED_PT, + INSTALLED_PT_EXPT, INSTALLED_TF, SKIP_FLAG, CommonTest, @@ -43,6 +44,10 @@ from deepmd.pd.model.model.ener_model import EnergyModel as EnergyModelPD else: EnergyModelPD = None +if INSTALLED_PT_EXPT: + from deepmd.pt_expt.model import EnergyModel as EnergyModelPTExpt +else: + EnergyModelPTExpt = None if INSTALLED_JAX: from deepmd.jax.model.ener_model import EnergyModel as EnergyModelJAX from deepmd.jax.model.model import get_model as get_model_jax @@ -97,6 +102,7 @@ def data(self) -> dict: dp_class = EnergyModelDP pt_class = EnergyModelPT pd_class = EnergyModelPD + pt_expt_class = EnergyModelPTExpt jax_class = EnergyModelJAX args = model_args() @@ -109,6 +115,8 @@ def get_reference_backend(self): return self.RefBackend.PT if not self.skip_tf: return self.RefBackend.TF + if not self.skip_pt_expt and self.pt_expt_class is not None: + return self.RefBackend.PT_EXPT if not self.skip_pd: return self.RefBackend.PD if not self.skip_jax: @@ -128,6 +136,9 @@ def pass_data_to_cls(self, cls, data) -> Any: return get_model_dp(data) elif cls is EnergyModelPT: return get_model_pt(data) + elif cls is EnergyModelPTExpt: + dp_model = get_model_dp(data) + return EnergyModelPTExpt.deserialize(dp_model.serialize()) elif cls is EnergyModelPD: return get_model_pd(data) elif cls is EnergyModelJAX: @@ -201,6 +212,15 @@ def eval_pt(self, pt_obj: Any) -> Any: self.box, ) + def eval_pt_expt(self, pt_expt_obj: Any) -> Any: + return self.eval_pt_expt_model( + pt_expt_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + def eval_pd(self, pd_obj: Any) -> Any: return self.eval_pd_model( pd_obj, @@ -239,6 +259,7 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: ) elif backend in { self.RefBackend.PT, + self.RefBackend.PT_EXPT, self.RefBackend.PD, self.RefBackend.JAX, }: @@ -250,3 +271,50 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: ret["atom_virial"].ravel(), ) raise ValueError(f"Unknown backend: {backend}") + + +@unittest.skipUnless(INSTALLED_PT and INSTALLED_PT_EXPT, "PyTorch is not installed") +class TestDPA1EnerModelAPIs(unittest.TestCase): + """Test translated_output_def consistency across dp, pt, and pt_expt backends.""" + + def setUp(self) -> None: + data = model_args().normalize_value( + { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_atten", + "sel": 40, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "seed": 1, + "attn": 128, + "attn_layer": 0, + "precision": "float64", + }, + "fitting_net": { + "neuron": [5, 5], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + }, + }, + trim_pattern="_*", + ) + self.dp_model = get_model_dp(data) + serialized = self.dp_model.serialize() + self.pt_model = EnergyModelPT.deserialize(serialized) + self.pt_expt_model = EnergyModelPTExpt.deserialize(serialized) + + def test_translated_output_def(self) -> None: + """translated_output_def should return the same keys on dp, pt, and pt_expt.""" + dp_def = self.dp_model.translated_output_def() + pt_def = self.pt_model.translated_output_def() + pt_expt_def = self.pt_expt_model.translated_output_def() + self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + self.assertEqual(set(dp_def.keys()), set(pt_expt_def.keys())) + for key in dp_def: + self.assertEqual(dp_def[key].shape, pt_def[key].shape) + self.assertEqual(dp_def[key].shape, pt_expt_def[key].shape) diff --git a/source/tests/consistent/model/test_ener.py b/source/tests/consistent/model/test_ener.py index 0bb1c343bd..f47365e2cf 100644 --- a/source/tests/consistent/model/test_ener.py +++ b/source/tests/consistent/model/test_ener.py @@ -34,6 +34,7 @@ ) from .common import ( ModelTest, + compare_variables_recursive, ) if INSTALLED_PT: @@ -146,7 +147,7 @@ def get_reference_backend(self): @property def skip_tf(self): - return ( + return not INSTALLED_TF or ( self.data["pair_exclude_types"] != [] or self.data["atom_exclude_types"] != [] ) @@ -546,7 +547,7 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: raise ValueError(f"Unknown backend: {backend}") -@unittest.skipUnless(INSTALLED_PT, "PyTorch is not installed") +@unittest.skipUnless(INSTALLED_PT and INSTALLED_PT_EXPT, "PyTorch is not installed") class TestEnerModelAPIs(unittest.TestCase): """Test consistency of model-level APIs between pt and dpmodel backends. @@ -580,14 +581,18 @@ def setUp(self) -> None: "resnet_dt": True, "precision": "float64", "seed": 1, + "numb_fparam": 2, + "numb_aparam": 3, + "default_fparam": [0.5, -0.3], }, }, trim_pattern="_*", ) - # Build dpmodel first, then deserialize into pt to share weights + # Build dpmodel first, then deserialize into pt/pt_expt to share weights self.dp_model = get_model_dp(data) serialized = self.dp_model.serialize() self.pt_model = EnergyModelPT.deserialize(serialized) + self.pt_expt_model = EnergyModelPTExpt.deserialize(serialized) # Coords / atype / box self.coords = np.array( @@ -642,13 +647,22 @@ def setUp(self) -> None: self.mapping = mapping self.nlist = nlist + # aparam for forward evaluation (1 frame, 6 atoms, 3 aparam) + rng = np.random.default_rng(42) + self.eval_aparam = rng.normal(size=(1, nloc, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + def test_translated_output_def(self) -> None: - """translated_output_def should return the same keys on dp and pt.""" + """translated_output_def should return the same keys on dp, pt, and pt_expt.""" dp_def = self.dp_model.translated_output_def() pt_def = self.pt_model.translated_output_def() + pt_expt_def = self.pt_expt_model.translated_output_def() self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + self.assertEqual(set(dp_def.keys()), set(pt_expt_def.keys())) for key in dp_def: self.assertEqual(dp_def[key].shape, pt_def[key].shape) + self.assertEqual(dp_def[key].shape, pt_expt_def[key].shape) def test_get_descriptor(self) -> None: """get_descriptor should return a non-None object on both backends.""" @@ -775,12 +789,12 @@ def test_need_sorted_nlist_for_lower(self) -> None: def test_get_dim_fparam(self) -> None: """get_dim_fparam should return the same value on dp and pt.""" self.assertEqual(self.dp_model.get_dim_fparam(), self.pt_model.get_dim_fparam()) - self.assertEqual(self.dp_model.get_dim_fparam(), 0) + self.assertEqual(self.dp_model.get_dim_fparam(), 2) def test_get_dim_aparam(self) -> None: """get_dim_aparam should return the same value on dp and pt.""" self.assertEqual(self.dp_model.get_dim_aparam(), self.pt_model.get_dim_aparam()) - self.assertEqual(self.dp_model.get_dim_aparam(), 0) + self.assertEqual(self.dp_model.get_dim_aparam(), 3) def test_get_sel_type(self) -> None: """get_sel_type should return the same list on dp and pt.""" @@ -793,6 +807,127 @@ def test_is_aparam_nall(self) -> None: self.assertEqual(self.dp_model.is_aparam_nall(), self.pt_model.is_aparam_nall()) self.assertFalse(self.dp_model.is_aparam_nall()) + def test_get_model_def_script(self) -> None: + """get_model_def_script should return the same value on dp, pt, and pt_expt.""" + dp_val = self.dp_model.get_model_def_script() + pt_val = self.pt_model.get_model_def_script() + pe_val = self.pt_expt_model.get_model_def_script() + self.assertEqual(dp_val, pt_val) + self.assertEqual(dp_val, pe_val) + + def test_get_min_nbor_dist(self) -> None: + """get_min_nbor_dist should return the same value on dp, pt, and pt_expt.""" + dp_val = self.dp_model.get_min_nbor_dist() + pt_val = self.pt_model.get_min_nbor_dist() + pe_val = self.pt_expt_model.get_min_nbor_dist() + self.assertEqual(dp_val, pt_val) + self.assertEqual(dp_val, pe_val) + + def test_set_case_embd(self) -> None: + """set_case_embd should produce consistent results across backends. + + Also verifies that different case indices produce different outputs, + confirming the embedding is actually used. + """ + from deepmd.utils.argcheck import ( + model_args, + ) + + # Build a model with dim_case_embd > 0 + data = model_args().normalize_value( + { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "neuron": [5, 5], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + "dim_case_embd": 3, + }, + }, + trim_pattern="_*", + ) + dp_model = get_model_dp(data) + serialized = dp_model.serialize() + pt_model = EnergyModelPT.deserialize(serialized) + pe_model = EnergyModelPTExpt.deserialize(serialized) + + def _eval(case_idx): + dp_model.set_case_embd(case_idx) + pt_model.set_case_embd(case_idx) + pe_model.set_case_embd(case_idx) + dp_ret = dp_model(self.coords, self.atype, box=self.box) + pt_ret = { + k: torch_to_numpy(v) + for k, v in pt_model( + numpy_to_torch(self.coords), + numpy_to_torch(self.atype), + box=numpy_to_torch(self.box), + ).items() + } + coord_t = pt_expt_numpy_to_torch(self.coords) + coord_t.requires_grad_(True) + pe_ret = { + k: v.detach().cpu().numpy() + for k, v in pe_model( + coord_t, + pt_expt_numpy_to_torch(self.atype), + box=pt_expt_numpy_to_torch(self.box), + ).items() + } + return dp_ret, pt_ret, pe_ret + + dp0, pt0, pe0 = _eval(0) + dp1, pt1, pe1 = _eval(1) + + # Cross-backend consistency for each case index + for key in ("energy", "atom_energy"): + np.testing.assert_allclose( + dp0[key], + pt0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 0: dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp0[key], + pe0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 0: dp vs pt_expt mismatch in {key}", + ) + np.testing.assert_allclose( + dp1[key], + pt1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 1: dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp1[key], + pe1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 1: dp vs pt_expt mismatch in {key}", + ) + # Different case indices should produce different outputs + self.assertFalse( + np.allclose(dp0["energy"], dp1["energy"]), + "set_case_embd(0) and set_case_embd(1) produced the same energy", + ) + def test_atomic_output_def(self) -> None: """atomic_output_def should return the same keys and shapes on dp and pt.""" dp_def = self.dp_model.atomic_output_def() @@ -827,12 +962,14 @@ def test_forward_common_atomic(self) -> None: self.extended_atype, self.nlist, mapping=self.mapping, + aparam=self.eval_aparam, ) pt_ret = self.pt_model.atomic_model.forward_common_atomic( numpy_to_torch(self.extended_coord), numpy_to_torch(self.extended_atype), numpy_to_torch(self.nlist), mapping=numpy_to_torch(self.mapping), + aparam=numpy_to_torch(self.eval_aparam), ) # Compare the common keys common_keys = set(dp_ret.keys()) & set(pt_ret.keys()) @@ -853,26 +990,40 @@ def test_has_default_fparam(self) -> None: self.dp_model.has_default_fparam(), self.pt_model.has_default_fparam(), ) - self.assertFalse(self.dp_model.has_default_fparam()) + self.assertTrue(self.dp_model.has_default_fparam()) def test_get_default_fparam(self) -> None: - """get_default_fparam should return None on both dp and pt (no fparam configured).""" + """get_default_fparam should return consistent values on dp and pt.""" dp_val = self.dp_model.get_default_fparam() pt_val = self.pt_model.get_default_fparam() - self.assertIsNone(dp_val) - self.assertIsNone(pt_val) - # Note: both return None because no default_fparam is configured. - # A non-trivial return requires configuring default_fparam in the fitting net. + np.testing.assert_allclose(dp_val, pt_val, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dp_val, [0.5, -0.3], rtol=1e-10, atol=1e-10) def test_change_out_bias(self) -> None: - """change_out_bias should produce consistent bias on dp and pt.""" + """change_out_bias should produce consistent bias on dp, pt, and pt_expt. + + Tests both set-by-statistic and change-by-statistic modes. + Note: change_out_bias only updates the output bias, not fitting input + stats (fparam/aparam). Fitting stats are updated by compute_or_load_stat. + """ nframes = 2 + nloc = 6 + numb_fparam = 2 + numb_aparam = 3 + rng = np.random.default_rng(123) + # Use realistic coords (from setUp, tiled for 2 frames) coords_2f = np.tile(self.coords, (nframes, 1, 1)) # (2, 6, 3) atype_2f = np.array([[0, 0, 1, 1, 1, 1], [0, 1, 1, 0, 1, 1]], dtype=np.int32) box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) natoms_data = np.array([[6, 6, 2, 4], [6, 6, 2, 4]], dtype=np.int32) energy_data = np.array([10.0, 20.0]).reshape(nframes, 1) + fparam_data = rng.normal(size=(nframes, numb_fparam)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + aparam_data = rng.normal(size=(nframes, nloc, numb_aparam)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) # dpmodel stat data (numpy) dp_merged = [ @@ -884,6 +1035,8 @@ def test_change_out_bias(self) -> None: "natoms": natoms_data, "energy": energy_data, "find_energy": np.float32(1.0), + "fparam": fparam_data, + "aparam": aparam_data, } ] # pt stat data (torch tensors) @@ -896,32 +1049,48 @@ def test_change_out_bias(self) -> None: "natoms": numpy_to_torch(natoms_data), "energy": numpy_to_torch(energy_data), "find_energy": np.float32(1.0), + "fparam": numpy_to_torch(fparam_data), + "aparam": numpy_to_torch(aparam_data), } ] + # pt_expt stat data (numpy, same as dp) + pe_merged = dp_merged # Save initial (zero) bias dp_bias_init = to_numpy_array(self.dp_model.get_out_bias()).copy() - # Test "set-by-statistic" mode + # --- Test "set-by-statistic" mode --- self.dp_model.change_out_bias(dp_merged, bias_adjust_mode="set-by-statistic") self.pt_model.change_out_bias(pt_merged, bias_adjust_mode="set-by-statistic") + self.pt_expt_model.change_out_bias( + pe_merged, bias_adjust_mode="set-by-statistic" + ) + + # Verify out bias consistency dp_bias = to_numpy_array(self.dp_model.get_out_bias()) pt_bias = torch_to_numpy(self.pt_model.get_out_bias()) + pe_bias = to_numpy_array(self.pt_expt_model.get_out_bias()) np.testing.assert_allclose(dp_bias, pt_bias, rtol=1e-10, atol=1e-10) - # Verify bias actually changed from initial zeros + np.testing.assert_allclose(dp_bias, pe_bias, rtol=1e-10, atol=1e-10) self.assertFalse( np.allclose(dp_bias, dp_bias_init), "set-by-statistic did not change the bias from initial values", ) - # Test "change-by-statistic" mode (adjusts bias based on model predictions) + # --- Test "change-by-statistic" mode --- dp_bias_before = dp_bias.copy() self.dp_model.change_out_bias(dp_merged, bias_adjust_mode="change-by-statistic") self.pt_model.change_out_bias(pt_merged, bias_adjust_mode="change-by-statistic") + self.pt_expt_model.change_out_bias( + pe_merged, bias_adjust_mode="change-by-statistic" + ) + + # Verify out bias consistency dp_bias2 = to_numpy_array(self.dp_model.get_out_bias()) pt_bias2 = torch_to_numpy(self.pt_model.get_out_bias()) + pe_bias2 = to_numpy_array(self.pt_expt_model.get_out_bias()) np.testing.assert_allclose(dp_bias2, pt_bias2, rtol=1e-10, atol=1e-10) - # Verify change-by-statistic further modified the bias + np.testing.assert_allclose(dp_bias2, pe_bias2, rtol=1e-10, atol=1e-10) self.assertFalse( np.allclose(dp_bias2, dp_bias_before), "change-by-statistic did not further change the bias", @@ -998,6 +1167,122 @@ def test_change_type_map(self) -> None: atol=1e-10, ) + def test_change_type_map_extend_stat(self) -> None: + """change_type_map with model_with_new_type_stat should propagate stats consistently across dp, pt, and pt_expt. + + Verifies that the model-level change_type_map correctly unwraps + model_with_new_type_stat.atomic_model before forwarding to the + atomic model. + """ + from deepmd.utils.argcheck import model_args as model_args_fn + + small_tm = ["O", "H"] + large_tm = ["O", "H", "Li"] + + small_data = model_args_fn().normalize_value( + { + "type_map": small_tm, + "descriptor": { + "type": "se_atten", + "sel": 20, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "seed": 1, + "attn": 6, + "attn_layer": 0, + }, + "fitting_net": { + "neuron": [5, 5], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + }, + }, + trim_pattern="_*", + ) + large_data = model_args_fn().normalize_value( + { + "type_map": large_tm, + "descriptor": { + "type": "se_atten", + "sel": 20, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "seed": 2, + "attn": 6, + "attn_layer": 0, + }, + "fitting_net": { + "neuron": [5, 5], + "resnet_dt": True, + "precision": "float64", + "seed": 2, + }, + }, + trim_pattern="_*", + ) + + dp_small = get_model_dp(small_data) + dp_large = get_model_dp(large_data) + + # Set distinguishable random stats on the large model's descriptor + rng = np.random.default_rng(42) + desc_large = dp_large.get_descriptor() + mean_large, std_large = desc_large.get_stat_mean_and_stddev() + mean_rand = rng.random(size=to_numpy_array(mean_large).shape) + std_rand = rng.random(size=to_numpy_array(std_large).shape) + desc_large.set_stat_mean_and_stddev(mean_rand, std_rand) + + # Build pt and pt_expt models from dp serialization + pt_small = EnergyModelPT.deserialize(dp_small.serialize()) + pt_large = EnergyModelPT.deserialize(dp_large.serialize()) + pt_expt_small = EnergyModelPTExpt.deserialize(dp_small.serialize()) + pt_expt_large = EnergyModelPTExpt.deserialize(dp_large.serialize()) + + # Extend type map with model_with_new_type_stat at the model level + dp_small.change_type_map(large_tm, model_with_new_type_stat=dp_large) + pt_small.change_type_map(large_tm, model_with_new_type_stat=pt_large) + pt_expt_small.change_type_map(large_tm, model_with_new_type_stat=pt_expt_large) + + # Descriptor stats should be consistent across backends + dp_mean, dp_std = dp_small.get_descriptor().get_stat_mean_and_stddev() + pt_mean, pt_std = pt_small.get_descriptor().get_stat_mean_and_stddev() + pt_expt_mean, pt_expt_std = ( + pt_expt_small.get_descriptor().get_stat_mean_and_stddev() + ) + np.testing.assert_allclose( + to_numpy_array(dp_mean), + torch_to_numpy(pt_mean), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + to_numpy_array(dp_std), + torch_to_numpy(pt_std), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + to_numpy_array(dp_mean), + to_numpy_array(pt_expt_mean), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + to_numpy_array(dp_std), + to_numpy_array(pt_expt_std), + rtol=1e-10, + atol=1e-10, + ) + def test_update_sel(self) -> None: """update_sel should return the same result on dp and pt.""" from unittest.mock import ( @@ -1177,3 +1462,388 @@ def raise_error(): np.testing.assert_allclose( dp_bias_loaded, dp_bias_after, rtol=1e-10, atol=1e-10 ) + + def test_get_observed_type_list(self) -> None: + """get_observed_type_list should be consistent across dp, pt, pt_expt. + + Uses mock data containing only type 0 ("O") so that type 1 ("H") is + unobserved and should be absent from the returned list. + """ + nframes = 2 + natoms = 6 + # All atoms are type 0 — type 1 is unobserved + atype_2f = np.zeros((nframes, natoms), dtype=np.int32) + coords_2f = np.tile(self.coords, (nframes, 1, 1)) + box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) + natoms_data = np.array([[natoms, natoms, natoms, 0]] * nframes, dtype=np.int32) + energy_data = np.array([10.0, 20.0]).reshape(nframes, 1) + + dp_merged = [ + { + "coord": coords_2f, + "atype": atype_2f, + "atype_ext": atype_2f, + "box": box_2f, + "natoms": natoms_data, + "energy": energy_data, + "find_energy": np.float32(1.0), + } + ] + pt_merged = [ + { + "coord": numpy_to_torch(coords_2f), + "atype": numpy_to_torch(atype_2f), + "atype_ext": numpy_to_torch(atype_2f), + "box": numpy_to_torch(box_2f), + "natoms": numpy_to_torch(natoms_data), + "energy": numpy_to_torch(energy_data), + "find_energy": np.float32(1.0), + } + ] + + self.dp_model.atomic_model.compute_or_load_out_stat(dp_merged) + self.pt_model.atomic_model.compute_or_load_out_stat(pt_merged) + self.pt_expt_model.atomic_model.compute_or_load_out_stat(dp_merged) + + dp_observed = self.dp_model.get_observed_type_list() + pt_observed = self.pt_model.get_observed_type_list() + pe_observed = self.pt_expt_model.get_observed_type_list() + + self.assertEqual(dp_observed, pt_observed) + self.assertEqual(dp_observed, pe_observed) + # Only type 0 ("O") should be observed + self.assertEqual(dp_observed, ["O"]) + + +@parameterized( + (([], []), ([[0, 1]], [1])), # (pair_exclude_types, atom_exclude_types) + (False, True), # fparam_in_data +) +@unittest.skipUnless(INSTALLED_PT and INSTALLED_PT_EXPT, "PT and PT_EXPT are required") +class TestEnerComputeOrLoadStat(unittest.TestCase): + """Test that compute_or_load_stat produces identical statistics on dp, pt, and pt_expt. + + Covers descriptor stats (dstd), fitting stats (fparam, aparam), and output bias. + Parameterized over exclusion types and whether fparam is explicitly provided or + injected via default_fparam. + """ + + def setUp(self) -> None: + (pair_exclude_types, atom_exclude_types), self.fparam_in_data = self.param + data = model_args().normalize_value( + { + "type_map": ["O", "H"], + "pair_exclude_types": pair_exclude_types, + "atom_exclude_types": atom_exclude_types, + "descriptor": { + "type": "dpa3", + "repflow": { + "n_dim": 20, + "e_dim": 10, + "a_dim": 8, + "nlayers": 3, + "e_rcut": 6.0, + "e_rcut_smth": 5.0, + "e_sel": 10, + "a_rcut": 4.0, + "a_rcut_smth": 3.5, + "a_sel": 8, + "axis_neuron": 4, + "update_angle": True, + "update_style": "res_residual", + "update_residual": 0.1, + "update_residual_init": "const", + }, + "precision": "float64", + "seed": 1, + }, + "fitting_net": { + "neuron": [10, 10], + "precision": "float64", + "seed": 1, + "numb_fparam": 2, + "default_fparam": [0.5, -0.3], + "numb_aparam": 3, + }, + }, + trim_pattern="_*", + ) + + # Save data for reuse in load-from-file test + self._model_data = data + + # Build dp model, then deserialize into pt and pt_expt to share weights + self.dp_model = get_model_dp(data) + serialized = self.dp_model.serialize() + self.pt_model = EnergyModelPT.deserialize(serialized) + self.pt_expt_model = EnergyModelPTExpt.deserialize(serialized) + + # Test coords / atype / box for forward evaluation + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 0.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, -1, 3) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32).reshape(1, -1) + self.box = np.array( + [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, 9) + + # Mock training data for compute_or_load_stat + natoms = 6 + nframes = 3 + rng = np.random.default_rng(42) + coords_stat = rng.normal(size=(nframes, natoms, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + atype_stat = np.array([[0, 0, 1, 1, 1, 1]] * nframes, dtype=np.int32) + box_stat = np.tile( + np.eye(3, dtype=GLOBAL_NP_FLOAT_PRECISION).reshape(1, 3, 3) * 13.0, + (nframes, 1, 1), + ) + natoms_stat = np.array([[natoms, natoms, 2, 4]] * nframes, dtype=np.int32) + energy_stat = rng.normal(size=(nframes, 1)).astype(GLOBAL_NP_FLOAT_PRECISION) + aparam_stat = rng.normal(size=(nframes, natoms, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + # dp / pt_expt sample (numpy) + np_sample = { + "coord": coords_stat, + "atype": atype_stat, + "atype_ext": atype_stat, + "box": box_stat, + "natoms": natoms_stat, + "energy": energy_stat, + "find_energy": np.float32(1.0), + "aparam": aparam_stat, + } + # pt sample (torch tensors) + pt_sample = { + "coord": numpy_to_torch(coords_stat), + "atype": numpy_to_torch(atype_stat), + "atype_ext": numpy_to_torch(atype_stat), + "box": numpy_to_torch(box_stat), + "natoms": numpy_to_torch(natoms_stat), + "energy": numpy_to_torch(energy_stat), + "find_energy": np.float32(1.0), + "aparam": numpy_to_torch(aparam_stat), + } + + if self.fparam_in_data: + fparam_stat = rng.normal(size=(nframes, 2)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + np_sample["fparam"] = fparam_stat + pt_sample["fparam"] = numpy_to_torch(fparam_stat) + self.expected_fparam_avg = np.mean(fparam_stat, axis=0) + else: + # No fparam → _make_wrapped_sampler injects default_fparam + self.expected_fparam_avg = np.array([0.5, -0.3]) + + self.np_sampled = [np_sample] + self.pt_sampled = [pt_sample] + + # aparam for forward evaluation (1 frame, 6 atoms, 3 aparam) + self.eval_aparam = rng.normal(size=(1, natoms, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + def _eval_dp(self) -> dict: + return self.dp_model( + self.coords, self.atype, box=self.box, aparam=self.eval_aparam + ) + + def _eval_pt(self) -> dict: + return { + kk: torch_to_numpy(vv) + for kk, vv in self.pt_model( + numpy_to_torch(self.coords), + numpy_to_torch(self.atype), + box=numpy_to_torch(self.box), + aparam=numpy_to_torch(self.eval_aparam), + do_atomic_virial=True, + ).items() + } + + def _eval_pt_expt(self) -> dict: + coord_t = pt_expt_numpy_to_torch(self.coords) + coord_t.requires_grad_(True) + return { + k: v.detach().cpu().numpy() + for k, v in self.pt_expt_model( + coord_t, + pt_expt_numpy_to_torch(self.atype), + box=pt_expt_numpy_to_torch(self.box), + aparam=pt_expt_numpy_to_torch(self.eval_aparam), + do_atomic_virial=True, + ).items() + } + + def test_compute_stat(self) -> None: + # 1. Pre-stat forward consistency + dp_ret0 = self._eval_dp() + pt_ret0 = self._eval_pt() + pe_ret0 = self._eval_pt_expt() + for key in ("energy", "atom_energy"): + np.testing.assert_allclose( + dp_ret0[key], + pt_ret0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Pre-stat dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp_ret0[key], + pe_ret0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Pre-stat dp vs pt_expt mismatch in {key}", + ) + + # 2. Run compute_or_load_stat on all three backends + # deepcopy because stat.py mutates natoms in-place when atom_exclude_types + # is non-empty (natoms[:, 2:] *= type_mask). + from copy import ( + deepcopy, + ) + + self.dp_model.compute_or_load_stat(lambda: deepcopy(self.np_sampled)) + self.pt_model.compute_or_load_stat(lambda: deepcopy(self.pt_sampled)) + self.pt_expt_model.compute_or_load_stat(lambda: deepcopy(self.np_sampled)) + + # 3. Serialize all three and compare @variables + dp_ser = self.dp_model.serialize() + pt_ser = self.pt_model.serialize() + pe_ser = self.pt_expt_model.serialize() + compare_variables_recursive(dp_ser, pt_ser) + compare_variables_recursive(dp_ser, pe_ser) + + # 4. Post-stat forward consistency + dp_ret1 = self._eval_dp() + pt_ret1 = self._eval_pt() + pe_ret1 = self._eval_pt_expt() + for key in ("energy", "atom_energy"): + np.testing.assert_allclose( + dp_ret1[key], + pt_ret1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Post-stat dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp_ret1[key], + pe_ret1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Post-stat dp vs pt_expt mismatch in {key}", + ) + + # 5. Non-triviality checks + fit_vars = dp_ser["fitting"]["@variables"] + # fparam stats were computed + fparam_avg = np.asarray(fit_vars["fparam_avg"]) + self.assertFalse( + np.allclose(fparam_avg, 0.0), + "fparam_avg is still zero — fparam stats were not computed", + ) + np.testing.assert_allclose( + fparam_avg, + self.expected_fparam_avg, + rtol=1e-10, + atol=1e-10, + err_msg="fparam_avg does not match expected values", + ) + # aparam stats were computed + aparam_avg = np.asarray(fit_vars["aparam_avg"]) + self.assertFalse( + np.allclose(aparam_avg, 0.0), + "aparam_avg is still zero — aparam stats were not computed", + ) + + def test_load_stat_from_file(self) -> None: + import tempfile + from pathlib import ( + Path, + ) + + import h5py + + from deepmd.utils.path import ( + DPPath, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + # Create separate stat files for each backend + dp_h5 = str((Path(tmpdir) / "dp_stat.h5").resolve()) + pt_h5 = str((Path(tmpdir) / "pt_stat.h5").resolve()) + pe_h5 = str((Path(tmpdir) / "pe_stat.h5").resolve()) + for p in (dp_h5, pt_h5, pe_h5): + with h5py.File(p, "w"): + pass + + # 1. Compute stats and save to file + self.dp_model.compute_or_load_stat( + lambda: self.np_sampled, stat_file_path=DPPath(dp_h5, "a") + ) + self.pt_model.compute_or_load_stat( + lambda: self.pt_sampled, stat_file_path=DPPath(pt_h5, "a") + ) + self.pt_expt_model.compute_or_load_stat( + lambda: self.np_sampled, stat_file_path=DPPath(pe_h5, "a") + ) + + # Save the computed serializations as reference + dp_ser_computed = self.dp_model.serialize() + pt_ser_computed = self.pt_model.serialize() + pe_ser_computed = self.pt_expt_model.serialize() + + # 2. Build fresh models from the same initial weights + dp_model2 = get_model_dp(self._model_data) + pt_model2 = EnergyModelPT.deserialize(dp_model2.serialize()) + pe_model2 = EnergyModelPTExpt.deserialize(dp_model2.serialize()) + + # 3. Load stats from file (should NOT call the sampled func) + def raise_error(): + raise RuntimeError("Should load from file, not recompute") + + dp_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(dp_h5, "a") + ) + pt_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(pt_h5, "a") + ) + pe_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(pe_h5, "a") + ) + + # 4. Loaded models should match the computed ones + dp_ser_loaded = dp_model2.serialize() + pt_ser_loaded = pt_model2.serialize() + pe_ser_loaded = pe_model2.serialize() + compare_variables_recursive(dp_ser_computed, dp_ser_loaded) + compare_variables_recursive(pt_ser_computed, pt_ser_loaded) + compare_variables_recursive(pe_ser_computed, pe_ser_loaded) + + # 5. Cross-backend consistency after loading + compare_variables_recursive(dp_ser_loaded, pt_ser_loaded) + compare_variables_recursive(dp_ser_loaded, pe_ser_loaded) diff --git a/source/tests/consistent/model/test_frozen.py b/source/tests/consistent/model/test_frozen.py index d7dfcfe735..d2c33f3cd9 100644 --- a/source/tests/consistent/model/test_frozen.py +++ b/source/tests/consistent/model/test_frozen.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import glob import os +import tempfile import unittest from typing import ( Any, @@ -57,6 +59,10 @@ def tearDownModule() -> None: os.remove(model_file) except FileNotFoundError: pass + # Clean up temporary .pb files created by TF FrozenModel + # (tempfile.NamedTemporaryFile(suffix=".pb", dir=os.curdir, delete=False)) + for tmp_pb in glob.glob(tempfile.gettempprefix() + "*.pb"): + os.remove(tmp_pb) @parameterized((pt_model, tf_model, dp_model)) diff --git a/source/tests/consistent/model/test_polar.py b/source/tests/consistent/model/test_polar.py index 0b9b94a599..93e696596e 100644 --- a/source/tests/consistent/model/test_polar.py +++ b/source/tests/consistent/model/test_polar.py @@ -6,8 +6,18 @@ import numpy as np +from deepmd.dpmodel.common import ( + to_numpy_array, +) from deepmd.dpmodel.model.model import get_model as get_model_dp from deepmd.dpmodel.model.polar_model import PolarModel as PolarModelDP +from deepmd.dpmodel.utils.nlist import ( + build_neighbor_list, + extend_coord_with_ghosts, +) +from deepmd.dpmodel.utils.region import ( + normalize_coord, +) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) @@ -15,16 +25,21 @@ from ..common import ( INSTALLED_JAX, INSTALLED_PT, + INSTALLED_PT_EXPT, INSTALLED_TF, CommonTest, + parameterized, ) from .common import ( ModelTest, + compare_variables_recursive, ) if INSTALLED_PT: from deepmd.pt.model.model import get_model as get_model_pt from deepmd.pt.model.model.polar_model import PolarModel as PolarModelPT + from deepmd.pt.utils.utils import to_numpy_array as torch_to_numpy + from deepmd.pt.utils.utils import to_torch_tensor as numpy_to_torch else: PolarModelPT = None if INSTALLED_TF: @@ -36,6 +51,11 @@ from deepmd.jax.model.polar_model import PolarModel as PolarModelJAX else: PolarModelJAX = None +if INSTALLED_PT_EXPT: + from deepmd.pt_expt.common import to_torch_array as pt_expt_numpy_to_torch + from deepmd.pt_expt.model import PolarModel as PolarModelPTExpt +else: + PolarModelPTExpt = None from deepmd.utils.argcheck import ( model_args, ) @@ -71,6 +91,7 @@ def data(self) -> dict: tf_class = PolarModelTF dp_class = PolarModelDP pt_class = PolarModelPT + pt_expt_class = PolarModelPTExpt jax_class = PolarModelJAX args = model_args() atol = 1e-8 @@ -84,6 +105,8 @@ def get_reference_backend(self): return self.RefBackend.PT if not self.skip_tf: return self.RefBackend.TF + if not self.skip_pt_expt and self.pt_expt_class is not None: + return self.RefBackend.PT_EXPT if not self.skip_dp: return self.RefBackend.DP raise ValueError("No available reference") @@ -105,6 +128,9 @@ def pass_data_to_cls(self, cls, data) -> Any: model = get_model_pt(data) model.atomic_model.out_bias.uniform_() return model + elif cls is PolarModelPTExpt: + dp_model = get_model_dp(data) + return PolarModelPTExpt.deserialize(dp_model.serialize()) elif cls is PolarModelJAX: return get_model_jax(data) return cls(**data, **self.additional_data) @@ -171,6 +197,15 @@ def eval_pt(self, pt_obj: Any) -> Any: self.box, ) + def eval_pt_expt(self, pt_expt_obj: Any) -> Any: + return self.eval_pt_expt_model( + pt_expt_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + def eval_jax(self, jax_obj: Any) -> Any: return self.eval_jax_model( jax_obj, @@ -190,6 +225,7 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: elif backend in { self.RefBackend.DP, self.RefBackend.PT, + self.RefBackend.PT_EXPT, self.RefBackend.JAX, }: return ( @@ -209,3 +245,1305 @@ def test_atom_exclude_types(self): tf_obj = self.tf_class.deserialize(data, suffix=self.unique_id) pt_obj = self.pt_class.deserialize(data) self.assertEqual(tf_obj.get_sel_type(), pt_obj.get_sel_type()) + + +@unittest.skipUnless(INSTALLED_PT and INSTALLED_PT_EXPT, "PyTorch is not installed") +class TestPolarModelAPIs(unittest.TestCase): + """Test consistency of model-level APIs between pt and dpmodel backends. + + Both models are constructed from the same serialized weights + (dpmodel -> serialize -> pt deserialize) so that numerical outputs + can be compared directly. + """ + + def setUp(self) -> None: + data = model_args().normalize_value( + { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20], + "rcut_smth": 1.8, + "rcut": 6.0, + "neuron": [2, 4, 8], + "resnet_dt": False, + "axis_neuron": 8, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "type": "polar", + "neuron": [4, 4, 4], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + "numb_fparam": 2, + "numb_aparam": 3, + "default_fparam": [0.5, -0.3], + }, + }, + trim_pattern="_*", + ) + # Build dpmodel first, then deserialize into pt/pt_expt to share weights + self.dp_model = get_model_dp(data) + serialized = self.dp_model.serialize() + self.pt_model = PolarModelPT.deserialize(serialized) + self.pt_expt_model = PolarModelPTExpt.deserialize(serialized) + + # Coords / atype / box + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, -1, 3) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32).reshape(1, -1) + self.box = np.array( + [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, 9) + + # Build extended coords + nlist for lower-level calls + rcut = 6.0 + nframes, nloc = self.atype.shape[:2] + coord_normalized = normalize_coord( + self.coords.reshape(nframes, nloc, 3), + self.box.reshape(nframes, 3, 3), + ) + extended_coord, extended_atype, mapping = extend_coord_with_ghosts( + coord_normalized, self.atype, self.box, rcut + ) + nlist = build_neighbor_list( + extended_coord, + extended_atype, + nloc, + rcut, + [20, 20], + distinguish_types=True, + ) + self.extended_coord = extended_coord.reshape(nframes, -1, 3) + self.extended_atype = extended_atype + self.mapping = mapping + self.nlist = nlist + + # aparam for forward evaluation (1 frame, 6 atoms, 3 aparam) + rng = np.random.default_rng(42) + self.eval_aparam = rng.normal(size=(1, nloc, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + def test_translated_output_def(self) -> None: + """translated_output_def should return the same keys on dp, pt, and pt_expt.""" + dp_def = self.dp_model.translated_output_def() + pt_def = self.pt_model.translated_output_def() + pt_expt_def = self.pt_expt_model.translated_output_def() + self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + self.assertEqual(set(dp_def.keys()), set(pt_expt_def.keys())) + for key in dp_def: + self.assertEqual(dp_def[key].shape, pt_def[key].shape) + self.assertEqual(dp_def[key].shape, pt_expt_def[key].shape) + + def test_get_descriptor(self) -> None: + """get_descriptor should return a non-None object on both backends.""" + self.assertIsNotNone(self.dp_model.get_descriptor()) + self.assertIsNotNone(self.pt_model.get_descriptor()) + + def test_get_fitting_net(self) -> None: + """get_fitting_net should return a non-None object on both backends.""" + self.assertIsNotNone(self.dp_model.get_fitting_net()) + self.assertIsNotNone(self.pt_model.get_fitting_net()) + + def test_get_out_bias(self) -> None: + """get_out_bias should return numerically equal values on dp and pt. + + PolarModel's apply_out_stat applies diagonal bias with scale, + so the bias storage should be consistent across backends. + """ + dp_bias = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias = torch_to_numpy(self.pt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias, pt_bias, rtol=1e-10, atol=1e-10) + # Verify shape: (n_output_keys x ntypes x 9) for polar + self.assertEqual(dp_bias.shape[1], 2) # ntypes + self.assertGreater(dp_bias.shape[0], 0) # at least one output key + + def test_set_out_bias(self) -> None: + """set_out_bias should update the bias on both backends.""" + dp_bias = to_numpy_array(self.dp_model.get_out_bias()) + new_bias = dp_bias + 1.0 + # dp + self.dp_model.set_out_bias(new_bias) + np.testing.assert_allclose( + to_numpy_array(self.dp_model.get_out_bias()), + new_bias, + rtol=1e-10, + atol=1e-10, + ) + # pt + self.pt_model.set_out_bias(numpy_to_torch(new_bias)) + np.testing.assert_allclose( + torch_to_numpy(self.pt_model.get_out_bias()), + new_bias, + rtol=1e-10, + atol=1e-10, + ) + + def test_model_output_def(self) -> None: + """model_output_def should return the same keys and shapes on dp and pt.""" + dp_def = self.dp_model.model_output_def().get_data() + pt_def = self.pt_model.model_output_def().get_data() + self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + for key in dp_def: + self.assertEqual(dp_def[key].shape, pt_def[key].shape) + + def test_model_output_type(self) -> None: + """model_output_type should return the same list on dp and pt.""" + self.assertEqual( + self.dp_model.model_output_type(), + self.pt_model.model_output_type(), + ) + + def test_do_grad_r(self) -> None: + """do_grad_r should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.do_grad_r("polarizability"), + self.pt_model.do_grad_r("polarizability"), + ) + self.assertFalse(self.dp_model.do_grad_r("polarizability")) + + def test_do_grad_c(self) -> None: + """do_grad_c should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.do_grad_c("polarizability"), + self.pt_model.do_grad_c("polarizability"), + ) + self.assertFalse(self.dp_model.do_grad_c("polarizability")) + + def test_get_rcut(self) -> None: + """get_rcut should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_rcut(), self.pt_model.get_rcut()) + self.assertAlmostEqual(self.dp_model.get_rcut(), 6.0) + + def test_get_type_map(self) -> None: + """get_type_map should return the same list on dp and pt.""" + self.assertEqual(self.dp_model.get_type_map(), self.pt_model.get_type_map()) + self.assertEqual(self.dp_model.get_type_map(), ["O", "H"]) + + def test_get_sel(self) -> None: + """get_sel should return the same list on dp and pt.""" + self.assertEqual(self.dp_model.get_sel(), self.pt_model.get_sel()) + self.assertEqual(self.dp_model.get_sel(), [20, 20]) + + def test_get_nsel(self) -> None: + """get_nsel should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_nsel(), self.pt_model.get_nsel()) + self.assertEqual(self.dp_model.get_nsel(), 40) + + def test_get_nnei(self) -> None: + """get_nnei should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_nnei(), self.pt_model.get_nnei()) + self.assertEqual(self.dp_model.get_nnei(), 40) + + def test_mixed_types(self) -> None: + """mixed_types should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.mixed_types(), self.pt_model.mixed_types()) + # se_e2_a is not mixed-types + self.assertFalse(self.dp_model.mixed_types()) + + def test_has_message_passing(self) -> None: + """has_message_passing should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.has_message_passing(), + self.pt_model.has_message_passing(), + ) + self.assertFalse(self.dp_model.has_message_passing()) + + def test_need_sorted_nlist_for_lower(self) -> None: + """need_sorted_nlist_for_lower should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.need_sorted_nlist_for_lower(), + self.pt_model.need_sorted_nlist_for_lower(), + ) + self.assertFalse(self.dp_model.need_sorted_nlist_for_lower()) + + def test_get_dim_fparam(self) -> None: + """get_dim_fparam should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_dim_fparam(), self.pt_model.get_dim_fparam()) + self.assertEqual(self.dp_model.get_dim_fparam(), 2) + + def test_get_dim_aparam(self) -> None: + """get_dim_aparam should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_dim_aparam(), self.pt_model.get_dim_aparam()) + self.assertEqual(self.dp_model.get_dim_aparam(), 3) + + def test_get_sel_type(self) -> None: + """get_sel_type should return the same list on dp and pt.""" + self.assertEqual(self.dp_model.get_sel_type(), self.pt_model.get_sel_type()) + self.assertEqual(self.dp_model.get_sel_type(), [0, 1]) + + def test_is_aparam_nall(self) -> None: + """is_aparam_nall should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.is_aparam_nall(), self.pt_model.is_aparam_nall()) + self.assertFalse(self.dp_model.is_aparam_nall()) + + def test_get_model_def_script(self) -> None: + """get_model_def_script should return the same value on dp, pt, and pt_expt.""" + dp_val = self.dp_model.get_model_def_script() + pt_val = self.pt_model.get_model_def_script() + pe_val = self.pt_expt_model.get_model_def_script() + self.assertEqual(dp_val, pt_val) + self.assertEqual(dp_val, pe_val) + + def test_get_min_nbor_dist(self) -> None: + """get_min_nbor_dist should return the same value on dp, pt, and pt_expt.""" + dp_val = self.dp_model.get_min_nbor_dist() + pt_val = self.pt_model.get_min_nbor_dist() + pe_val = self.pt_expt_model.get_min_nbor_dist() + self.assertEqual(dp_val, pt_val) + self.assertEqual(dp_val, pe_val) + + def test_set_case_embd(self) -> None: + """set_case_embd should produce consistent results across backends. + + Also verifies that different case indices produce different outputs, + confirming the embedding is actually used. + """ + from deepmd.utils.argcheck import ( + model_args, + ) + + # Build a model with dim_case_embd > 0 + data = model_args().normalize_value( + { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "type": "polar", + "neuron": [5, 5], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + "dim_case_embd": 3, + }, + }, + trim_pattern="_*", + ) + dp_model = get_model_dp(data) + serialized = dp_model.serialize() + pt_model = PolarModelPT.deserialize(serialized) + pe_model = PolarModelPTExpt.deserialize(serialized) + + def _eval(case_idx): + dp_model.set_case_embd(case_idx) + pt_model.set_case_embd(case_idx) + pe_model.set_case_embd(case_idx) + dp_ret = dp_model(self.coords, self.atype, box=self.box) + pt_ret = { + k: torch_to_numpy(v) + for k, v in pt_model( + numpy_to_torch(self.coords), + numpy_to_torch(self.atype), + box=numpy_to_torch(self.box), + ).items() + } + coord_t = pt_expt_numpy_to_torch(self.coords) + coord_t.requires_grad_(True) + pe_ret = { + k: v.detach().cpu().numpy() + for k, v in pe_model( + coord_t, + pt_expt_numpy_to_torch(self.atype), + box=pt_expt_numpy_to_torch(self.box), + ).items() + } + return dp_ret, pt_ret, pe_ret + + dp0, pt0, pe0 = _eval(0) + dp1, pt1, pe1 = _eval(1) + + # Cross-backend consistency for each case index + for key in ("polar", "global_polar"): + np.testing.assert_allclose( + dp0[key], + pt0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 0: dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp0[key], + pe0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 0: dp vs pt_expt mismatch in {key}", + ) + np.testing.assert_allclose( + dp1[key], + pt1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 1: dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp1[key], + pe1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 1: dp vs pt_expt mismatch in {key}", + ) + # Different case indices should produce different outputs + self.assertFalse( + np.allclose(dp0["global_polar"], dp1["global_polar"]), + "set_case_embd(0) and set_case_embd(1) produced the same global_polar", + ) + + def test_atomic_output_def(self) -> None: + """atomic_output_def should return the same keys and shapes on dp and pt.""" + dp_def = self.dp_model.atomic_output_def() + pt_def = self.pt_model.atomic_output_def() + self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + for key in dp_def.keys(): + self.assertEqual(dp_def[key].shape, pt_def[key].shape) + + def test_format_nlist(self) -> None: + """format_nlist should produce the same result on dp and pt.""" + dp_nlist = self.dp_model.format_nlist( + self.extended_coord, + self.extended_atype, + self.nlist, + ) + pt_nlist = torch_to_numpy( + self.pt_model.format_nlist( + numpy_to_torch(self.extended_coord), + numpy_to_torch(self.extended_atype), + numpy_to_torch(self.nlist), + ) + ) + np.testing.assert_equal(dp_nlist, pt_nlist) + + def test_forward_common_atomic(self) -> None: + """forward_common_atomic should produce consistent results on dp and pt. + + Compares at the atomic_model level, where both backends define this method. + """ + dp_ret = self.dp_model.atomic_model.forward_common_atomic( + self.extended_coord, + self.extended_atype, + self.nlist, + mapping=self.mapping, + aparam=self.eval_aparam, + ) + pt_ret = self.pt_model.atomic_model.forward_common_atomic( + numpy_to_torch(self.extended_coord), + numpy_to_torch(self.extended_atype), + numpy_to_torch(self.nlist), + mapping=numpy_to_torch(self.mapping), + aparam=numpy_to_torch(self.eval_aparam), + ) + # Compare the common keys + common_keys = set(dp_ret.keys()) & set(pt_ret.keys()) + self.assertTrue(len(common_keys) > 0) + for key in common_keys: + if dp_ret[key] is not None and pt_ret[key] is not None: + np.testing.assert_allclose( + dp_ret[key], + torch_to_numpy(pt_ret[key]), + rtol=1e-10, + atol=1e-10, + err_msg=f"Mismatch in forward_common_atomic key '{key}'", + ) + + def test_has_default_fparam(self) -> None: + """has_default_fparam should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.has_default_fparam(), + self.pt_model.has_default_fparam(), + ) + self.assertTrue(self.dp_model.has_default_fparam()) + + def test_get_default_fparam(self) -> None: + """get_default_fparam should return consistent values on dp and pt.""" + dp_val = self.dp_model.get_default_fparam() + pt_val = self.pt_model.get_default_fparam() + np.testing.assert_allclose(dp_val, pt_val, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dp_val, [0.5, -0.3], rtol=1e-10, atol=1e-10) + + def test_change_out_bias(self) -> None: + """change_out_bias should produce consistent bias on dp, pt, and pt_expt. + + Tests both set-by-statistic and change-by-statistic modes. + Note: change_out_bias only updates the output bias, not fitting input + stats (fparam/aparam). Fitting stats are updated by compute_or_load_stat. + """ + nframes = 2 + nloc = 6 + numb_fparam = 2 + numb_aparam = 3 + rng = np.random.default_rng(123) + + # Use realistic coords (from setUp, tiled for 2 frames) + coords_2f = np.tile(self.coords, (nframes, 1, 1)) # (2, 6, 3) + atype_2f = np.array([[0, 0, 1, 1, 1, 1], [0, 1, 1, 0, 1, 1]], dtype=np.int32) + box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) + natoms_data = np.array([[6, 6, 2, 4], [6, 6, 2, 4]], dtype=np.int32) + polar_data = rng.normal(size=(nframes, 9)).astype(GLOBAL_NP_FLOAT_PRECISION) + fparam_data = rng.normal(size=(nframes, numb_fparam)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + aparam_data = rng.normal(size=(nframes, nloc, numb_aparam)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + # dpmodel stat data (numpy) + dp_merged = [ + { + "coord": coords_2f, + "atype": atype_2f, + "atype_ext": atype_2f, + "box": box_2f, + "natoms": natoms_data, + "polarizability": polar_data, + "find_polarizability": np.float32(1.0), + "fparam": fparam_data, + "aparam": aparam_data, + } + ] + # pt stat data (torch tensors) + pt_merged = [ + { + "coord": numpy_to_torch(coords_2f), + "atype": numpy_to_torch(atype_2f), + "atype_ext": numpy_to_torch(atype_2f), + "box": numpy_to_torch(box_2f), + "natoms": numpy_to_torch(natoms_data), + "polarizability": numpy_to_torch(polar_data), + "find_polarizability": np.float32(1.0), + "fparam": numpy_to_torch(fparam_data), + "aparam": numpy_to_torch(aparam_data), + } + ] + # pt_expt stat data (numpy, same as dp) + pe_merged = dp_merged + + # Save initial (zero) bias + dp_bias_init = to_numpy_array(self.dp_model.get_out_bias()).copy() + + # --- Test "set-by-statistic" mode --- + self.dp_model.change_out_bias(dp_merged, bias_adjust_mode="set-by-statistic") + self.pt_model.change_out_bias(pt_merged, bias_adjust_mode="set-by-statistic") + self.pt_expt_model.change_out_bias( + pe_merged, bias_adjust_mode="set-by-statistic" + ) + + # Verify out bias consistency + dp_bias = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias = torch_to_numpy(self.pt_model.get_out_bias()) + pe_bias = to_numpy_array(self.pt_expt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias, pt_bias, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dp_bias, pe_bias, rtol=1e-10, atol=1e-10) + self.assertFalse( + np.allclose(dp_bias, dp_bias_init), + "set-by-statistic did not change the bias from initial values", + ) + + # --- Test "change-by-statistic" mode --- + dp_bias_before = dp_bias.copy() + self.dp_model.change_out_bias(dp_merged, bias_adjust_mode="change-by-statistic") + self.pt_model.change_out_bias(pt_merged, bias_adjust_mode="change-by-statistic") + self.pt_expt_model.change_out_bias( + pe_merged, bias_adjust_mode="change-by-statistic" + ) + + # Verify out bias consistency + dp_bias2 = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias2 = torch_to_numpy(self.pt_model.get_out_bias()) + pe_bias2 = to_numpy_array(self.pt_expt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias2, pt_bias2, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dp_bias2, pe_bias2, rtol=1e-10, atol=1e-10) + self.assertFalse( + np.allclose(dp_bias2, dp_bias_before), + "change-by-statistic did not further change the bias", + ) + + def test_change_type_map(self) -> None: + """change_type_map should produce consistent results on dp and pt. + + Uses a DPA1 (se_atten) descriptor since se_e2_a does not support + change_type_map (non-mixed-types descriptors raise NotImplementedError). + """ + from deepmd.utils.argcheck import model_args as model_args_fn + + data = model_args_fn().normalize_value( + { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_atten", + "sel": 20, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "seed": 1, + "attn": 6, + "attn_layer": 0, + }, + "fitting_net": { + "type": "polar", + "neuron": [4, 4, 4], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + }, + }, + trim_pattern="_*", + ) + dp_model = get_model_dp(data) + pt_model = PolarModelPT.deserialize(dp_model.serialize()) + + # Set non-zero out_bias so the swap is non-trivial + dp_bias_orig = to_numpy_array(dp_model.get_out_bias()).copy() + new_bias = dp_bias_orig.copy() + new_bias[:, 0, :] = 1.5 # type 0 ("O") + new_bias[:, 1, :] = -3.7 # type 1 ("H") + dp_model.set_out_bias(new_bias) + pt_model.set_out_bias(numpy_to_torch(new_bias)) + + new_type_map = ["H", "O"] + dp_model.change_type_map(new_type_map) + pt_model.change_type_map(new_type_map) + + # Both should have the new type_map + self.assertEqual(dp_model.get_type_map(), new_type_map) + self.assertEqual(pt_model.get_type_map(), new_type_map) + + # Out_bias should be reordered consistently between backends + dp_bias_new = to_numpy_array(dp_model.get_out_bias()) + pt_bias_new = torch_to_numpy(pt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias_new, pt_bias_new, rtol=1e-10, atol=1e-10) + + # Verify the reorder is correct: old type 0 -> new type 1, old type 1 -> new type 0 + np.testing.assert_allclose( + dp_bias_new[:, 0, :], + new_bias[:, 1, :], + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + dp_bias_new[:, 1, :], + new_bias[:, 0, :], + rtol=1e-10, + atol=1e-10, + ) + + def test_change_type_map_extend_stat(self) -> None: + """change_type_map with model_with_new_type_stat should propagate stats consistently across dp, pt, and pt_expt. + + Verifies that the model-level change_type_map correctly unwraps + model_with_new_type_stat.atomic_model before forwarding to the + atomic model. + """ + from deepmd.utils.argcheck import model_args as model_args_fn + + small_tm = ["O", "H"] + large_tm = ["O", "H", "Li"] + + small_data = model_args_fn().normalize_value( + { + "type_map": small_tm, + "descriptor": { + "type": "se_atten", + "sel": 20, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "seed": 1, + "attn": 6, + "attn_layer": 0, + }, + "fitting_net": { + "type": "polar", + "neuron": [4, 4, 4], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + }, + }, + trim_pattern="_*", + ) + large_data = model_args_fn().normalize_value( + { + "type_map": large_tm, + "descriptor": { + "type": "se_atten", + "sel": 20, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "seed": 2, + "attn": 6, + "attn_layer": 0, + }, + "fitting_net": { + "type": "polar", + "neuron": [4, 4, 4], + "resnet_dt": True, + "precision": "float64", + "seed": 2, + }, + }, + trim_pattern="_*", + ) + + dp_small = get_model_dp(small_data) + dp_large = get_model_dp(large_data) + + # Set distinguishable random stats on the large model's descriptor + rng = np.random.default_rng(42) + desc_large = dp_large.get_descriptor() + mean_large, std_large = desc_large.get_stat_mean_and_stddev() + mean_rand = rng.random(size=to_numpy_array(mean_large).shape) + std_rand = rng.random(size=to_numpy_array(std_large).shape) + desc_large.set_stat_mean_and_stddev(mean_rand, std_rand) + + # Build pt and pt_expt models from dp serialization + pt_small = PolarModelPT.deserialize(dp_small.serialize()) + pt_large = PolarModelPT.deserialize(dp_large.serialize()) + pt_expt_small = PolarModelPTExpt.deserialize(dp_small.serialize()) + pt_expt_large = PolarModelPTExpt.deserialize(dp_large.serialize()) + + # Extend type map with model_with_new_type_stat at the model level + dp_small.change_type_map(large_tm, model_with_new_type_stat=dp_large) + pt_small.change_type_map(large_tm, model_with_new_type_stat=pt_large) + pt_expt_small.change_type_map(large_tm, model_with_new_type_stat=pt_expt_large) + + # Descriptor stats should be consistent across backends + dp_mean, dp_std = dp_small.get_descriptor().get_stat_mean_and_stddev() + pt_mean, pt_std = pt_small.get_descriptor().get_stat_mean_and_stddev() + pt_expt_mean, pt_expt_std = ( + pt_expt_small.get_descriptor().get_stat_mean_and_stddev() + ) + np.testing.assert_allclose( + to_numpy_array(dp_mean), + torch_to_numpy(pt_mean), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + to_numpy_array(dp_std), + torch_to_numpy(pt_std), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + to_numpy_array(dp_mean), + to_numpy_array(pt_expt_mean), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + to_numpy_array(dp_std), + to_numpy_array(pt_expt_std), + rtol=1e-10, + atol=1e-10, + ) + + def test_update_sel(self) -> None: + """update_sel should return the same result on dp and pt.""" + from unittest.mock import ( + patch, + ) + + from deepmd.dpmodel.model.dp_model import DPModelCommon as DPModelCommonDP + from deepmd.pt.model.model.dp_model import DPModelCommon as DPModelCommonPT + + mock_min_nbor_dist = 0.5 + mock_sel = [10, 20] + local_jdata = { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": "auto", + "rcut_smth": 0.50, + "rcut": 6.00, + }, + "fitting_net": { + "type": "polar", + "neuron": [4, 4, 4], + }, + } + type_map = ["O", "H"] + + with patch( + "deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat", + return_value=(mock_min_nbor_dist, mock_sel), + ): + dp_result, dp_min_dist = DPModelCommonDP.update_sel( + None, type_map, local_jdata + ) + + with patch( + "deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat", + return_value=(mock_min_nbor_dist, mock_sel), + ): + pt_result, pt_min_dist = DPModelCommonPT.update_sel( + None, type_map, local_jdata + ) + + self.assertEqual(dp_result, pt_result) + self.assertEqual(dp_min_dist, pt_min_dist) + # Verify sel was actually updated (not still "auto") + self.assertIsInstance(dp_result["descriptor"]["sel"], list) + self.assertNotEqual(dp_result["descriptor"]["sel"], "auto") + + def test_get_ntypes(self) -> None: + """get_ntypes should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_ntypes(), self.pt_model.get_ntypes()) + self.assertEqual(self.dp_model.get_ntypes(), 2) + + def test_compute_or_load_out_stat(self) -> None: + """compute_or_load_out_stat should produce consistent bias on dp and pt. + + Tests both the compute path (from data) and the load path (from file). + PolarModel's apply_out_stat applies diagonal bias with scale, + so the stored bias should be consistent across backends. + """ + import tempfile + from pathlib import ( + Path, + ) + + import h5py + + from deepmd.utils.path import ( + DPPath, + ) + + nframes = 2 + coords_2f = np.tile(self.coords, (nframes, 1, 1)) + atype_2f = np.array([[0, 0, 1, 1, 1, 1], [0, 1, 1, 0, 1, 1]], dtype=np.int32) + box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) + natoms_data = np.array([[6, 6, 2, 4], [6, 6, 2, 4]], dtype=np.int32) + polar_data = ( + np.random.default_rng(42) + .normal(size=(nframes, 9)) + .astype(GLOBAL_NP_FLOAT_PRECISION) + ) + + dp_merged = [ + { + "coord": coords_2f, + "atype": atype_2f, + "atype_ext": atype_2f, + "box": box_2f, + "natoms": natoms_data, + "polarizability": polar_data, + "find_polarizability": np.float32(1.0), + } + ] + pt_merged = [ + { + "coord": numpy_to_torch(coords_2f), + "atype": numpy_to_torch(atype_2f), + "atype_ext": numpy_to_torch(atype_2f), + "box": numpy_to_torch(box_2f), + "natoms": numpy_to_torch(natoms_data), + "polarizability": numpy_to_torch(polar_data), + "find_polarizability": np.float32(1.0), + } + ] + + # Verify bias is initially identical + dp_bias_before = to_numpy_array(self.dp_model.get_out_bias()).copy() + pt_bias_before = torch_to_numpy(self.pt_model.get_out_bias()).copy() + np.testing.assert_allclose( + dp_bias_before, pt_bias_before, rtol=1e-10, atol=1e-10 + ) + + with tempfile.TemporaryDirectory() as tmpdir: + # Create separate h5 files for dp and pt + dp_h5 = str((Path(tmpdir) / "dp_stat.h5").resolve()) + pt_h5 = str((Path(tmpdir) / "pt_stat.h5").resolve()) + with h5py.File(dp_h5, "w"): + pass + with h5py.File(pt_h5, "w"): + pass + dp_stat_path = DPPath(dp_h5, "a") + pt_stat_path = DPPath(pt_h5, "a") + + # 1. Compute stats and save to file + self.dp_model.atomic_model.compute_or_load_out_stat( + dp_merged, stat_file_path=dp_stat_path + ) + self.pt_model.atomic_model.compute_or_load_out_stat( + pt_merged, stat_file_path=pt_stat_path + ) + + dp_bias_after = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias_after = torch_to_numpy(self.pt_model.get_out_bias()) + np.testing.assert_allclose( + dp_bias_after, pt_bias_after, rtol=1e-10, atol=1e-10 + ) + + # 2. Verify both backends saved the same file content + with h5py.File(dp_h5, "r") as dp_f, h5py.File(pt_h5, "r") as pt_f: + dp_keys = sorted(dp_f.keys()) + pt_keys = sorted(pt_f.keys()) + self.assertEqual(dp_keys, pt_keys) + for key in dp_keys: + np.testing.assert_allclose( + np.array(dp_f[key]), + np.array(pt_f[key]), + rtol=1e-10, + atol=1e-10, + err_msg=f"Stat file content mismatch for key {key}", + ) + + # 3. Reset biases to zero, then load from file + zero_bias = np.zeros_like(dp_bias_after) + self.dp_model.set_out_bias(zero_bias) + self.pt_model.set_out_bias(numpy_to_torch(zero_bias)) + + # Use a callable that raises to ensure it loads from file, not recomputes + def raise_error(): + raise RuntimeError("Should not recompute — should load from file") + + self.dp_model.atomic_model.compute_or_load_out_stat( + raise_error, stat_file_path=dp_stat_path + ) + self.pt_model.atomic_model.compute_or_load_out_stat( + raise_error, stat_file_path=pt_stat_path + ) + + dp_bias_loaded = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias_loaded = torch_to_numpy(self.pt_model.get_out_bias()) + + # Loaded biases should match between backends + np.testing.assert_allclose( + dp_bias_loaded, pt_bias_loaded, rtol=1e-10, atol=1e-10 + ) + # Loaded biases should match the originally computed biases + np.testing.assert_allclose( + dp_bias_loaded, dp_bias_after, rtol=1e-10, atol=1e-10 + ) + + def test_get_observed_type_list(self) -> None: + """get_observed_type_list should be consistent across dp, pt, pt_expt. + + Uses mock data containing only type 0 ("O") so that type 1 ("H") is + unobserved and should be absent from the returned list. + """ + nframes = 2 + natoms = 6 + # All atoms are type 0 — type 1 is unobserved + atype_2f = np.zeros((nframes, natoms), dtype=np.int32) + coords_2f = np.tile(self.coords, (nframes, 1, 1)) + box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) + natoms_data = np.array([[natoms, natoms, natoms, 0]] * nframes, dtype=np.int32) + polar_data = ( + np.random.default_rng(42) + .normal(size=(nframes, 9)) + .astype(GLOBAL_NP_FLOAT_PRECISION) + ) + + dp_merged = [ + { + "coord": coords_2f, + "atype": atype_2f, + "atype_ext": atype_2f, + "box": box_2f, + "natoms": natoms_data, + "polarizability": polar_data, + "find_polarizability": np.float32(1.0), + } + ] + pt_merged = [ + { + "coord": numpy_to_torch(coords_2f), + "atype": numpy_to_torch(atype_2f), + "atype_ext": numpy_to_torch(atype_2f), + "box": numpy_to_torch(box_2f), + "natoms": numpy_to_torch(natoms_data), + "polarizability": numpy_to_torch(polar_data), + "find_polarizability": np.float32(1.0), + } + ] + + self.dp_model.atomic_model.compute_or_load_out_stat(dp_merged) + self.pt_model.atomic_model.compute_or_load_out_stat(pt_merged) + self.pt_expt_model.atomic_model.compute_or_load_out_stat(dp_merged) + + dp_observed = self.dp_model.get_observed_type_list() + pt_observed = self.pt_model.get_observed_type_list() + pe_observed = self.pt_expt_model.get_observed_type_list() + + self.assertEqual(dp_observed, pt_observed) + self.assertEqual(dp_observed, pe_observed) + # Only type 0 ("O") should be observed + self.assertEqual(dp_observed, ["O"]) + + +@parameterized( + (([], []), ([[0, 1]], [1])), # (pair_exclude_types, atom_exclude_types) + (False, True), # fparam_in_data +) +@unittest.skipUnless(INSTALLED_PT and INSTALLED_PT_EXPT, "PT and PT_EXPT are required") +class TestPolarComputeOrLoadStat(unittest.TestCase): + """Test that compute_or_load_stat produces identical statistics on dp, pt, and pt_expt. + + Covers descriptor stats (dstd), fitting stats (fparam, aparam), and output bias. + Parameterized over exclusion types and whether fparam is explicitly provided or + injected via default_fparam. + """ + + def setUp(self) -> None: + (pair_exclude_types, atom_exclude_types), self.fparam_in_data = self.param + data = model_args().normalize_value( + { + "type_map": ["O", "H"], + "pair_exclude_types": pair_exclude_types, + "atom_exclude_types": atom_exclude_types, + "descriptor": { + "type": "dpa3", + "repflow": { + "n_dim": 20, + "e_dim": 10, + "a_dim": 8, + "nlayers": 3, + "e_rcut": 6.0, + "e_rcut_smth": 5.0, + "e_sel": 10, + "a_rcut": 4.0, + "a_rcut_smth": 3.5, + "a_sel": 8, + "axis_neuron": 4, + "update_angle": True, + "update_style": "res_residual", + "update_residual": 0.1, + "update_residual_init": "const", + }, + "precision": "float64", + "seed": 1, + }, + "fitting_net": { + "type": "polar", + "neuron": [10, 10], + "precision": "float64", + "seed": 1, + "numb_fparam": 2, + "default_fparam": [0.5, -0.3], + "numb_aparam": 3, + }, + }, + trim_pattern="_*", + ) + + # Save data for reuse in load-from-file test + self._model_data = data + + # Build dp model, then deserialize into pt and pt_expt to share weights + self.dp_model = get_model_dp(data) + serialized = self.dp_model.serialize() + self.pt_model = PolarModelPT.deserialize(serialized) + self.pt_expt_model = PolarModelPTExpt.deserialize(serialized) + + # Test coords / atype / box for forward evaluation + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 0.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, -1, 3) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32).reshape(1, -1) + self.box = np.array( + [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, 9) + + # Mock training data for compute_or_load_stat + natoms = 6 + nframes = 3 + rng = np.random.default_rng(42) + coords_stat = rng.normal(size=(nframes, natoms, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + atype_stat = np.array([[0, 0, 1, 1, 1, 1]] * nframes, dtype=np.int32) + box_stat = np.tile( + np.eye(3, dtype=GLOBAL_NP_FLOAT_PRECISION).reshape(1, 3, 3) * 13.0, + (nframes, 1, 1), + ) + natoms_stat = np.array([[natoms, natoms, 2, 4]] * nframes, dtype=np.int32) + polar_stat = rng.normal(size=(nframes, 9)).astype(GLOBAL_NP_FLOAT_PRECISION) + aparam_stat = rng.normal(size=(nframes, natoms, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + # dp / pt_expt sample (numpy) + np_sample = { + "coord": coords_stat, + "atype": atype_stat, + "atype_ext": atype_stat, + "box": box_stat, + "natoms": natoms_stat, + "polarizability": polar_stat, + "find_polarizability": np.float32(1.0), + "aparam": aparam_stat, + } + # pt sample (torch tensors) + pt_sample = { + "coord": numpy_to_torch(coords_stat), + "atype": numpy_to_torch(atype_stat), + "atype_ext": numpy_to_torch(atype_stat), + "box": numpy_to_torch(box_stat), + "natoms": numpy_to_torch(natoms_stat), + "polarizability": numpy_to_torch(polar_stat), + "find_polarizability": np.float32(1.0), + "aparam": numpy_to_torch(aparam_stat), + } + + if self.fparam_in_data: + fparam_stat = rng.normal(size=(nframes, 2)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + np_sample["fparam"] = fparam_stat + pt_sample["fparam"] = numpy_to_torch(fparam_stat) + self.expected_fparam_avg = np.mean(fparam_stat, axis=0) + else: + # No fparam → _make_wrapped_sampler injects default_fparam + self.expected_fparam_avg = np.array([0.5, -0.3]) + + self.np_sampled = [np_sample] + self.pt_sampled = [pt_sample] + + # aparam for forward evaluation (1 frame, 6 atoms, 3 aparam) + self.eval_aparam = rng.normal(size=(1, natoms, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + def _eval_dp(self) -> dict: + return self.dp_model( + self.coords, self.atype, box=self.box, aparam=self.eval_aparam + ) + + def _eval_pt(self) -> dict: + return { + kk: torch_to_numpy(vv) + for kk, vv in self.pt_model( + numpy_to_torch(self.coords), + numpy_to_torch(self.atype), + box=numpy_to_torch(self.box), + aparam=numpy_to_torch(self.eval_aparam), + do_atomic_virial=True, + ).items() + } + + def _eval_pt_expt(self) -> dict: + coord_t = pt_expt_numpy_to_torch(self.coords) + coord_t.requires_grad_(True) + return { + k: v.detach().cpu().numpy() + for k, v in self.pt_expt_model( + coord_t, + pt_expt_numpy_to_torch(self.atype), + box=pt_expt_numpy_to_torch(self.box), + aparam=pt_expt_numpy_to_torch(self.eval_aparam), + do_atomic_virial=True, + ).items() + } + + def test_compute_stat(self) -> None: + # 1. Pre-stat forward consistency + dp_ret0 = self._eval_dp() + pt_ret0 = self._eval_pt() + pe_ret0 = self._eval_pt_expt() + for key in ("global_polar", "polar"): + np.testing.assert_allclose( + dp_ret0[key], + pt_ret0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Pre-stat dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp_ret0[key], + pe_ret0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Pre-stat dp vs pt_expt mismatch in {key}", + ) + + # 2. Run compute_or_load_stat on all three backends + self.dp_model.compute_or_load_stat(lambda: self.np_sampled) + self.pt_model.compute_or_load_stat(lambda: self.pt_sampled) + self.pt_expt_model.compute_or_load_stat(lambda: self.np_sampled) + + # 3. Serialize all three and compare @variables + dp_ser = self.dp_model.serialize() + pt_ser = self.pt_model.serialize() + pe_ser = self.pt_expt_model.serialize() + compare_variables_recursive(dp_ser, pt_ser) + compare_variables_recursive(dp_ser, pe_ser) + + # 4. Post-stat forward consistency + # PolarModel's apply_out_stat applies diagonal bias with scale, so + # output values WILL differ from pre-stat after stat computation. + dp_ret1 = self._eval_dp() + pt_ret1 = self._eval_pt() + pe_ret1 = self._eval_pt_expt() + for key in ("global_polar", "polar"): + np.testing.assert_allclose( + dp_ret1[key], + pt_ret1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Post-stat dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp_ret1[key], + pe_ret1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Post-stat dp vs pt_expt mismatch in {key}", + ) + + # 5. Non-triviality checks + fit_vars = dp_ser["fitting"]["@variables"] + # fparam stats were computed + fparam_avg = np.asarray(fit_vars["fparam_avg"]) + self.assertFalse( + np.allclose(fparam_avg, 0.0), + "fparam_avg is still zero — fparam stats were not computed", + ) + np.testing.assert_allclose( + fparam_avg, + self.expected_fparam_avg, + rtol=1e-10, + atol=1e-10, + err_msg="fparam_avg does not match expected values", + ) + # aparam stats were computed + aparam_avg = np.asarray(fit_vars["aparam_avg"]) + self.assertFalse( + np.allclose(aparam_avg, 0.0), + "aparam_avg is still zero — aparam stats were not computed", + ) + + def test_load_stat_from_file(self) -> None: + import tempfile + from pathlib import ( + Path, + ) + + import h5py + + from deepmd.utils.path import ( + DPPath, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + # Create separate stat files for each backend + dp_h5 = str((Path(tmpdir) / "dp_stat.h5").resolve()) + pt_h5 = str((Path(tmpdir) / "pt_stat.h5").resolve()) + pe_h5 = str((Path(tmpdir) / "pe_stat.h5").resolve()) + for p in (dp_h5, pt_h5, pe_h5): + with h5py.File(p, "w"): + pass + + # 1. Compute stats and save to file + self.dp_model.compute_or_load_stat( + lambda: self.np_sampled, stat_file_path=DPPath(dp_h5, "a") + ) + self.pt_model.compute_or_load_stat( + lambda: self.pt_sampled, stat_file_path=DPPath(pt_h5, "a") + ) + self.pt_expt_model.compute_or_load_stat( + lambda: self.np_sampled, stat_file_path=DPPath(pe_h5, "a") + ) + + # Save the computed serializations as reference + dp_ser_computed = self.dp_model.serialize() + pt_ser_computed = self.pt_model.serialize() + pe_ser_computed = self.pt_expt_model.serialize() + + # 2. Build fresh models from the same initial weights + dp_model2 = get_model_dp(self._model_data) + pt_model2 = PolarModelPT.deserialize(dp_model2.serialize()) + pe_model2 = PolarModelPTExpt.deserialize(dp_model2.serialize()) + + # 3. Load stats from file (should NOT call the sampled func) + def raise_error(): + raise RuntimeError("Should load from file, not recompute") + + dp_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(dp_h5, "a") + ) + pt_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(pt_h5, "a") + ) + pe_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(pe_h5, "a") + ) + + # 4. Loaded models should match the computed ones + dp_ser_loaded = dp_model2.serialize() + pt_ser_loaded = pt_model2.serialize() + pe_ser_loaded = pe_model2.serialize() + compare_variables_recursive(dp_ser_computed, dp_ser_loaded) + compare_variables_recursive(pt_ser_computed, pt_ser_loaded) + compare_variables_recursive(pe_ser_computed, pe_ser_loaded) + + # 5. Cross-backend consistency after loading + compare_variables_recursive(dp_ser_loaded, pt_ser_loaded) + compare_variables_recursive(dp_ser_loaded, pe_ser_loaded) diff --git a/source/tests/consistent/model/test_property.py b/source/tests/consistent/model/test_property.py index 33e63af98e..cf4f7f1de9 100644 --- a/source/tests/consistent/model/test_property.py +++ b/source/tests/consistent/model/test_property.py @@ -6,8 +6,18 @@ import numpy as np +from deepmd.dpmodel.common import ( + to_numpy_array, +) from deepmd.dpmodel.model.model import get_model as get_model_dp from deepmd.dpmodel.model.property_model import PropertyModel as PropertyModelDP +from deepmd.dpmodel.utils.nlist import ( + build_neighbor_list, + extend_coord_with_ghosts, +) +from deepmd.dpmodel.utils.region import ( + normalize_coord, +) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) @@ -15,15 +25,20 @@ from ..common import ( INSTALLED_JAX, INSTALLED_PT, + INSTALLED_PT_EXPT, CommonTest, + parameterized, ) from .common import ( ModelTest, + compare_variables_recursive, ) if INSTALLED_PT: from deepmd.pt.model.model import get_model as get_model_pt from deepmd.pt.model.model.property_model import PropertyModel as PropertyModelPT + from deepmd.pt.utils.utils import to_numpy_array as torch_to_numpy + from deepmd.pt.utils.utils import to_torch_tensor as numpy_to_torch else: PropertyModelPT = None if INSTALLED_JAX: @@ -31,6 +46,11 @@ from deepmd.jax.model.property_model import PropertyModel as PropertyModelJAX else: PropertyModelJAX = None +if INSTALLED_PT_EXPT: + from deepmd.pt_expt.common import to_torch_array as pt_expt_numpy_to_torch + from deepmd.pt_expt.model import PropertyModel as PropertyModelPTExpt +else: + PropertyModelPTExpt = None from deepmd.utils.argcheck import ( model_args, ) @@ -67,6 +87,7 @@ def data(self) -> dict: tf_class = None dp_class = PropertyModelDP pt_class = PropertyModelPT + pt_expt_class = PropertyModelPTExpt jax_class = PropertyModelJAX args = model_args() @@ -79,6 +100,8 @@ def get_reference_backend(self): return self.RefBackend.PT if not self.skip_tf: return self.RefBackend.TF + if not self.skip_pt_expt and self.pt_expt_class is not None: + return self.RefBackend.PT_EXPT if not self.skip_dp: return self.RefBackend.DP raise ValueError("No available reference") @@ -100,6 +123,9 @@ def pass_data_to_cls(self, cls, data) -> Any: model = get_model_pt(data) model.atomic_model.out_bias.uniform_() return model + elif cls is PropertyModelPTExpt: + dp_model = get_model_dp(data) + return PropertyModelPTExpt.deserialize(dp_model.serialize()) elif cls is PropertyModelJAX: return get_model_jax(data) return cls(**data, **self.additional_data) @@ -172,6 +198,15 @@ def eval_pt(self, pt_obj: Any) -> Any: self.box, ) + def eval_pt_expt(self, pt_expt_obj: Any) -> Any: + return self.eval_pt_expt_model( + pt_expt_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + def eval_jax(self, jax_obj: Any) -> Any: return self.eval_jax_model( jax_obj, @@ -184,9 +219,1324 @@ def eval_jax(self, jax_obj: Any) -> Any: def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: # shape not matched. ravel... property_name = self.data["fitting_net"]["property_name"] - if backend in {self.RefBackend.DP, self.RefBackend.PT, self.RefBackend.JAX}: + if backend in { + self.RefBackend.DP, + self.RefBackend.PT, + self.RefBackend.PT_EXPT, + self.RefBackend.JAX, + }: return ( ret[property_name].ravel(), ret[f"atom_{property_name}"].ravel(), ) raise ValueError(f"Unknown backend: {backend}") + + +@unittest.skipUnless(INSTALLED_PT and INSTALLED_PT_EXPT, "PyTorch is not installed") +class TestPropertyModelAPIs(unittest.TestCase): + """Test consistency of model-level APIs between pt and dpmodel backends. + + Both models are constructed from the same serialized weights + (dpmodel -> serialize -> pt deserialize) so that numerical outputs + can be compared directly. + """ + + def setUp(self) -> None: + data = model_args().normalize_value( + { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20], + "rcut_smth": 1.8, + "rcut": 6.0, + "neuron": [2, 4, 8], + "resnet_dt": False, + "axis_neuron": 8, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "type": "property", + "neuron": [4, 4, 4], + "property_name": "foo", + "resnet_dt": True, + "precision": "float64", + "seed": 1, + "numb_fparam": 2, + "numb_aparam": 3, + "default_fparam": [0.5, -0.3], + }, + }, + trim_pattern="_*", + ) + # Build dpmodel first, then deserialize into pt/pt_expt to share weights + self.dp_model = get_model_dp(data) + serialized = self.dp_model.serialize() + self.pt_model = PropertyModelPT.deserialize(serialized) + self.pt_expt_model = PropertyModelPTExpt.deserialize(serialized) + + # Coords / atype / box + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, -1, 3) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32).reshape(1, -1) + self.box = np.array( + [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, 9) + + # Build extended coords + nlist for lower-level calls + rcut = 6.0 + nframes, nloc = self.atype.shape[:2] + coord_normalized = normalize_coord( + self.coords.reshape(nframes, nloc, 3), + self.box.reshape(nframes, 3, 3), + ) + extended_coord, extended_atype, mapping = extend_coord_with_ghosts( + coord_normalized, self.atype, self.box, rcut + ) + nlist = build_neighbor_list( + extended_coord, + extended_atype, + nloc, + rcut, + [20, 20], + distinguish_types=True, + ) + self.extended_coord = extended_coord.reshape(nframes, -1, 3) + self.extended_atype = extended_atype + self.mapping = mapping + self.nlist = nlist + + # aparam for forward evaluation (1 frame, 6 atoms, 3 aparam) + rng = np.random.default_rng(42) + self.eval_aparam = rng.normal(size=(1, nloc, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + def test_translated_output_def(self) -> None: + """translated_output_def should return the same keys on dp, pt, and pt_expt.""" + dp_def = self.dp_model.translated_output_def() + pt_def = self.pt_model.translated_output_def() + pt_expt_def = self.pt_expt_model.translated_output_def() + self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + self.assertEqual(set(dp_def.keys()), set(pt_expt_def.keys())) + for key in dp_def: + self.assertEqual(dp_def[key].shape, pt_def[key].shape) + self.assertEqual(dp_def[key].shape, pt_expt_def[key].shape) + + def test_get_descriptor(self) -> None: + """get_descriptor should return a non-None object on both backends.""" + self.assertIsNotNone(self.dp_model.get_descriptor()) + self.assertIsNotNone(self.pt_model.get_descriptor()) + + def test_get_fitting_net(self) -> None: + """get_fitting_net should return a non-None object on both backends.""" + self.assertIsNotNone(self.dp_model.get_fitting_net()) + self.assertIsNotNone(self.pt_model.get_fitting_net()) + + def test_get_out_bias(self) -> None: + """get_out_bias should return numerically equal values on dp and pt. + + PropertyModel's apply_out_stat applies output * std + bias, + so the bias storage should be consistent across backends. + """ + dp_bias = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias = torch_to_numpy(self.pt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias, pt_bias, rtol=1e-10, atol=1e-10) + # Verify shape: (n_output_keys x ntypes x task_dim) for property + self.assertEqual(dp_bias.shape[1], 2) # ntypes + self.assertGreater(dp_bias.shape[0], 0) # at least one output key + + def test_set_out_bias(self) -> None: + """set_out_bias should update the bias on both backends.""" + dp_bias = to_numpy_array(self.dp_model.get_out_bias()) + new_bias = dp_bias + 1.0 + # dp + self.dp_model.set_out_bias(new_bias) + np.testing.assert_allclose( + to_numpy_array(self.dp_model.get_out_bias()), + new_bias, + rtol=1e-10, + atol=1e-10, + ) + # pt + self.pt_model.set_out_bias(numpy_to_torch(new_bias)) + np.testing.assert_allclose( + torch_to_numpy(self.pt_model.get_out_bias()), + new_bias, + rtol=1e-10, + atol=1e-10, + ) + + def test_model_output_def(self) -> None: + """model_output_def should return the same keys and shapes on dp and pt.""" + dp_def = self.dp_model.model_output_def().get_data() + pt_def = self.pt_model.model_output_def().get_data() + self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + for key in dp_def: + self.assertEqual(dp_def[key].shape, pt_def[key].shape) + + def test_model_output_type(self) -> None: + """model_output_type should return the same list on dp and pt.""" + self.assertEqual( + self.dp_model.model_output_type(), + self.pt_model.model_output_type(), + ) + + def test_do_grad_r(self) -> None: + """do_grad_r should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.do_grad_r("foo"), + self.pt_model.do_grad_r("foo"), + ) + self.assertFalse(self.dp_model.do_grad_r("foo")) + + def test_do_grad_c(self) -> None: + """do_grad_c should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.do_grad_c("foo"), + self.pt_model.do_grad_c("foo"), + ) + self.assertFalse(self.dp_model.do_grad_c("foo")) + + def test_get_rcut(self) -> None: + """get_rcut should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_rcut(), self.pt_model.get_rcut()) + self.assertAlmostEqual(self.dp_model.get_rcut(), 6.0) + + def test_get_type_map(self) -> None: + """get_type_map should return the same list on dp and pt.""" + self.assertEqual(self.dp_model.get_type_map(), self.pt_model.get_type_map()) + self.assertEqual(self.dp_model.get_type_map(), ["O", "H"]) + + def test_get_sel(self) -> None: + """get_sel should return the same list on dp and pt.""" + self.assertEqual(self.dp_model.get_sel(), self.pt_model.get_sel()) + self.assertEqual(self.dp_model.get_sel(), [20, 20]) + + def test_get_nsel(self) -> None: + """get_nsel should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_nsel(), self.pt_model.get_nsel()) + self.assertEqual(self.dp_model.get_nsel(), 40) + + def test_get_nnei(self) -> None: + """get_nnei should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_nnei(), self.pt_model.get_nnei()) + self.assertEqual(self.dp_model.get_nnei(), 40) + + def test_mixed_types(self) -> None: + """mixed_types should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.mixed_types(), self.pt_model.mixed_types()) + # se_e2_a is not mixed-types + self.assertFalse(self.dp_model.mixed_types()) + + def test_has_message_passing(self) -> None: + """has_message_passing should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.has_message_passing(), + self.pt_model.has_message_passing(), + ) + self.assertFalse(self.dp_model.has_message_passing()) + + def test_need_sorted_nlist_for_lower(self) -> None: + """need_sorted_nlist_for_lower should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.need_sorted_nlist_for_lower(), + self.pt_model.need_sorted_nlist_for_lower(), + ) + self.assertFalse(self.dp_model.need_sorted_nlist_for_lower()) + + def test_get_dim_fparam(self) -> None: + """get_dim_fparam should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_dim_fparam(), self.pt_model.get_dim_fparam()) + self.assertEqual(self.dp_model.get_dim_fparam(), 2) + + def test_get_dim_aparam(self) -> None: + """get_dim_aparam should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_dim_aparam(), self.pt_model.get_dim_aparam()) + self.assertEqual(self.dp_model.get_dim_aparam(), 3) + + def test_get_sel_type(self) -> None: + """get_sel_type should return the same list on dp and pt.""" + self.assertEqual(self.dp_model.get_sel_type(), self.pt_model.get_sel_type()) + self.assertEqual(self.dp_model.get_sel_type(), [0, 1]) + + def test_is_aparam_nall(self) -> None: + """is_aparam_nall should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.is_aparam_nall(), self.pt_model.is_aparam_nall()) + self.assertFalse(self.dp_model.is_aparam_nall()) + + def test_get_model_def_script(self) -> None: + """get_model_def_script should return the same value on dp, pt, and pt_expt.""" + dp_val = self.dp_model.get_model_def_script() + pt_val = self.pt_model.get_model_def_script() + pe_val = self.pt_expt_model.get_model_def_script() + self.assertEqual(dp_val, pt_val) + self.assertEqual(dp_val, pe_val) + + def test_get_min_nbor_dist(self) -> None: + """get_min_nbor_dist should return the same value on dp, pt, and pt_expt.""" + dp_val = self.dp_model.get_min_nbor_dist() + pt_val = self.pt_model.get_min_nbor_dist() + pe_val = self.pt_expt_model.get_min_nbor_dist() + self.assertEqual(dp_val, pt_val) + self.assertEqual(dp_val, pe_val) + + def test_set_case_embd(self) -> None: + """set_case_embd should produce consistent results across backends. + + Also verifies that different case indices produce different outputs, + confirming the embedding is actually used. + """ + from deepmd.utils.argcheck import ( + model_args, + ) + + # Build a model with dim_case_embd > 0 + data = model_args().normalize_value( + { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "type": "property", + "property_name": "foo", + "neuron": [5, 5], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + "dim_case_embd": 3, + }, + }, + trim_pattern="_*", + ) + dp_model = get_model_dp(data) + serialized = dp_model.serialize() + pt_model = PropertyModelPT.deserialize(serialized) + pe_model = PropertyModelPTExpt.deserialize(serialized) + + def _eval(case_idx): + dp_model.set_case_embd(case_idx) + pt_model.set_case_embd(case_idx) + pe_model.set_case_embd(case_idx) + dp_ret = dp_model(self.coords, self.atype, box=self.box) + pt_ret = { + k: torch_to_numpy(v) + for k, v in pt_model( + numpy_to_torch(self.coords), + numpy_to_torch(self.atype), + box=numpy_to_torch(self.box), + ).items() + } + coord_t = pt_expt_numpy_to_torch(self.coords) + coord_t.requires_grad_(True) + pe_ret = { + k: v.detach().cpu().numpy() + for k, v in pe_model( + coord_t, + pt_expt_numpy_to_torch(self.atype), + box=pt_expt_numpy_to_torch(self.box), + ).items() + } + return dp_ret, pt_ret, pe_ret + + dp0, pt0, pe0 = _eval(0) + dp1, pt1, pe1 = _eval(1) + + # Cross-backend consistency for each case index + for key in ("foo", "atom_foo"): + np.testing.assert_allclose( + dp0[key], + pt0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 0: dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp0[key], + pe0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 0: dp vs pt_expt mismatch in {key}", + ) + np.testing.assert_allclose( + dp1[key], + pt1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 1: dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp1[key], + pe1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 1: dp vs pt_expt mismatch in {key}", + ) + # Different case indices should produce different outputs + self.assertFalse( + np.allclose(dp0["foo"], dp1["foo"]), + "set_case_embd(0) and set_case_embd(1) produced the same foo", + ) + + def test_atomic_output_def(self) -> None: + """atomic_output_def should return the same keys and shapes on dp and pt.""" + dp_def = self.dp_model.atomic_output_def() + pt_def = self.pt_model.atomic_output_def() + self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + for key in dp_def.keys(): + self.assertEqual(dp_def[key].shape, pt_def[key].shape) + + def test_format_nlist(self) -> None: + """format_nlist should produce the same result on dp and pt.""" + dp_nlist = self.dp_model.format_nlist( + self.extended_coord, + self.extended_atype, + self.nlist, + ) + pt_nlist = torch_to_numpy( + self.pt_model.format_nlist( + numpy_to_torch(self.extended_coord), + numpy_to_torch(self.extended_atype), + numpy_to_torch(self.nlist), + ) + ) + np.testing.assert_equal(dp_nlist, pt_nlist) + + def test_forward_common_atomic(self) -> None: + """forward_common_atomic should produce consistent results on dp and pt. + + Compares at the atomic_model level, where both backends define this method. + """ + dp_ret = self.dp_model.atomic_model.forward_common_atomic( + self.extended_coord, + self.extended_atype, + self.nlist, + mapping=self.mapping, + aparam=self.eval_aparam, + ) + pt_ret = self.pt_model.atomic_model.forward_common_atomic( + numpy_to_torch(self.extended_coord), + numpy_to_torch(self.extended_atype), + numpy_to_torch(self.nlist), + mapping=numpy_to_torch(self.mapping), + aparam=numpy_to_torch(self.eval_aparam), + ) + # Compare the common keys + common_keys = set(dp_ret.keys()) & set(pt_ret.keys()) + self.assertTrue(len(common_keys) > 0) + for key in common_keys: + if dp_ret[key] is not None and pt_ret[key] is not None: + np.testing.assert_allclose( + dp_ret[key], + torch_to_numpy(pt_ret[key]), + rtol=1e-10, + atol=1e-10, + err_msg=f"Mismatch in forward_common_atomic key '{key}'", + ) + + def test_has_default_fparam(self) -> None: + """has_default_fparam should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.has_default_fparam(), + self.pt_model.has_default_fparam(), + ) + self.assertTrue(self.dp_model.has_default_fparam()) + + def test_get_default_fparam(self) -> None: + """get_default_fparam should return consistent values on dp and pt.""" + dp_val = self.dp_model.get_default_fparam() + pt_val = self.pt_model.get_default_fparam() + np.testing.assert_allclose(dp_val, pt_val, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dp_val, [0.5, -0.3], rtol=1e-10, atol=1e-10) + + def test_change_out_bias(self) -> None: + """change_out_bias should produce consistent bias on dp, pt, and pt_expt. + + Tests both set-by-statistic and change-by-statistic modes. + Note: change_out_bias only updates the output bias, not fitting input + stats (fparam/aparam). Fitting stats are updated by compute_or_load_stat. + """ + nframes = 2 + nloc = 6 + numb_fparam = 2 + numb_aparam = 3 + rng = np.random.default_rng(123) + + # Use realistic coords (from setUp, tiled for 2 frames) + coords_2f = np.tile(self.coords, (nframes, 1, 1)) # (2, 6, 3) + atype_2f = np.array([[0, 0, 1, 1, 1, 1], [0, 1, 1, 0, 1, 1]], dtype=np.int32) + box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) + natoms_data = np.array([[6, 6, 2, 4], [6, 6, 2, 4]], dtype=np.int32) + foo_data = rng.normal(size=(nframes, 1)).astype(GLOBAL_NP_FLOAT_PRECISION) + fparam_data = rng.normal(size=(nframes, numb_fparam)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + aparam_data = rng.normal(size=(nframes, nloc, numb_aparam)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + # dpmodel stat data (numpy) + dp_merged = [ + { + "coord": coords_2f, + "atype": atype_2f, + "atype_ext": atype_2f, + "box": box_2f, + "natoms": natoms_data, + "foo": foo_data, + "find_foo": np.float32(1.0), + "fparam": fparam_data, + "aparam": aparam_data, + } + ] + # pt stat data (torch tensors) + pt_merged = [ + { + "coord": numpy_to_torch(coords_2f), + "atype": numpy_to_torch(atype_2f), + "atype_ext": numpy_to_torch(atype_2f), + "box": numpy_to_torch(box_2f), + "natoms": numpy_to_torch(natoms_data), + "foo": numpy_to_torch(foo_data), + "find_foo": np.float32(1.0), + "fparam": numpy_to_torch(fparam_data), + "aparam": numpy_to_torch(aparam_data), + } + ] + # pt_expt stat data (numpy, same as dp) + pe_merged = dp_merged + + # Save initial (zero) bias + dp_bias_init = to_numpy_array(self.dp_model.get_out_bias()).copy() + + # --- Test "set-by-statistic" mode --- + self.dp_model.change_out_bias(dp_merged, bias_adjust_mode="set-by-statistic") + self.pt_model.change_out_bias(pt_merged, bias_adjust_mode="set-by-statistic") + self.pt_expt_model.change_out_bias( + pe_merged, bias_adjust_mode="set-by-statistic" + ) + + # Verify out bias consistency + dp_bias = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias = torch_to_numpy(self.pt_model.get_out_bias()) + pe_bias = to_numpy_array(self.pt_expt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias, pt_bias, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dp_bias, pe_bias, rtol=1e-10, atol=1e-10) + self.assertFalse( + np.allclose(dp_bias, dp_bias_init), + "set-by-statistic did not change the bias from initial values", + ) + + # --- Test "change-by-statistic" mode --- + dp_bias_before = dp_bias.copy() + self.dp_model.change_out_bias(dp_merged, bias_adjust_mode="change-by-statistic") + self.pt_model.change_out_bias(pt_merged, bias_adjust_mode="change-by-statistic") + self.pt_expt_model.change_out_bias( + pe_merged, bias_adjust_mode="change-by-statistic" + ) + + # Verify out bias consistency + dp_bias2 = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias2 = torch_to_numpy(self.pt_model.get_out_bias()) + pe_bias2 = to_numpy_array(self.pt_expt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias2, pt_bias2, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dp_bias2, pe_bias2, rtol=1e-10, atol=1e-10) + self.assertFalse( + np.allclose(dp_bias2, dp_bias_before), + "change-by-statistic did not further change the bias", + ) + + def test_change_type_map(self) -> None: + """change_type_map should produce consistent results on dp and pt. + + Uses a DPA1 (se_atten) descriptor since se_e2_a does not support + change_type_map (non-mixed-types descriptors raise NotImplementedError). + """ + from deepmd.utils.argcheck import model_args as model_args_fn + + data = model_args_fn().normalize_value( + { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_atten", + "sel": 20, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "seed": 1, + "attn": 6, + "attn_layer": 0, + }, + "fitting_net": { + "type": "property", + "neuron": [4, 4, 4], + "property_name": "foo", + "resnet_dt": True, + "precision": "float64", + "seed": 1, + }, + }, + trim_pattern="_*", + ) + dp_model = get_model_dp(data) + pt_model = PropertyModelPT.deserialize(dp_model.serialize()) + + # Set non-zero out_bias so the swap is non-trivial + dp_bias_orig = to_numpy_array(dp_model.get_out_bias()).copy() + new_bias = dp_bias_orig.copy() + new_bias[:, 0, :] = 1.5 # type 0 ("O") + new_bias[:, 1, :] = -3.7 # type 1 ("H") + dp_model.set_out_bias(new_bias) + pt_model.set_out_bias(numpy_to_torch(new_bias)) + + new_type_map = ["H", "O"] + dp_model.change_type_map(new_type_map) + pt_model.change_type_map(new_type_map) + + # Both should have the new type_map + self.assertEqual(dp_model.get_type_map(), new_type_map) + self.assertEqual(pt_model.get_type_map(), new_type_map) + + # Out_bias should be reordered consistently between backends + dp_bias_new = to_numpy_array(dp_model.get_out_bias()) + pt_bias_new = torch_to_numpy(pt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias_new, pt_bias_new, rtol=1e-10, atol=1e-10) + + # Verify the reorder is correct: old type 0 -> new type 1, old type 1 -> new type 0 + np.testing.assert_allclose( + dp_bias_new[:, 0, :], + new_bias[:, 1, :], + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + dp_bias_new[:, 1, :], + new_bias[:, 0, :], + rtol=1e-10, + atol=1e-10, + ) + + def test_change_type_map_extend_stat(self) -> None: + """change_type_map with model_with_new_type_stat should propagate stats consistently across dp, pt, and pt_expt. + + Verifies that the model-level change_type_map correctly unwraps + model_with_new_type_stat.atomic_model before forwarding to the + atomic model. + """ + from deepmd.utils.argcheck import model_args as model_args_fn + + small_tm = ["O", "H"] + large_tm = ["O", "H", "Li"] + + small_data = model_args_fn().normalize_value( + { + "type_map": small_tm, + "descriptor": { + "type": "se_atten", + "sel": 20, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "seed": 1, + "attn": 6, + "attn_layer": 0, + }, + "fitting_net": { + "type": "property", + "neuron": [4, 4, 4], + "property_name": "foo", + "resnet_dt": True, + "precision": "float64", + "seed": 1, + }, + }, + trim_pattern="_*", + ) + large_data = model_args_fn().normalize_value( + { + "type_map": large_tm, + "descriptor": { + "type": "se_atten", + "sel": 20, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "seed": 2, + "attn": 6, + "attn_layer": 0, + }, + "fitting_net": { + "type": "property", + "neuron": [4, 4, 4], + "property_name": "foo", + "resnet_dt": True, + "precision": "float64", + "seed": 2, + }, + }, + trim_pattern="_*", + ) + + dp_small = get_model_dp(small_data) + dp_large = get_model_dp(large_data) + + # Set distinguishable random stats on the large model's descriptor + rng = np.random.default_rng(42) + desc_large = dp_large.get_descriptor() + mean_large, std_large = desc_large.get_stat_mean_and_stddev() + mean_rand = rng.random(size=to_numpy_array(mean_large).shape) + std_rand = rng.random(size=to_numpy_array(std_large).shape) + desc_large.set_stat_mean_and_stddev(mean_rand, std_rand) + + # Build pt and pt_expt models from dp serialization + pt_small = PropertyModelPT.deserialize(dp_small.serialize()) + pt_large = PropertyModelPT.deserialize(dp_large.serialize()) + pt_expt_small = PropertyModelPTExpt.deserialize(dp_small.serialize()) + pt_expt_large = PropertyModelPTExpt.deserialize(dp_large.serialize()) + + # Extend type map with model_with_new_type_stat at the model level + dp_small.change_type_map(large_tm, model_with_new_type_stat=dp_large) + pt_small.change_type_map(large_tm, model_with_new_type_stat=pt_large) + pt_expt_small.change_type_map(large_tm, model_with_new_type_stat=pt_expt_large) + + # Descriptor stats should be consistent across backends + dp_mean, dp_std = dp_small.get_descriptor().get_stat_mean_and_stddev() + pt_mean, pt_std = pt_small.get_descriptor().get_stat_mean_and_stddev() + pt_expt_mean, pt_expt_std = ( + pt_expt_small.get_descriptor().get_stat_mean_and_stddev() + ) + np.testing.assert_allclose( + to_numpy_array(dp_mean), + torch_to_numpy(pt_mean), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + to_numpy_array(dp_std), + torch_to_numpy(pt_std), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + to_numpy_array(dp_mean), + to_numpy_array(pt_expt_mean), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + to_numpy_array(dp_std), + to_numpy_array(pt_expt_std), + rtol=1e-10, + atol=1e-10, + ) + + def test_update_sel(self) -> None: + """update_sel should return the same result on dp and pt.""" + from unittest.mock import ( + patch, + ) + + from deepmd.dpmodel.model.dp_model import DPModelCommon as DPModelCommonDP + from deepmd.pt.model.model.dp_model import DPModelCommon as DPModelCommonPT + + mock_min_nbor_dist = 0.5 + mock_sel = [10, 20] + local_jdata = { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": "auto", + "rcut_smth": 0.50, + "rcut": 6.00, + }, + "fitting_net": { + "type": "property", + "neuron": [4, 4, 4], + "property_name": "foo", + }, + } + type_map = ["O", "H"] + + with patch( + "deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat", + return_value=(mock_min_nbor_dist, mock_sel), + ): + dp_result, dp_min_dist = DPModelCommonDP.update_sel( + None, type_map, local_jdata + ) + + with patch( + "deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat", + return_value=(mock_min_nbor_dist, mock_sel), + ): + pt_result, pt_min_dist = DPModelCommonPT.update_sel( + None, type_map, local_jdata + ) + + self.assertEqual(dp_result, pt_result) + self.assertEqual(dp_min_dist, pt_min_dist) + # Verify sel was actually updated (not still "auto") + self.assertIsInstance(dp_result["descriptor"]["sel"], list) + self.assertNotEqual(dp_result["descriptor"]["sel"], "auto") + + def test_get_ntypes(self) -> None: + """get_ntypes should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_ntypes(), self.pt_model.get_ntypes()) + self.assertEqual(self.dp_model.get_ntypes(), 2) + + def test_compute_or_load_out_stat(self) -> None: + """compute_or_load_out_stat should produce consistent bias on dp and pt. + + Tests both the compute path (from data) and the load path (from file). + PropertyModel's apply_out_stat applies output * std + bias, + so the stored bias should be consistent across backends. + """ + import tempfile + from pathlib import ( + Path, + ) + + import h5py + + from deepmd.utils.path import ( + DPPath, + ) + + nframes = 2 + coords_2f = np.tile(self.coords, (nframes, 1, 1)) + atype_2f = np.array([[0, 0, 1, 1, 1, 1], [0, 1, 1, 0, 1, 1]], dtype=np.int32) + box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) + natoms_data = np.array([[6, 6, 2, 4], [6, 6, 2, 4]], dtype=np.int32) + foo_data = ( + np.random.default_rng(42) + .normal(size=(nframes, 1)) + .astype(GLOBAL_NP_FLOAT_PRECISION) + ) + + dp_merged = [ + { + "coord": coords_2f, + "atype": atype_2f, + "atype_ext": atype_2f, + "box": box_2f, + "natoms": natoms_data, + "foo": foo_data, + "find_foo": np.float32(1.0), + } + ] + pt_merged = [ + { + "coord": numpy_to_torch(coords_2f), + "atype": numpy_to_torch(atype_2f), + "atype_ext": numpy_to_torch(atype_2f), + "box": numpy_to_torch(box_2f), + "natoms": numpy_to_torch(natoms_data), + "foo": numpy_to_torch(foo_data), + "find_foo": np.float32(1.0), + } + ] + + # Verify bias is initially identical + dp_bias_before = to_numpy_array(self.dp_model.get_out_bias()).copy() + pt_bias_before = torch_to_numpy(self.pt_model.get_out_bias()).copy() + np.testing.assert_allclose( + dp_bias_before, pt_bias_before, rtol=1e-10, atol=1e-10 + ) + + with tempfile.TemporaryDirectory() as tmpdir: + # Create separate h5 files for dp and pt + dp_h5 = str((Path(tmpdir) / "dp_stat.h5").resolve()) + pt_h5 = str((Path(tmpdir) / "pt_stat.h5").resolve()) + with h5py.File(dp_h5, "w"): + pass + with h5py.File(pt_h5, "w"): + pass + dp_stat_path = DPPath(dp_h5, "a") + pt_stat_path = DPPath(pt_h5, "a") + + # 1. Compute stats and save to file + self.dp_model.atomic_model.compute_or_load_out_stat( + dp_merged, stat_file_path=dp_stat_path + ) + self.pt_model.atomic_model.compute_or_load_out_stat( + pt_merged, stat_file_path=pt_stat_path + ) + + dp_bias_after = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias_after = torch_to_numpy(self.pt_model.get_out_bias()) + np.testing.assert_allclose( + dp_bias_after, pt_bias_after, rtol=1e-10, atol=1e-10 + ) + + # 2. Verify both backends saved the same file content + with h5py.File(dp_h5, "r") as dp_f, h5py.File(pt_h5, "r") as pt_f: + dp_keys = sorted(dp_f.keys()) + pt_keys = sorted(pt_f.keys()) + self.assertEqual(dp_keys, pt_keys) + for key in dp_keys: + np.testing.assert_allclose( + np.array(dp_f[key]), + np.array(pt_f[key]), + rtol=1e-10, + atol=1e-10, + err_msg=f"Stat file content mismatch for key {key}", + ) + + # 3. Reset biases to zero, then load from file + zero_bias = np.zeros_like(dp_bias_after) + self.dp_model.set_out_bias(zero_bias) + self.pt_model.set_out_bias(numpy_to_torch(zero_bias)) + + # Use a callable that raises to ensure it loads from file, not recomputes + def raise_error(): + raise RuntimeError("Should not recompute — should load from file") + + self.dp_model.atomic_model.compute_or_load_out_stat( + raise_error, stat_file_path=dp_stat_path + ) + self.pt_model.atomic_model.compute_or_load_out_stat( + raise_error, stat_file_path=pt_stat_path + ) + + dp_bias_loaded = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias_loaded = torch_to_numpy(self.pt_model.get_out_bias()) + + # Loaded biases should match between backends + np.testing.assert_allclose( + dp_bias_loaded, pt_bias_loaded, rtol=1e-10, atol=1e-10 + ) + # Loaded biases should match the originally computed biases + np.testing.assert_allclose( + dp_bias_loaded, dp_bias_after, rtol=1e-10, atol=1e-10 + ) + + def test_get_observed_type_list(self) -> None: + """get_observed_type_list should be consistent across dp, pt, pt_expt. + + Uses mock data containing only type 0 ("O") so that type 1 ("H") is + unobserved and should be absent from the returned list. + """ + nframes = 2 + natoms = 6 + # All atoms are type 0 — type 1 is unobserved + atype_2f = np.zeros((nframes, natoms), dtype=np.int32) + coords_2f = np.tile(self.coords, (nframes, 1, 1)) + box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) + natoms_data = np.array([[natoms, natoms, natoms, 0]] * nframes, dtype=np.int32) + foo_data = ( + np.random.default_rng(42) + .normal(size=(nframes, 1)) + .astype(GLOBAL_NP_FLOAT_PRECISION) + ) + + dp_merged = [ + { + "coord": coords_2f, + "atype": atype_2f, + "atype_ext": atype_2f, + "box": box_2f, + "natoms": natoms_data, + "foo": foo_data, + "find_foo": np.float32(1.0), + } + ] + pt_merged = [ + { + "coord": numpy_to_torch(coords_2f), + "atype": numpy_to_torch(atype_2f), + "atype_ext": numpy_to_torch(atype_2f), + "box": numpy_to_torch(box_2f), + "natoms": numpy_to_torch(natoms_data), + "foo": numpy_to_torch(foo_data), + "find_foo": np.float32(1.0), + } + ] + + self.dp_model.atomic_model.compute_or_load_out_stat(dp_merged) + self.pt_model.atomic_model.compute_or_load_out_stat(pt_merged) + self.pt_expt_model.atomic_model.compute_or_load_out_stat(dp_merged) + + dp_observed = self.dp_model.get_observed_type_list() + pt_observed = self.pt_model.get_observed_type_list() + pe_observed = self.pt_expt_model.get_observed_type_list() + + self.assertEqual(dp_observed, pt_observed) + self.assertEqual(dp_observed, pe_observed) + # Property model uses stats_distinguish_types=False, so all types + # get the same (non-zero) bias — both types appear observed. + self.assertEqual(dp_observed, ["O", "H"]) + + +@parameterized( + (([], []), ([[0, 1]], [1])), # (pair_exclude_types, atom_exclude_types) + (False, True), # fparam_in_data +) +@unittest.skipUnless(INSTALLED_PT and INSTALLED_PT_EXPT, "PT and PT_EXPT are required") +class TestPropertyComputeOrLoadStat(unittest.TestCase): + """Test that compute_or_load_stat produces identical statistics on dp, pt, and pt_expt. + + Covers descriptor stats (dstd), fitting stats (fparam, aparam), and output bias. + Parameterized over exclusion types and whether fparam is explicitly provided or + injected via default_fparam. + """ + + def setUp(self) -> None: + (pair_exclude_types, atom_exclude_types), self.fparam_in_data = self.param + data = model_args().normalize_value( + { + "type_map": ["O", "H"], + "pair_exclude_types": pair_exclude_types, + "atom_exclude_types": atom_exclude_types, + "descriptor": { + "type": "dpa3", + "repflow": { + "n_dim": 20, + "e_dim": 10, + "a_dim": 8, + "nlayers": 3, + "e_rcut": 6.0, + "e_rcut_smth": 5.0, + "e_sel": 10, + "a_rcut": 4.0, + "a_rcut_smth": 3.5, + "a_sel": 8, + "axis_neuron": 4, + "update_angle": True, + "update_style": "res_residual", + "update_residual": 0.1, + "update_residual_init": "const", + }, + "precision": "float64", + "seed": 1, + }, + "fitting_net": { + "type": "property", + "neuron": [10, 10], + "property_name": "foo", + "precision": "float64", + "seed": 1, + "numb_fparam": 2, + "default_fparam": [0.5, -0.3], + "numb_aparam": 3, + }, + }, + trim_pattern="_*", + ) + + # Save data for reuse in load-from-file test + self._model_data = data + + # Build dp model, then deserialize into pt and pt_expt to share weights + self.dp_model = get_model_dp(data) + serialized = self.dp_model.serialize() + self.pt_model = PropertyModelPT.deserialize(serialized) + self.pt_expt_model = PropertyModelPTExpt.deserialize(serialized) + + # Test coords / atype / box for forward evaluation + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 0.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, -1, 3) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32).reshape(1, -1) + self.box = np.array( + [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, 9) + + # Mock training data for compute_or_load_stat + natoms = 6 + nframes = 3 + rng = np.random.default_rng(42) + coords_stat = rng.normal(size=(nframes, natoms, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + atype_stat = np.array([[0, 0, 1, 1, 1, 1]] * nframes, dtype=np.int32) + box_stat = np.tile( + np.eye(3, dtype=GLOBAL_NP_FLOAT_PRECISION).reshape(1, 3, 3) * 13.0, + (nframes, 1, 1), + ) + natoms_stat = np.array([[natoms, natoms, 2, 4]] * nframes, dtype=np.int32) + foo_stat = rng.normal(size=(nframes, 1)).astype(GLOBAL_NP_FLOAT_PRECISION) + aparam_stat = rng.normal(size=(nframes, natoms, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + # dp / pt_expt sample (numpy) + np_sample = { + "coord": coords_stat, + "atype": atype_stat, + "atype_ext": atype_stat, + "box": box_stat, + "natoms": natoms_stat, + "foo": foo_stat, + "find_foo": np.float32(1.0), + "aparam": aparam_stat, + } + # pt sample (torch tensors) + pt_sample = { + "coord": numpy_to_torch(coords_stat), + "atype": numpy_to_torch(atype_stat), + "atype_ext": numpy_to_torch(atype_stat), + "box": numpy_to_torch(box_stat), + "natoms": numpy_to_torch(natoms_stat), + "foo": numpy_to_torch(foo_stat), + "find_foo": np.float32(1.0), + "aparam": numpy_to_torch(aparam_stat), + } + + if self.fparam_in_data: + fparam_stat = rng.normal(size=(nframes, 2)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + np_sample["fparam"] = fparam_stat + pt_sample["fparam"] = numpy_to_torch(fparam_stat) + self.expected_fparam_avg = np.mean(fparam_stat, axis=0) + else: + # No fparam -> _make_wrapped_sampler injects default_fparam + self.expected_fparam_avg = np.array([0.5, -0.3]) + + self.np_sampled = [np_sample] + self.pt_sampled = [pt_sample] + + # aparam for forward evaluation (1 frame, 6 atoms, 3 aparam) + self.eval_aparam = rng.normal(size=(1, natoms, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + + def _eval_dp(self) -> dict: + return self.dp_model( + self.coords, self.atype, box=self.box, aparam=self.eval_aparam + ) + + def _eval_pt(self) -> dict: + return { + kk: torch_to_numpy(vv) + for kk, vv in self.pt_model( + numpy_to_torch(self.coords), + numpy_to_torch(self.atype), + box=numpy_to_torch(self.box), + aparam=numpy_to_torch(self.eval_aparam), + do_atomic_virial=True, + ).items() + } + + def _eval_pt_expt(self) -> dict: + coord_t = pt_expt_numpy_to_torch(self.coords) + coord_t.requires_grad_(True) + return { + k: v.detach().cpu().numpy() + for k, v in self.pt_expt_model( + coord_t, + pt_expt_numpy_to_torch(self.atype), + box=pt_expt_numpy_to_torch(self.box), + aparam=pt_expt_numpy_to_torch(self.eval_aparam), + do_atomic_virial=True, + ).items() + } + + def test_compute_stat(self) -> None: + # 1. Pre-stat forward consistency + dp_ret0 = self._eval_dp() + pt_ret0 = self._eval_pt() + pe_ret0 = self._eval_pt_expt() + for key in ("foo", "atom_foo"): + np.testing.assert_allclose( + dp_ret0[key], + pt_ret0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Pre-stat dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp_ret0[key], + pe_ret0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Pre-stat dp vs pt_expt mismatch in {key}", + ) + + # 2. Run compute_or_load_stat on all three backends + self.dp_model.compute_or_load_stat(lambda: self.np_sampled) + self.pt_model.compute_or_load_stat(lambda: self.pt_sampled) + self.pt_expt_model.compute_or_load_stat(lambda: self.np_sampled) + + # 3. Serialize all three and compare @variables + dp_ser = self.dp_model.serialize() + pt_ser = self.pt_model.serialize() + pe_ser = self.pt_expt_model.serialize() + compare_variables_recursive(dp_ser, pt_ser) + compare_variables_recursive(dp_ser, pe_ser) + + # 4. Post-stat forward consistency + # PropertyModel's apply_out_stat applies output * std + bias, so + # output values WILL differ from pre-stat after stat computation. + dp_ret1 = self._eval_dp() + pt_ret1 = self._eval_pt() + pe_ret1 = self._eval_pt_expt() + for key in ("foo", "atom_foo"): + np.testing.assert_allclose( + dp_ret1[key], + pt_ret1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Post-stat dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp_ret1[key], + pe_ret1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Post-stat dp vs pt_expt mismatch in {key}", + ) + + # 5. Non-triviality checks + fit_vars = dp_ser["fitting"]["@variables"] + # fparam stats were computed + fparam_avg = np.asarray(fit_vars["fparam_avg"]) + self.assertFalse( + np.allclose(fparam_avg, 0.0), + "fparam_avg is still zero — fparam stats were not computed", + ) + np.testing.assert_allclose( + fparam_avg, + self.expected_fparam_avg, + rtol=1e-10, + atol=1e-10, + err_msg="fparam_avg does not match expected values", + ) + # aparam stats were computed + aparam_avg = np.asarray(fit_vars["aparam_avg"]) + self.assertFalse( + np.allclose(aparam_avg, 0.0), + "aparam_avg is still zero — aparam stats were not computed", + ) + + def test_load_stat_from_file(self) -> None: + import tempfile + from pathlib import ( + Path, + ) + + import h5py + + from deepmd.utils.path import ( + DPPath, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + # Create separate stat files for each backend + dp_h5 = str((Path(tmpdir) / "dp_stat.h5").resolve()) + pt_h5 = str((Path(tmpdir) / "pt_stat.h5").resolve()) + pe_h5 = str((Path(tmpdir) / "pe_stat.h5").resolve()) + for p in (dp_h5, pt_h5, pe_h5): + with h5py.File(p, "w"): + pass + + # 1. Compute stats and save to file + self.dp_model.compute_or_load_stat( + lambda: self.np_sampled, stat_file_path=DPPath(dp_h5, "a") + ) + self.pt_model.compute_or_load_stat( + lambda: self.pt_sampled, stat_file_path=DPPath(pt_h5, "a") + ) + self.pt_expt_model.compute_or_load_stat( + lambda: self.np_sampled, stat_file_path=DPPath(pe_h5, "a") + ) + + # Save the computed serializations as reference + dp_ser_computed = self.dp_model.serialize() + pt_ser_computed = self.pt_model.serialize() + pe_ser_computed = self.pt_expt_model.serialize() + + # 2. Build fresh models from the same initial weights + dp_model2 = get_model_dp(self._model_data) + pt_model2 = PropertyModelPT.deserialize(dp_model2.serialize()) + pe_model2 = PropertyModelPTExpt.deserialize(dp_model2.serialize()) + + # 3. Load stats from file (should NOT call the sampled func) + def raise_error(): + raise RuntimeError("Should load from file, not recompute") + + dp_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(dp_h5, "a") + ) + pt_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(pt_h5, "a") + ) + pe_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(pe_h5, "a") + ) + + # 4. Loaded models should match the computed ones + dp_ser_loaded = dp_model2.serialize() + pt_ser_loaded = pt_model2.serialize() + pe_ser_loaded = pe_model2.serialize() + compare_variables_recursive(dp_ser_computed, dp_ser_loaded) + compare_variables_recursive(pt_ser_computed, pt_ser_loaded) + compare_variables_recursive(pe_ser_computed, pe_ser_loaded) + + # 5. Cross-backend consistency after loading + compare_variables_recursive(dp_ser_loaded, pt_ser_loaded) + compare_variables_recursive(dp_ser_loaded, pe_ser_loaded) diff --git a/source/tests/consistent/model/test_zbl_ener.py b/source/tests/consistent/model/test_zbl_ener.py index 2783bb4a02..0cf907bee2 100644 --- a/source/tests/consistent/model/test_zbl_ener.py +++ b/source/tests/consistent/model/test_zbl_ener.py @@ -1,4 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import copy +import os import unittest from typing import ( Any, @@ -6,8 +8,18 @@ import numpy as np +from deepmd.dpmodel.common import ( + to_numpy_array, +) from deepmd.dpmodel.model.dp_zbl_model import DPZBLModel as DPZBLModelDP from deepmd.dpmodel.model.model import get_model as get_model_dp +from deepmd.dpmodel.utils.nlist import ( + build_neighbor_list, + extend_coord_with_ghosts, +) +from deepmd.dpmodel.utils.region import ( + normalize_coord, +) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) @@ -15,17 +27,21 @@ from ..common import ( INSTALLED_JAX, INSTALLED_PT, + INSTALLED_PT_EXPT, SKIP_FLAG, CommonTest, parameterized, ) from .common import ( ModelTest, + compare_variables_recursive, ) if INSTALLED_PT: from deepmd.pt.model.model import get_model as get_model_pt from deepmd.pt.model.model.dp_zbl_model import DPZBLModel as DPZBLModelPT + from deepmd.pt.utils.utils import to_numpy_array as torch_to_numpy + from deepmd.pt.utils.utils import to_torch_tensor as numpy_to_torch else: DPZBLModelPT = None if INSTALLED_JAX: @@ -33,8 +49,11 @@ from deepmd.jax.model.model import get_model as get_model_jax else: DPZBLModelJAX = None -import os - +if INSTALLED_PT_EXPT: + from deepmd.pt_expt.common import to_torch_array as pt_expt_numpy_to_torch + from deepmd.pt_expt.model import DPZBLModel as DPZBLModelPTExpt +else: + DPZBLModelPTExpt = None from deepmd.utils.argcheck import ( model_args, ) @@ -92,6 +111,7 @@ def data(self) -> dict: dp_class = DPZBLModelDP pt_class = DPZBLModelPT + pt_expt_class = DPZBLModelPTExpt jax_class = DPZBLModelJAX args = model_args() @@ -104,6 +124,8 @@ def get_reference_backend(self): return self.RefBackend.PT if not self.skip_tf: return self.RefBackend.TF + if not self.skip_pt_expt and self.pt_expt_class is not None: + return self.RefBackend.PT_EXPT if not self.skip_jax: return self.RefBackend.JAX if not self.skip_dp: @@ -125,6 +147,9 @@ def pass_data_to_cls(self, cls, data) -> Any: return get_model_dp(data) elif cls is DPZBLModelPT: return get_model_pt(data) + elif cls is DPZBLModelPTExpt: + dp_model = get_model_dp(data) + return DPZBLModelPTExpt.deserialize(dp_model.serialize()) elif cls is DPZBLModelJAX: return get_model_jax(data) return cls(**data, **self.additional_data) @@ -196,6 +221,15 @@ def eval_pt(self, pt_obj: Any) -> Any: self.box, ) + def eval_pt_expt(self, pt_expt_obj: Any) -> Any: + return self.eval_pt_expt_model( + pt_expt_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + def eval_jax(self, jax_obj: Any) -> Any: return self.eval_jax_model( jax_obj, @@ -218,6 +252,7 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: ) elif backend in { self.RefBackend.PT, + self.RefBackend.PT_EXPT, self.RefBackend.JAX, }: return ( @@ -227,3 +262,1044 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: ret["virial"].ravel(), ) raise ValueError(f"Unknown backend: {backend}") + + +@unittest.skipUnless(INSTALLED_PT and INSTALLED_PT_EXPT, "PyTorch is not installed") +class TestZBLEnerModelAPIs(unittest.TestCase): + """Test consistency of model-level APIs between pt and dpmodel backends. + + Both models are constructed from the same serialized weights + (dpmodel -> serialize -> pt deserialize) so that numerical outputs + can be compared directly. + + DPZBLModel is a linear combination model (DP + ZBL) and does NOT + support get_descriptor() or get_fitting_net() at the top level. + """ + + def setUp(self) -> None: + data = model_args().normalize_value( + { + "type_map": ["O", "H", "B"], + "use_srtab": f"{TESTS_DIR}/pt/water/data/zbl_tab_potential/H2O_tab_potential.txt", + "smin_alpha": 0.1, + "sw_rmin": 0.2, + "sw_rmax": 4.0, + "descriptor": { + "type": "se_atten", + "sel": 40, + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [3, 6], + "axis_neuron": 2, + "attn": 8, + "attn_layer": 2, + "attn_dotr": True, + "attn_mask": False, + "activation_function": "tanh", + "scaling_factor": 1.0, + "normalize": False, + "temperature": 1.0, + "set_davg_zero": True, + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "neuron": [5, 5], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + }, + }, + trim_pattern="_*", + ) + # Build dpmodel first, then deserialize into pt/pt_expt to share weights + self.dp_model = get_model_dp(data) + serialized = self.dp_model.serialize() + self.pt_model = DPZBLModelPT.deserialize(serialized) + self.pt_expt_model = DPZBLModelPTExpt.deserialize(serialized) + + # Coords / atype / box + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, -1, 3) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32).reshape(1, -1) + self.box = np.array( + [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, 9) + + # Build extended coords + nlist for lower-level calls + rcut = self.dp_model.get_rcut() + sel = self.dp_model.get_sel() + nframes, nloc = self.atype.shape[:2] + coord_normalized = normalize_coord( + self.coords.reshape(nframes, nloc, 3), + self.box.reshape(nframes, 3, 3), + ) + extended_coord, extended_atype, mapping = extend_coord_with_ghosts( + coord_normalized, self.atype, self.box, rcut + ) + nlist = build_neighbor_list( + extended_coord, + extended_atype, + nloc, + rcut, + sel, + distinguish_types=False, + ) + self.extended_coord = extended_coord.reshape(nframes, -1, 3) + self.extended_atype = extended_atype + self.mapping = mapping + self.nlist = nlist + + def test_translated_output_def(self) -> None: + """translated_output_def should return the same keys on dp, pt, and pt_expt.""" + dp_def = self.dp_model.translated_output_def() + pt_def = self.pt_model.translated_output_def() + pt_expt_def = self.pt_expt_model.translated_output_def() + self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + self.assertEqual(set(dp_def.keys()), set(pt_expt_def.keys())) + for key in dp_def: + self.assertEqual(dp_def[key].shape, pt_def[key].shape) + self.assertEqual(dp_def[key].shape, pt_expt_def[key].shape) + + def test_get_out_bias(self) -> None: + """get_out_bias should return numerically equal values on dp and pt.""" + dp_bias = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias = torch_to_numpy(self.pt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias, pt_bias, rtol=1e-10, atol=1e-10) + # Verify shape: ntypes=3 for ZBL model + self.assertEqual(dp_bias.shape[1], 3) + self.assertGreater(dp_bias.shape[0], 0) + + def test_set_out_bias(self) -> None: + """set_out_bias should update the bias on both backends.""" + dp_bias = to_numpy_array(self.dp_model.get_out_bias()) + new_bias = dp_bias + 1.0 + # dp + self.dp_model.set_out_bias(new_bias) + np.testing.assert_allclose( + to_numpy_array(self.dp_model.get_out_bias()), + new_bias, + rtol=1e-10, + atol=1e-10, + ) + # pt + self.pt_model.set_out_bias(numpy_to_torch(new_bias)) + np.testing.assert_allclose( + torch_to_numpy(self.pt_model.get_out_bias()), + new_bias, + rtol=1e-10, + atol=1e-10, + ) + + def test_model_output_def(self) -> None: + """model_output_def should return the same keys and shapes on dp and pt.""" + dp_def = self.dp_model.model_output_def().get_data() + pt_def = self.pt_model.model_output_def().get_data() + self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + for key in dp_def: + self.assertEqual(dp_def[key].shape, pt_def[key].shape) + + def test_model_output_type(self) -> None: + """model_output_type should return the same list on dp and pt.""" + self.assertEqual( + self.dp_model.model_output_type(), + self.pt_model.model_output_type(), + ) + + def test_do_grad_r(self) -> None: + """do_grad_r should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.do_grad_r("energy"), + self.pt_model.do_grad_r("energy"), + ) + self.assertTrue(self.dp_model.do_grad_r("energy")) + + def test_do_grad_c(self) -> None: + """do_grad_c should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.do_grad_c("energy"), + self.pt_model.do_grad_c("energy"), + ) + self.assertTrue(self.dp_model.do_grad_c("energy")) + + def test_get_rcut(self) -> None: + """get_rcut should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_rcut(), self.pt_model.get_rcut()) + + def test_get_type_map(self) -> None: + """get_type_map should return the same list on dp and pt.""" + self.assertEqual(self.dp_model.get_type_map(), self.pt_model.get_type_map()) + self.assertEqual(self.dp_model.get_type_map(), ["O", "H", "B"]) + + def test_get_sel(self) -> None: + """get_sel should return the same list on dp and pt.""" + self.assertEqual(self.dp_model.get_sel(), self.pt_model.get_sel()) + + def test_get_nsel(self) -> None: + """get_nsel should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_nsel(), self.pt_model.get_nsel()) + + def test_get_nnei(self) -> None: + """get_nnei should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_nnei(), self.pt_model.get_nnei()) + + def test_mixed_types(self) -> None: + """mixed_types should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.mixed_types(), self.pt_model.mixed_types()) + # DPZBLModel (LinearEnergyAtomicModel) always uses mixed types + self.assertTrue(self.dp_model.mixed_types()) + + def test_has_message_passing(self) -> None: + """has_message_passing should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.has_message_passing(), + self.pt_model.has_message_passing(), + ) + self.assertFalse(self.dp_model.has_message_passing()) + + def test_need_sorted_nlist_for_lower(self) -> None: + """need_sorted_nlist_for_lower should return the same value on dp and pt.""" + self.assertEqual( + self.dp_model.need_sorted_nlist_for_lower(), + self.pt_model.need_sorted_nlist_for_lower(), + ) + + def test_get_dim_fparam(self) -> None: + """get_dim_fparam should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_dim_fparam(), self.pt_model.get_dim_fparam()) + self.assertEqual(self.dp_model.get_dim_fparam(), 0) + + def test_get_dim_aparam(self) -> None: + """get_dim_aparam should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_dim_aparam(), self.pt_model.get_dim_aparam()) + self.assertEqual(self.dp_model.get_dim_aparam(), 0) + + def test_get_sel_type(self) -> None: + """get_sel_type should return the same list on dp and pt.""" + self.assertEqual(self.dp_model.get_sel_type(), self.pt_model.get_sel_type()) + + def test_is_aparam_nall(self) -> None: + """is_aparam_nall should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.is_aparam_nall(), self.pt_model.is_aparam_nall()) + self.assertFalse(self.dp_model.is_aparam_nall()) + + def test_get_model_def_script(self) -> None: + """get_model_def_script should return the same value on dp, pt, and pt_expt.""" + dp_val = self.dp_model.get_model_def_script() + pt_val = self.pt_model.get_model_def_script() + pe_val = self.pt_expt_model.get_model_def_script() + self.assertEqual(dp_val, pt_val) + self.assertEqual(dp_val, pe_val) + + def test_get_min_nbor_dist(self) -> None: + """get_min_nbor_dist should return the same value on dp, pt, and pt_expt.""" + dp_val = self.dp_model.get_min_nbor_dist() + pt_val = self.pt_model.get_min_nbor_dist() + pe_val = self.pt_expt_model.get_min_nbor_dist() + self.assertEqual(dp_val, pt_val) + self.assertEqual(dp_val, pe_val) + + def test_set_case_embd(self) -> None: + """set_case_embd should produce consistent results across backends. + + Also verifies that different case indices produce different outputs, + confirming the embedding is actually used. + """ + from deepmd.utils.argcheck import ( + model_args, + ) + + # Build a model with dim_case_embd > 0 + data = model_args().normalize_value( + { + "type_map": ["O", "H", "B"], + "use_srtab": f"{TESTS_DIR}/pt/water/data/zbl_tab_potential/H2O_tab_potential.txt", + "smin_alpha": 0.1, + "sw_rmin": 0.2, + "sw_rmax": 4.0, + "descriptor": { + "type": "se_atten", + "sel": 40, + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [3, 6], + "axis_neuron": 2, + "attn": 8, + "attn_layer": 2, + "attn_dotr": True, + "attn_mask": False, + "activation_function": "tanh", + "scaling_factor": 1.0, + "normalize": False, + "temperature": 1.0, + "set_davg_zero": True, + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "neuron": [5, 5], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + "dim_case_embd": 3, + }, + }, + trim_pattern="_*", + ) + dp_model = get_model_dp(data) + serialized = dp_model.serialize() + pt_model = DPZBLModelPT.deserialize(serialized) + pe_model = DPZBLModelPTExpt.deserialize(serialized) + + def _eval(case_idx): + dp_model.set_case_embd(case_idx) + pt_model.set_case_embd(case_idx) + pe_model.set_case_embd(case_idx) + dp_ret = dp_model(self.coords, self.atype, box=self.box) + pt_ret = { + k: torch_to_numpy(v) + for k, v in pt_model( + numpy_to_torch(self.coords), + numpy_to_torch(self.atype), + box=numpy_to_torch(self.box), + ).items() + } + coord_t = pt_expt_numpy_to_torch(self.coords) + coord_t.requires_grad_(True) + pe_ret = { + k: v.detach().cpu().numpy() + for k, v in pe_model( + coord_t, + pt_expt_numpy_to_torch(self.atype), + box=pt_expt_numpy_to_torch(self.box), + ).items() + } + return dp_ret, pt_ret, pe_ret + + dp0, pt0, pe0 = _eval(0) + dp1, pt1, pe1 = _eval(1) + + # Cross-backend consistency for each case index + for key in ("energy", "atom_energy"): + np.testing.assert_allclose( + dp0[key], + pt0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 0: dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp0[key], + pe0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 0: dp vs pt_expt mismatch in {key}", + ) + np.testing.assert_allclose( + dp1[key], + pt1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 1: dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp1[key], + pe1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"case 1: dp vs pt_expt mismatch in {key}", + ) + # Different case indices should produce different outputs + self.assertFalse( + np.allclose(dp0["energy"], dp1["energy"]), + "set_case_embd(0) and set_case_embd(1) produced the same energy", + ) + + def test_atomic_output_def(self) -> None: + """atomic_output_def should return the same keys and shapes on dp and pt.""" + dp_def = self.dp_model.atomic_output_def() + pt_def = self.pt_model.atomic_output_def() + self.assertEqual(set(dp_def.keys()), set(pt_def.keys())) + for key in dp_def.keys(): + self.assertEqual(dp_def[key].shape, pt_def[key].shape) + + def test_get_ntypes(self) -> None: + """get_ntypes should return the same value on dp and pt.""" + self.assertEqual(self.dp_model.get_ntypes(), self.pt_model.get_ntypes()) + self.assertEqual(self.dp_model.get_ntypes(), 3) + + def test_format_nlist(self) -> None: + """format_nlist should produce the same result on dp and pt.""" + dp_nlist = self.dp_model.format_nlist( + self.extended_coord, + self.extended_atype, + self.nlist, + ) + pt_nlist = torch_to_numpy( + self.pt_model.format_nlist( + numpy_to_torch(self.extended_coord), + numpy_to_torch(self.extended_atype), + numpy_to_torch(self.nlist), + ) + ) + np.testing.assert_equal(dp_nlist, pt_nlist) + + def test_forward_common_atomic(self) -> None: + """forward_common_atomic should produce consistent results on dp and pt. + + Compares at the atomic_model level, where both backends define this method. + DPZBLModel has no aparam, so we don't pass aparam here. + """ + dp_ret = self.dp_model.atomic_model.forward_common_atomic( + self.extended_coord, + self.extended_atype, + self.nlist, + mapping=self.mapping, + ) + pt_ret = self.pt_model.atomic_model.forward_common_atomic( + numpy_to_torch(self.extended_coord), + numpy_to_torch(self.extended_atype), + numpy_to_torch(self.nlist), + mapping=numpy_to_torch(self.mapping), + ) + # Compare the common keys + common_keys = set(dp_ret.keys()) & set(pt_ret.keys()) + self.assertTrue(len(common_keys) > 0) + for key in common_keys: + if dp_ret[key] is not None and pt_ret[key] is not None: + np.testing.assert_allclose( + dp_ret[key], + torch_to_numpy(pt_ret[key]), + rtol=1e-10, + atol=1e-10, + err_msg=f"Mismatch in forward_common_atomic key '{key}'", + ) + + def test_change_out_bias(self) -> None: + """change_out_bias should produce consistent bias on dp, pt, and pt_expt. + + Tests both set-by-statistic and change-by-statistic modes. + DPZBLModel has no fparam/aparam, so fitting stats checks are skipped. + """ + nframes = 2 + + # Use realistic coords (from setUp, tiled for 2 frames) + coords_2f = np.tile(self.coords, (nframes, 1, 1)) # (2, 6, 3) + atype_2f = np.array([[0, 0, 1, 1, 1, 1], [0, 1, 1, 0, 1, 1]], dtype=np.int32) + box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) + # natoms: [nloc, nloc, n_type0, n_type1, n_type2] — 3 types + natoms_data = np.array([[6, 6, 2, 4, 0], [6, 6, 2, 4, 0]], dtype=np.int32) + energy_data = np.array([10.0, 20.0]).reshape(nframes, 1) + + # dpmodel stat data (numpy) + dp_merged = [ + { + "coord": coords_2f, + "atype": atype_2f, + "atype_ext": atype_2f, + "box": box_2f, + "natoms": natoms_data, + "energy": energy_data, + "find_energy": np.float32(1.0), + } + ] + # pt stat data (torch tensors) + pt_merged = [ + { + "coord": numpy_to_torch(coords_2f), + "atype": numpy_to_torch(atype_2f), + "atype_ext": numpy_to_torch(atype_2f), + "box": numpy_to_torch(box_2f), + "natoms": numpy_to_torch(natoms_data), + "energy": numpy_to_torch(energy_data), + "find_energy": np.float32(1.0), + } + ] + # pt_expt stat data (numpy, same as dp) + pe_merged = dp_merged + + # Save initial (zero) bias + dp_bias_init = to_numpy_array(self.dp_model.get_out_bias()).copy() + + # --- Test "set-by-statistic" mode --- + self.dp_model.change_out_bias(dp_merged, bias_adjust_mode="set-by-statistic") + self.pt_model.change_out_bias(pt_merged, bias_adjust_mode="set-by-statistic") + self.pt_expt_model.change_out_bias( + pe_merged, bias_adjust_mode="set-by-statistic" + ) + + # Verify out bias consistency + dp_bias = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias = torch_to_numpy(self.pt_model.get_out_bias()) + pe_bias = to_numpy_array(self.pt_expt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias, pt_bias, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dp_bias, pe_bias, rtol=1e-10, atol=1e-10) + self.assertFalse( + np.allclose(dp_bias, dp_bias_init), + "set-by-statistic did not change the bias from initial values", + ) + + # --- Test "change-by-statistic" mode --- + dp_bias_before = dp_bias.copy() + self.dp_model.change_out_bias(dp_merged, bias_adjust_mode="change-by-statistic") + self.pt_model.change_out_bias(pt_merged, bias_adjust_mode="change-by-statistic") + self.pt_expt_model.change_out_bias( + pe_merged, bias_adjust_mode="change-by-statistic" + ) + + # Verify out bias consistency + dp_bias2 = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias2 = torch_to_numpy(self.pt_model.get_out_bias()) + pe_bias2 = to_numpy_array(self.pt_expt_model.get_out_bias()) + np.testing.assert_allclose(dp_bias2, pt_bias2, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dp_bias2, pe_bias2, rtol=1e-10, atol=1e-10) + self.assertFalse( + np.allclose(dp_bias2, dp_bias_before), + "change-by-statistic did not further change the bias", + ) + + # test_change_type_map: NOT applicable — PairTabAtomicModel does not + # support changing type map (would require rebuilding the tab file), + # so LinearEnergyAtomicModel.change_type_map always fails for DPZBLModel + # when the new type_map differs from the original. + + # test_change_type_map_extend_stat: NOT applicable — same reason. + + def test_update_sel(self) -> None: + """update_sel should return the same result on dp and pt.""" + from unittest.mock import ( + patch, + ) + + from deepmd.dpmodel.model.dp_model import DPModelCommon as DPModelCommonDP + from deepmd.pt.model.model.dp_model import DPModelCommon as DPModelCommonPT + + mock_min_nbor_dist = 0.5 + mock_sel = [30] + local_jdata = { + "type_map": ["O", "H", "B"], + "use_srtab": f"{TESTS_DIR}/pt/water/data/zbl_tab_potential/H2O_tab_potential.txt", + "smin_alpha": 0.1, + "sw_rmin": 0.2, + "sw_rmax": 4.0, + "descriptor": { + "type": "se_atten", + "sel": "auto", + "rcut_smth": 0.5, + "rcut": 4.0, + }, + "fitting_net": { + "neuron": [5, 5], + }, + } + type_map = ["O", "H", "B"] + + with patch( + "deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat", + return_value=(mock_min_nbor_dist, mock_sel), + ): + dp_result, dp_min_dist = DPModelCommonDP.update_sel( + None, type_map, local_jdata + ) + + with patch( + "deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat", + return_value=(mock_min_nbor_dist, mock_sel), + ): + pt_result, pt_min_dist = DPModelCommonPT.update_sel( + None, type_map, local_jdata + ) + + self.assertEqual(dp_result, pt_result) + self.assertEqual(dp_min_dist, pt_min_dist) + # Verify sel was actually updated (not still "auto") + self.assertNotEqual(dp_result["descriptor"]["sel"], "auto") + + def test_compute_or_load_out_stat(self) -> None: + """compute_or_load_out_stat should produce consistent bias on dp and pt. + + Tests both the compute path (from data) and the load path (from file). + """ + import tempfile + from pathlib import ( + Path, + ) + + import h5py + + from deepmd.utils.path import ( + DPPath, + ) + + nframes = 2 + coords_2f = np.tile(self.coords, (nframes, 1, 1)) + atype_2f = np.array([[0, 0, 1, 1, 1, 1], [0, 1, 1, 0, 1, 1]], dtype=np.int32) + box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) + # natoms: [nloc, nloc, n_type0, n_type1, n_type2] — 3 types + natoms_data = np.array([[6, 6, 2, 4, 0], [6, 6, 2, 4, 0]], dtype=np.int32) + energy_data = np.array([10.0, 20.0]).reshape(nframes, 1) + + dp_merged = [ + { + "coord": coords_2f, + "atype": atype_2f, + "atype_ext": atype_2f, + "box": box_2f, + "natoms": natoms_data, + "energy": energy_data, + "find_energy": np.float32(1.0), + } + ] + pt_merged = [ + { + "coord": numpy_to_torch(coords_2f), + "atype": numpy_to_torch(atype_2f), + "atype_ext": numpy_to_torch(atype_2f), + "box": numpy_to_torch(box_2f), + "natoms": numpy_to_torch(natoms_data), + "energy": numpy_to_torch(energy_data), + "find_energy": np.float32(1.0), + } + ] + + # Verify bias is initially identical + dp_bias_before = to_numpy_array(self.dp_model.get_out_bias()).copy() + pt_bias_before = torch_to_numpy(self.pt_model.get_out_bias()).copy() + np.testing.assert_allclose( + dp_bias_before, pt_bias_before, rtol=1e-10, atol=1e-10 + ) + + with tempfile.TemporaryDirectory() as tmpdir: + # Create separate h5 files for dp and pt + dp_h5 = str((Path(tmpdir) / "dp_stat.h5").resolve()) + pt_h5 = str((Path(tmpdir) / "pt_stat.h5").resolve()) + with h5py.File(dp_h5, "w"): + pass + with h5py.File(pt_h5, "w"): + pass + dp_stat_path = DPPath(dp_h5, "a") + pt_stat_path = DPPath(pt_h5, "a") + + # 1. Compute stats and save to file + self.dp_model.atomic_model.compute_or_load_out_stat( + dp_merged, stat_file_path=dp_stat_path + ) + self.pt_model.atomic_model.compute_or_load_out_stat( + pt_merged, stat_file_path=pt_stat_path + ) + + dp_bias_after = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias_after = torch_to_numpy(self.pt_model.get_out_bias()) + np.testing.assert_allclose( + dp_bias_after, pt_bias_after, rtol=1e-10, atol=1e-10 + ) + + # Verify bias actually changed (not still all zeros) + self.assertFalse( + np.allclose(dp_bias_after, dp_bias_before), + "compute_or_load_out_stat did not change the bias", + ) + + # 2. Verify both backends saved the same file content + with h5py.File(dp_h5, "r") as dp_f, h5py.File(pt_h5, "r") as pt_f: + dp_keys = sorted(dp_f.keys()) + pt_keys = sorted(pt_f.keys()) + self.assertEqual(dp_keys, pt_keys) + for key in dp_keys: + np.testing.assert_allclose( + np.array(dp_f[key]), + np.array(pt_f[key]), + rtol=1e-10, + atol=1e-10, + err_msg=f"Stat file content mismatch for key {key}", + ) + + # 3. Reset biases to zero, then load from file + zero_bias = np.zeros_like(dp_bias_after) + self.dp_model.set_out_bias(zero_bias) + self.pt_model.set_out_bias(numpy_to_torch(zero_bias)) + + # Use a callable that raises to ensure it loads from file, not recomputes + def raise_error(): + raise RuntimeError("Should not recompute — should load from file") + + self.dp_model.atomic_model.compute_or_load_out_stat( + raise_error, stat_file_path=dp_stat_path + ) + self.pt_model.atomic_model.compute_or_load_out_stat( + raise_error, stat_file_path=pt_stat_path + ) + + dp_bias_loaded = to_numpy_array(self.dp_model.get_out_bias()) + pt_bias_loaded = torch_to_numpy(self.pt_model.get_out_bias()) + + # Loaded biases should match between backends + np.testing.assert_allclose( + dp_bias_loaded, pt_bias_loaded, rtol=1e-10, atol=1e-10 + ) + # Loaded biases should match the originally computed biases + np.testing.assert_allclose( + dp_bias_loaded, dp_bias_after, rtol=1e-10, atol=1e-10 + ) + + def test_get_observed_type_list(self) -> None: + """get_observed_type_list should be consistent across dp, pt, pt_expt. + + Uses mock data containing only type 0 ("O") so that types 1 ("H") + and 2 ("B") are unobserved and should be absent from the returned list. + """ + nframes = 2 + natoms = 6 + # All atoms are type 0 — types 1, 2 are unobserved + atype_2f = np.zeros((nframes, natoms), dtype=np.int32) + coords_2f = np.tile(self.coords, (nframes, 1, 1)) + box_2f = np.tile(self.box.reshape(1, 3, 3), (nframes, 1, 1)) + natoms_data = np.array( + [[natoms, natoms, natoms, 0, 0]] * nframes, dtype=np.int32 + ) + energy_data = np.array([10.0, 20.0]).reshape(nframes, 1) + + dp_merged = [ + { + "coord": coords_2f, + "atype": atype_2f, + "atype_ext": atype_2f, + "box": box_2f, + "natoms": natoms_data, + "energy": energy_data, + "find_energy": np.float32(1.0), + } + ] + pt_merged = [ + { + "coord": numpy_to_torch(coords_2f), + "atype": numpy_to_torch(atype_2f), + "atype_ext": numpy_to_torch(atype_2f), + "box": numpy_to_torch(box_2f), + "natoms": numpy_to_torch(natoms_data), + "energy": numpy_to_torch(energy_data), + "find_energy": np.float32(1.0), + } + ] + + self.dp_model.atomic_model.compute_or_load_out_stat(dp_merged) + self.pt_model.atomic_model.compute_or_load_out_stat(pt_merged) + self.pt_expt_model.atomic_model.compute_or_load_out_stat(dp_merged) + + dp_observed = self.dp_model.get_observed_type_list() + pt_observed = self.pt_model.get_observed_type_list() + pe_observed = self.pt_expt_model.get_observed_type_list() + + self.assertEqual(dp_observed, pt_observed) + self.assertEqual(dp_observed, pe_observed) + # Only type 0 ("O") should be observed + self.assertEqual(dp_observed, ["O"]) + + +@parameterized( + (([], []), ([[0, 1]], [1])), # (pair_exclude_types, atom_exclude_types) +) +@unittest.skipUnless(INSTALLED_PT and INSTALLED_PT_EXPT, "PT and PT_EXPT are required") +class TestZBLComputeOrLoadStat(unittest.TestCase): + """Test that compute_or_load_stat produces identical statistics on dp, pt, and pt_expt. + + Covers descriptor stats (dstd) and output bias for DPZBLModel. + Parameterized over exclusion types only (no fparam — LinearEnergyAtomicModel + does not expose fitting_net for param stats). + """ + + def setUp(self) -> None: + ((pair_exclude_types, atom_exclude_types),) = self.param + data = model_args().normalize_value( + { + "type_map": ["O", "H", "B"], + "use_srtab": f"{TESTS_DIR}/pt/water/data/zbl_tab_potential/H2O_tab_potential.txt", + "smin_alpha": 0.1, + "sw_rmin": 0.2, + "sw_rmax": 4.0, + "pair_exclude_types": pair_exclude_types, + "atom_exclude_types": atom_exclude_types, + "descriptor": { + "type": "se_atten", + "sel": 40, + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [3, 6], + "axis_neuron": 2, + "attn": 8, + "attn_layer": 2, + "attn_dotr": True, + "attn_mask": False, + "activation_function": "tanh", + "scaling_factor": 1.0, + "normalize": False, + "temperature": 1.0, + "set_davg_zero": True, + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "neuron": [5, 5], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + }, + }, + trim_pattern="_*", + ) + + # Save data for reuse in load-from-file test + self._model_data = data + + # Build dp model, then deserialize into pt and pt_expt to share weights + self.dp_model = get_model_dp(data) + serialized = self.dp_model.serialize() + self.pt_model = DPZBLModelPT.deserialize(serialized) + self.pt_expt_model = DPZBLModelPTExpt.deserialize(serialized) + + # Test coords / atype / box for forward evaluation + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 0.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, -1, 3) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32).reshape(1, -1) + self.box = np.array( + [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, 9) + + # Mock training data for compute_or_load_stat + natoms = 6 + nframes = 3 + rng = np.random.default_rng(42) + coords_stat = rng.normal(size=(nframes, natoms, 3)).astype( + GLOBAL_NP_FLOAT_PRECISION + ) + atype_stat = np.array([[0, 0, 1, 1, 1, 1]] * nframes, dtype=np.int32) + box_stat = np.tile( + np.eye(3, dtype=GLOBAL_NP_FLOAT_PRECISION).reshape(1, 3, 3) * 13.0, + (nframes, 1, 1), + ) + natoms_stat = np.array( + [[natoms, natoms, 2, 4, 0]] * nframes, dtype=np.int32 + ) # 3 types: O=2, H=4, B=0 + energy_stat = rng.normal(size=(nframes, 1)).astype(GLOBAL_NP_FLOAT_PRECISION) + + # dp / pt_expt sample (numpy) + np_sample = { + "coord": coords_stat, + "atype": atype_stat, + "atype_ext": atype_stat, + "box": box_stat, + "natoms": natoms_stat, + "energy": energy_stat, + "find_energy": np.float32(1.0), + } + # pt sample (torch tensors) + pt_sample = { + "coord": numpy_to_torch(coords_stat), + "atype": numpy_to_torch(atype_stat), + "atype_ext": numpy_to_torch(atype_stat), + "box": numpy_to_torch(box_stat), + "natoms": numpy_to_torch(natoms_stat), + "energy": numpy_to_torch(energy_stat), + "find_energy": np.float32(1.0), + } + + self.np_sampled = [np_sample] + self.pt_sampled = [pt_sample] + + def _eval_dp(self) -> dict: + return self.dp_model(self.coords, self.atype, box=self.box) + + def _eval_pt(self) -> dict: + return { + kk: torch_to_numpy(vv) + for kk, vv in self.pt_model( + numpy_to_torch(self.coords), + numpy_to_torch(self.atype), + box=numpy_to_torch(self.box), + do_atomic_virial=True, + ).items() + } + + def _eval_pt_expt(self) -> dict: + coord_t = pt_expt_numpy_to_torch(self.coords) + coord_t.requires_grad_(True) + return { + k: v.detach().cpu().numpy() + for k, v in self.pt_expt_model( + coord_t, + pt_expt_numpy_to_torch(self.atype), + box=pt_expt_numpy_to_torch(self.box), + do_atomic_virial=True, + ).items() + } + + def test_compute_stat(self) -> None: + # 1. Pre-stat forward consistency + dp_ret0 = self._eval_dp() + pt_ret0 = self._eval_pt() + pe_ret0 = self._eval_pt_expt() + for key in ("energy", "atom_energy"): + np.testing.assert_allclose( + dp_ret0[key], + pt_ret0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Pre-stat dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp_ret0[key], + pe_ret0[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Pre-stat dp vs pt_expt mismatch in {key}", + ) + + # 2. Run compute_or_load_stat on all three backends + # deepcopy samples: the ZBL model's stat path mutates natoms in-place + # (stat.py applies atom_exclude_types mask via natoms[:, 2:] *= type_mask), + # so each backend must receive its own copy. + self.dp_model.compute_or_load_stat(lambda: copy.deepcopy(self.np_sampled)) + self.pt_model.compute_or_load_stat(lambda: copy.deepcopy(self.pt_sampled)) + self.pt_expt_model.compute_or_load_stat(lambda: copy.deepcopy(self.np_sampled)) + + # 3. Serialize all three and compare @variables + dp_ser = self.dp_model.serialize() + pt_ser = self.pt_model.serialize() + pe_ser = self.pt_expt_model.serialize() + compare_variables_recursive(dp_ser, pt_ser) + compare_variables_recursive(dp_ser, pe_ser) + + # 4. Post-stat forward consistency + dp_ret1 = self._eval_dp() + pt_ret1 = self._eval_pt() + pe_ret1 = self._eval_pt_expt() + for key in ("energy", "atom_energy"): + np.testing.assert_allclose( + dp_ret1[key], + pt_ret1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Post-stat dp vs pt mismatch in {key}", + ) + np.testing.assert_allclose( + dp_ret1[key], + pe_ret1[key], + rtol=1e-10, + atol=1e-10, + err_msg=f"Post-stat dp vs pt_expt mismatch in {key}", + ) + + def test_load_stat_from_file(self) -> None: + import tempfile + from pathlib import ( + Path, + ) + + import h5py + + from deepmd.utils.path import ( + DPPath, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + # Create separate stat files for each backend + dp_h5 = str((Path(tmpdir) / "dp_stat.h5").resolve()) + pt_h5 = str((Path(tmpdir) / "pt_stat.h5").resolve()) + pe_h5 = str((Path(tmpdir) / "pe_stat.h5").resolve()) + for p in (dp_h5, pt_h5, pe_h5): + with h5py.File(p, "w"): + pass + + # 1. Compute stats and save to file (deepcopy: stat path mutates natoms) + self.dp_model.compute_or_load_stat( + lambda: copy.deepcopy(self.np_sampled), + stat_file_path=DPPath(dp_h5, "a"), + ) + self.pt_model.compute_or_load_stat( + lambda: copy.deepcopy(self.pt_sampled), + stat_file_path=DPPath(pt_h5, "a"), + ) + self.pt_expt_model.compute_or_load_stat( + lambda: copy.deepcopy(self.np_sampled), + stat_file_path=DPPath(pe_h5, "a"), + ) + + # Save the computed serializations as reference + dp_ser_computed = self.dp_model.serialize() + pt_ser_computed = self.pt_model.serialize() + pe_ser_computed = self.pt_expt_model.serialize() + + # 2. Build fresh models from the same initial weights + dp_model2 = get_model_dp(self._model_data) + pt_model2 = DPZBLModelPT.deserialize(dp_model2.serialize()) + pe_model2 = DPZBLModelPTExpt.deserialize(dp_model2.serialize()) + + # 3. Load stats from file (should NOT call the sampled func) + def raise_error(): + raise RuntimeError("Should load from file, not recompute") + + dp_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(dp_h5, "a") + ) + pt_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(pt_h5, "a") + ) + pe_model2.compute_or_load_stat( + raise_error, stat_file_path=DPPath(pe_h5, "a") + ) + + # 4. Loaded models should match the computed ones + dp_ser_loaded = dp_model2.serialize() + pt_ser_loaded = pt_model2.serialize() + pe_ser_loaded = pe_model2.serialize() + compare_variables_recursive(dp_ser_computed, dp_ser_loaded) + compare_variables_recursive(pt_ser_computed, pt_ser_loaded) + compare_variables_recursive(pe_ser_computed, pe_ser_loaded) + + # 5. Cross-backend consistency after loading + compare_variables_recursive(dp_ser_loaded, pt_ser_loaded) + compare_variables_recursive(dp_ser_loaded, pe_ser_loaded) diff --git a/source/tests/pt_expt/atomic_model/test_dp_atomic_model.py b/source/tests/pt_expt/atomic_model/test_dp_atomic_model.py deleted file mode 100644 index 0196170cd0..0000000000 --- a/source/tests/pt_expt/atomic_model/test_dp_atomic_model.py +++ /dev/null @@ -1,288 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import itertools -import unittest - -import numpy as np -import torch - -from deepmd.dpmodel.atomic_model import DPAtomicModel as DPDPAtomicModel -from deepmd.dpmodel.descriptor import DescrptSeA as DPDescrptSeA -from deepmd.dpmodel.fitting import InvarFitting as DPInvarFitting -from deepmd.pt_expt.atomic_model import ( - DPAtomicModel, -) -from deepmd.pt_expt.descriptor.se_e2_a import ( - DescrptSeA, -) -from deepmd.pt_expt.fitting import ( - InvarFitting, -) -from deepmd.pt_expt.utils import ( - env, -) - -from ...pt.model.test_env_mat import ( - TestCaseSingleFrameWithNlist, - TestCaseSingleFrameWithNlistWithVirtual, -) -from ...seed import ( - GLOBAL_SEED, -) - - -class TestDPAtomicModel(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self) -> None: - TestCaseSingleFrameWithNlist.setUp(self) - self.device = env.DEVICE - - def test_self_consistency(self) -> None: - """Test that pt_expt atomic model serialize/deserialize preserves behavior.""" - nf, nloc, nnei = self.nlist.shape - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ).to(self.device) - ft = InvarFitting( - "energy", - self.nt, - ds.get_dim_out(), - 1, - mixed_types=ds.mixed_types(), - seed=GLOBAL_SEED, - ).to(self.device) - type_map = ["foo", "bar"] - - # test the case of exclusion - for atom_excl, pair_excl in itertools.product([[], [1]], [[], [[0, 1]]]): - md0 = DPAtomicModel( - ds, - ft, - type_map=type_map, - ).to(self.device) - md0.reinit_atom_exclude(atom_excl) - md0.reinit_pair_exclude(pair_excl) - md1 = DPAtomicModel.deserialize(md0.serialize()).to(self.device) - - # Test forward pass - args = [ - torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device), - torch.tensor(self.atype_ext, dtype=torch.int64, device=self.device), - torch.tensor(self.nlist, dtype=torch.int64, device=self.device), - ] - ret0 = md0.forward_common_atomic(*args) - ret1 = md1.forward_common_atomic(*args) - np.testing.assert_allclose( - ret0["energy"].detach().cpu().numpy(), - ret1["energy"].detach().cpu().numpy(), - ) - - def test_dp_consistency(self) -> None: - """Test numerical consistency between dpmodel and pt_expt atomic models.""" - nf, nloc, nnei = self.nlist.shape - ds = DPDescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ) - ft = DPInvarFitting( - "energy", - self.nt, - ds.get_dim_out(), - 1, - mixed_types=ds.mixed_types(), - seed=GLOBAL_SEED, - ) - type_map = ["foo", "bar"] - md0 = DPDPAtomicModel(ds, ft, type_map=type_map) - md1 = DPAtomicModel.deserialize(md0.serialize()).to(self.device) - - # dpmodel uses numpy arrays - args0 = [self.coord_ext, self.atype_ext, self.nlist] - # pt_expt uses torch tensors - args1 = [ - torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device), - torch.tensor(self.atype_ext, dtype=torch.int64, device=self.device), - torch.tensor(self.nlist, dtype=torch.int64, device=self.device), - ] - ret0 = md0.forward_common_atomic(*args0) - ret1 = md1.forward_common_atomic(*args1) - np.testing.assert_allclose( - ret0["energy"], - ret1["energy"].detach().cpu().numpy(), - ) - - def test_exportable(self) -> None: - """Test that pt_expt atomic model can be exported with torch.export.""" - nf, nloc, nnei = self.nlist.shape - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ).to(self.device) - ft = InvarFitting( - "energy", - self.nt, - ds.get_dim_out(), - 1, - mixed_types=ds.mixed_types(), - seed=GLOBAL_SEED, - ).to(self.device) - type_map = ["foo", "bar"] - md0 = DPAtomicModel(ds, ft, type_map=type_map).to(self.device) - md0 = md0.eval() - - # Prepare inputs for export - coord = torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device) - atype = torch.tensor(self.atype_ext, dtype=torch.int64, device=self.device) - nlist = torch.tensor(self.nlist, dtype=torch.int64, device=self.device) - - # Test forward pass - ret0 = md0(coord, atype, nlist) - self.assertIn("energy", ret0) - - # Test torch.export - # Use strict=False for now to handle dynamic shapes - exported = torch.export.export( - md0, - (coord, atype, nlist), - strict=False, - ) - self.assertIsNotNone(exported) - - # Test exported model produces same output - ret1 = exported.module()(coord, atype, nlist) - np.testing.assert_allclose( - ret0["energy"].detach().cpu().numpy(), - ret1["energy"].detach().cpu().numpy(), - rtol=1e-10, - atol=1e-10, - ) - - def test_excl_consistency(self) -> None: - """Test that exclusion masks work correctly after serialize/deserialize.""" - type_map = ["foo", "bar"] - - # test the case of exclusion - for atom_excl, pair_excl in itertools.product([[], [1]], [[], [[0, 1]]]): - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ).to(self.device) - ft = InvarFitting( - "energy", - self.nt, - ds.get_dim_out(), - 1, - mixed_types=ds.mixed_types(), - seed=GLOBAL_SEED, - ).to(self.device) - md0 = DPAtomicModel( - ds, - ft, - type_map=type_map, - ).to(self.device) - md1 = DPAtomicModel.deserialize(md0.serialize()).to(self.device) - - md0.reinit_atom_exclude(atom_excl) - md0.reinit_pair_exclude(pair_excl) - # hacking! - md1.descriptor.reinit_exclude(pair_excl) - md1.fitting.reinit_exclude(atom_excl) - - # check energy consistency - args = [ - torch.tensor(self.coord_ext, dtype=torch.float64, device=self.device), - torch.tensor(self.atype_ext, dtype=torch.int64, device=self.device), - torch.tensor(self.nlist, dtype=torch.int64, device=self.device), - ] - ret0 = md0.forward_common_atomic(*args) - ret1 = md1.forward_common_atomic(*args) - np.testing.assert_allclose( - ret0["energy"].detach().cpu().numpy(), - ret1["energy"].detach().cpu().numpy(), - ) - - # check output def - out_names = [vv.name for vv in md0.atomic_output_def().get_data().values()] - self.assertEqual(out_names, ["energy", "mask"]) - if atom_excl != []: - for ii in md0.atomic_output_def().get_data().values(): - if ii.name == "mask": - self.assertEqual(ii.shape, [1]) - self.assertFalse(ii.reducible) - self.assertFalse(ii.r_differentiable) - self.assertFalse(ii.c_differentiable) - - # check mask - if atom_excl == []: - pass - elif atom_excl == [1]: - self.assertIn("mask", ret0.keys()) - expected = np.array([1, 1, 0], dtype=int) - expected = np.concatenate( - [expected, expected[self.perm[: self.nloc]]] - ).reshape(2, 3) - np.testing.assert_array_equal( - ret0["mask"].detach().cpu().numpy(), expected - ) - else: - raise ValueError(f"not expected atom_excl {atom_excl}") - - -class TestDPAtomicModelVirtualConsistency(unittest.TestCase): - def setUp(self) -> None: - self.case0 = TestCaseSingleFrameWithNlist() - self.case1 = TestCaseSingleFrameWithNlistWithVirtual() - self.case0.setUp() - self.case1.setUp() - self.device = env.DEVICE - - def test_virtual_consistency(self) -> None: - nf, _, _ = self.case0.nlist.shape - ds = DescrptSeA( - self.case0.rcut, - self.case0.rcut_smth, - self.case0.sel, - ) - ft = InvarFitting( - "energy", - self.case0.nt, - ds.get_dim_out(), - 1, - mixed_types=ds.mixed_types(), - seed=GLOBAL_SEED, - ) - type_map = ["foo", "bar"] - md1 = DPAtomicModel(ds, ft, type_map=type_map).to(self.device) - - args0 = [ - torch.tensor(self.case0.coord_ext, dtype=torch.float64, device=self.device), - torch.tensor(self.case0.atype_ext, dtype=torch.int64, device=self.device), - torch.tensor(self.case0.nlist, dtype=torch.int64, device=self.device), - ] - args1 = [ - torch.tensor(self.case1.coord_ext, dtype=torch.float64, device=self.device), - torch.tensor(self.case1.atype_ext, dtype=torch.int64, device=self.device), - torch.tensor(self.case1.nlist, dtype=torch.int64, device=self.device), - ] - - ret0 = md1.forward_common_atomic(*args0) - ret1 = md1.forward_common_atomic(*args1) - - for dd in range(self.case0.nf): - np.testing.assert_allclose( - ret0["energy"][dd].detach().cpu().numpy(), - ret1["energy"][dd, self.case1.get_real_mapping[dd], :] - .detach() - .cpu() - .numpy(), - ) - expected_mask = np.array( - [ - [1, 0, 1, 1], - [1, 1, 0, 1], - ] - ) - np.testing.assert_equal(ret1["mask"].detach().cpu().numpy(), expected_mask) diff --git a/source/tests/pt_expt/model/test_dipole_model.py b/source/tests/pt_expt/model/test_dipole_model.py new file mode 100644 index 0000000000..4dafd9d0ae --- /dev/null +++ b/source/tests/pt_expt/model/test_dipole_model.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np +import torch + +from deepmd.dpmodel.descriptor import DescrptSeA as DPDescrptSeA +from deepmd.dpmodel.fitting import DipoleFitting as DPDipoleFitting +from deepmd.dpmodel.model.dipole_model import DipoleModel as DPDipoleModel +from deepmd.dpmodel.utils.nlist import ( + build_neighbor_list, + extend_coord_with_ghosts, +) +from deepmd.dpmodel.utils.region import ( + normalize_coord, +) +from deepmd.pt_expt.model import ( + DipoleModel, +) +from deepmd.pt_expt.utils import ( + env, +) + +from ...seed import ( + GLOBAL_SEED, +) + + +class TestDipoleModel(unittest.TestCase): + def setUp(self) -> None: + self.device = env.DEVICE + self.natoms = 5 + self.rcut = 4.0 + self.rcut_smth = 0.5 + self.sel = [8, 6] + self.nt = 2 + self.type_map = ["foo", "bar"] + + generator = torch.Generator(device=self.device).manual_seed(GLOBAL_SEED) + cell = torch.rand( + [3, 3], dtype=torch.float64, device=self.device, generator=generator + ) + cell = (cell + cell.T) + 5.0 * torch.eye(3, device=self.device) + self.cell = cell.unsqueeze(0) + coord = torch.rand( + [self.natoms, 3], + dtype=torch.float64, + device=self.device, + generator=generator, + ) + coord = torch.matmul(coord, cell) + self.coord = coord.unsqueeze(0).to(self.device) + self.atype = torch.tensor( + [[0, 0, 0, 1, 1]], dtype=torch.int64, device=self.device + ) + + def _make_dp_model(self): + ds = DPDescrptSeA(self.rcut, self.rcut_smth, self.sel) + ft = DPDipoleFitting( + self.nt, + ds.get_dim_out(), + embedding_width=ds.get_dim_emb(), + seed=GLOBAL_SEED, + ) + return DPDipoleModel(ds, ft, type_map=self.type_map) + + def _prepare_lower_inputs(self): + coord_np = self.coord.detach().cpu().numpy() + atype_np = self.atype.detach().cpu().numpy() + cell_np = self.cell.reshape(1, 9).detach().cpu().numpy() + coord_normalized = normalize_coord( + coord_np.reshape(1, self.natoms, 3), + cell_np.reshape(1, 3, 3), + ) + extended_coord, extended_atype, mapping = extend_coord_with_ghosts( + coord_normalized, atype_np, cell_np, self.rcut + ) + nlist = build_neighbor_list( + extended_coord, + extended_atype, + self.natoms, + self.rcut, + self.sel, + distinguish_types=True, + ) + extended_coord = extended_coord.reshape(1, -1, 3) + return ( + torch.tensor(extended_coord, dtype=torch.float64, device=self.device), + torch.tensor(extended_atype, dtype=torch.int64, device=self.device), + torch.tensor(nlist, dtype=torch.int64, device=self.device), + torch.tensor(mapping, dtype=torch.int64, device=self.device), + ) + + def test_dp_consistency(self) -> None: + md_dp = self._make_dp_model() + md_pt = DipoleModel.deserialize(md_dp.serialize()).to(self.device) + md_pt.eval() + + coord_np = self.coord.detach().cpu().numpy() + atype_np = self.atype.detach().cpu().numpy() + cell_np = self.cell.reshape(1, 9).detach().cpu().numpy() + ret_dp = md_dp(coord_np.reshape(1, -1), atype_np, cell_np) + + coord = self.coord.clone().requires_grad_(True) + ret_pt = md_pt(coord, self.atype, self.cell.reshape(1, 9)) + + np.testing.assert_allclose( + ret_dp["global_dipole"], + ret_pt["global_dipole"].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + ret_dp["dipole"], + ret_pt["dipole"].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + ) + + def test_output_keys(self) -> None: + md_dp = self._make_dp_model() + md_pt = DipoleModel.deserialize(md_dp.serialize()).to(self.device) + md_pt.eval() + coord = self.coord.clone().requires_grad_(True) + ret = md_pt(coord, self.atype, self.cell.reshape(1, 9)) + self.assertIn("dipole", ret) + self.assertIn("global_dipole", ret) + + def test_forward_lower_exportable(self) -> None: + md_dp = self._make_dp_model() + md_pt = DipoleModel.deserialize(md_dp.serialize()).to(self.device) + md_pt.eval() + + ext_coord, ext_atype, nlist_t, mapping_t = self._prepare_lower_inputs() + fparam = None + aparam = None + + ret_eager = md_pt.forward_lower( + ext_coord.requires_grad_(True), + ext_atype, + nlist_t, + mapping_t, + fparam=fparam, + aparam=aparam, + ) + + traced = md_pt.forward_lower_exportable( + ext_coord, + ext_atype, + nlist_t, + mapping_t, + fparam=fparam, + aparam=aparam, + ) + self.assertIsInstance(traced, torch.nn.Module) + + exported = torch.export.export( + traced, + (ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam), + strict=False, + ) + self.assertIsNotNone(exported) + + ret_traced = traced(ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam) + ret_exported = exported.module()( + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam + ) + + for key in ("dipole", "global_dipole"): + np.testing.assert_allclose( + ret_eager[key].detach().cpu().numpy(), + ret_traced[key].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=f"traced vs eager: {key}", + ) + np.testing.assert_allclose( + ret_eager[key].detach().cpu().numpy(), + ret_exported[key].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=f"exported vs eager: {key}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt_expt/model/test_dos_model.py b/source/tests/pt_expt/model/test_dos_model.py new file mode 100644 index 0000000000..993c55972b --- /dev/null +++ b/source/tests/pt_expt/model/test_dos_model.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np +import torch + +from deepmd.dpmodel.descriptor import DescrptSeA as DPDescrptSeA +from deepmd.dpmodel.fitting import DOSFittingNet as DPDOSFittingNet +from deepmd.dpmodel.model.dos_model import DOSModel as DPDOSModel +from deepmd.dpmodel.utils.nlist import ( + build_neighbor_list, + extend_coord_with_ghosts, +) +from deepmd.dpmodel.utils.region import ( + normalize_coord, +) +from deepmd.pt_expt.model import ( + DOSModel, +) +from deepmd.pt_expt.utils import ( + env, +) + +from ...seed import ( + GLOBAL_SEED, +) + + +class TestDOSModel(unittest.TestCase): + def setUp(self) -> None: + self.device = env.DEVICE + self.natoms = 5 + self.rcut = 4.0 + self.rcut_smth = 0.5 + self.sel = [8, 6] + self.nt = 2 + self.type_map = ["foo", "bar"] + + generator = torch.Generator(device=self.device).manual_seed(GLOBAL_SEED) + cell = torch.rand( + [3, 3], dtype=torch.float64, device=self.device, generator=generator + ) + cell = (cell + cell.T) + 5.0 * torch.eye(3, device=self.device) + self.cell = cell.unsqueeze(0) + coord = torch.rand( + [self.natoms, 3], + dtype=torch.float64, + device=self.device, + generator=generator, + ) + coord = torch.matmul(coord, cell) + self.coord = coord.unsqueeze(0).to(self.device) + self.atype = torch.tensor( + [[0, 0, 0, 1, 1]], dtype=torch.int64, device=self.device + ) + + def _make_dp_model(self): + ds = DPDescrptSeA(self.rcut, self.rcut_smth, self.sel) + ft = DPDOSFittingNet( + self.nt, + ds.get_dim_out(), + numb_dos=10, + seed=GLOBAL_SEED, + ) + return DPDOSModel(ds, ft, type_map=self.type_map) + + def _prepare_lower_inputs(self): + coord_np = self.coord.detach().cpu().numpy() + atype_np = self.atype.detach().cpu().numpy() + cell_np = self.cell.reshape(1, 9).detach().cpu().numpy() + coord_normalized = normalize_coord( + coord_np.reshape(1, self.natoms, 3), + cell_np.reshape(1, 3, 3), + ) + extended_coord, extended_atype, mapping = extend_coord_with_ghosts( + coord_normalized, atype_np, cell_np, self.rcut + ) + nlist = build_neighbor_list( + extended_coord, + extended_atype, + self.natoms, + self.rcut, + self.sel, + distinguish_types=True, + ) + extended_coord = extended_coord.reshape(1, -1, 3) + return ( + torch.tensor(extended_coord, dtype=torch.float64, device=self.device), + torch.tensor(extended_atype, dtype=torch.int64, device=self.device), + torch.tensor(nlist, dtype=torch.int64, device=self.device), + torch.tensor(mapping, dtype=torch.int64, device=self.device), + ) + + def test_dp_consistency(self) -> None: + md_dp = self._make_dp_model() + md_pt = DOSModel.deserialize(md_dp.serialize()).to(self.device) + md_pt.eval() + + coord_np = self.coord.detach().cpu().numpy() + atype_np = self.atype.detach().cpu().numpy() + cell_np = self.cell.reshape(1, 9).detach().cpu().numpy() + ret_dp = md_dp(coord_np.reshape(1, -1), atype_np, cell_np) + + coord = self.coord.clone().requires_grad_(True) + ret_pt = md_pt(coord, self.atype, self.cell.reshape(1, 9)) + + np.testing.assert_allclose( + ret_dp["dos"], + ret_pt["dos"].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + ret_dp["atom_dos"], + ret_pt["atom_dos"].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + ) + + def test_output_keys(self) -> None: + md_dp = self._make_dp_model() + md_pt = DOSModel.deserialize(md_dp.serialize()).to(self.device) + md_pt.eval() + coord = self.coord.clone().requires_grad_(True) + ret = md_pt(coord, self.atype, self.cell.reshape(1, 9)) + self.assertIn("dos", ret) + self.assertIn("atom_dos", ret) + + def test_forward_lower_exportable(self) -> None: + md_dp = self._make_dp_model() + md_pt = DOSModel.deserialize(md_dp.serialize()).to(self.device) + md_pt.eval() + + ext_coord, ext_atype, nlist_t, mapping_t = self._prepare_lower_inputs() + fparam = None + aparam = None + + ret_eager = md_pt.forward_lower( + ext_coord.requires_grad_(True), + ext_atype, + nlist_t, + mapping_t, + fparam=fparam, + aparam=aparam, + ) + + traced = md_pt.forward_lower_exportable( + ext_coord, + ext_atype, + nlist_t, + mapping_t, + fparam=fparam, + aparam=aparam, + ) + self.assertIsInstance(traced, torch.nn.Module) + + exported = torch.export.export( + traced, + (ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam), + strict=False, + ) + self.assertIsNotNone(exported) + + ret_traced = traced(ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam) + ret_exported = exported.module()( + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam + ) + + for key in ("atom_dos", "dos"): + np.testing.assert_allclose( + ret_eager[key].detach().cpu().numpy(), + ret_traced[key].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=f"traced vs eager: {key}", + ) + np.testing.assert_allclose( + ret_eager[key].detach().cpu().numpy(), + ret_exported[key].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=f"exported vs eager: {key}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt_expt/model/test_dp_zbl_model.py b/source/tests/pt_expt/model/test_dp_zbl_model.py new file mode 100644 index 0000000000..1fa1d332e5 --- /dev/null +++ b/source/tests/pt_expt/model/test_dp_zbl_model.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import os +import unittest + +import numpy as np +import torch + +from deepmd.dpmodel.atomic_model import ( + DPAtomicModel, +) +from deepmd.dpmodel.atomic_model.pairtab_atomic_model import ( + PairTabAtomicModel, +) +from deepmd.dpmodel.descriptor import DescrptDPA1 as DPDescrptDPA1 +from deepmd.dpmodel.fitting import InvarFitting as DPInvarFitting +from deepmd.dpmodel.model.dp_zbl_model import DPZBLModel as DPDPZBLModel +from deepmd.dpmodel.utils.nlist import ( + build_neighbor_list, + extend_coord_with_ghosts, +) +from deepmd.dpmodel.utils.region import ( + normalize_coord, +) +from deepmd.pt_expt.model import ( + DPZBLModel, +) +from deepmd.pt_expt.utils import ( + env, +) + +from ...seed import ( + GLOBAL_SEED, +) + +TESTS_DIR = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) +TAB_FILE = os.path.join( + TESTS_DIR, + "pt", + "model", + "water", + "data", + "zbl_tab_potential", + "H2O_tab_potential.txt", +) + + +class TestDPZBLModel(unittest.TestCase): + def setUp(self) -> None: + self.device = env.DEVICE + self.natoms = 5 + self.rcut = 4.0 + self.rcut_smth = 0.5 + self.sel = 20 + self.nt = 3 + self.type_map = ["O", "H", "B"] + + generator = torch.Generator(device=self.device).manual_seed(GLOBAL_SEED) + cell = torch.rand( + [3, 3], dtype=torch.float64, device=self.device, generator=generator + ) + cell = (cell + cell.T) + 5.0 * torch.eye(3, device=self.device) + self.cell = cell.unsqueeze(0) + coord = torch.rand( + [self.natoms, 3], + dtype=torch.float64, + device=self.device, + generator=generator, + ) + coord = torch.matmul(coord, cell) + self.coord = coord.unsqueeze(0).to(self.device) + self.atype = torch.tensor( + [[0, 0, 1, 1, 2]], dtype=torch.int64, device=self.device + ) + + def _make_dp_model(self): + ds = DPDescrptDPA1( + rcut_smth=self.rcut_smth, + rcut=self.rcut, + sel=self.sel, + ntypes=self.nt, + neuron=[3, 6], + axis_neuron=2, + attn=4, + attn_layer=2, + attn_dotr=True, + attn_mask=False, + activation_function="tanh", + set_davg_zero=True, + type_one_side=True, + seed=GLOBAL_SEED, + ) + ft = DPInvarFitting( + "energy", + self.nt, + ds.get_dim_out(), + 1, + mixed_types=ds.mixed_types(), + seed=GLOBAL_SEED, + ) + dp_model = DPAtomicModel(ds, ft, type_map=self.type_map) + zbl_model = PairTabAtomicModel( + tab_file=TAB_FILE, + rcut=self.rcut, + sel=self.sel, + type_map=self.type_map, + ) + return DPDPZBLModel( + dp_model, + zbl_model, + sw_rmin=0.2, + sw_rmax=4.0, + type_map=self.type_map, + ) + + def _prepare_lower_inputs(self): + coord_np = self.coord.detach().cpu().numpy() + atype_np = self.atype.detach().cpu().numpy() + cell_np = self.cell.reshape(1, 9).detach().cpu().numpy() + coord_normalized = normalize_coord( + coord_np.reshape(1, self.natoms, 3), + cell_np.reshape(1, 3, 3), + ) + extended_coord, extended_atype, mapping = extend_coord_with_ghosts( + coord_normalized, atype_np, cell_np, self.rcut + ) + nlist = build_neighbor_list( + extended_coord, + extended_atype, + self.natoms, + self.rcut, + [self.sel], + distinguish_types=False, + ) + extended_coord = extended_coord.reshape(1, -1, 3) + return ( + torch.tensor(extended_coord, dtype=torch.float64, device=self.device), + torch.tensor(extended_atype, dtype=torch.int64, device=self.device), + torch.tensor(nlist, dtype=torch.int64, device=self.device), + torch.tensor(mapping, dtype=torch.int64, device=self.device), + ) + + def test_dp_consistency(self) -> None: + md_dp = self._make_dp_model() + md_pt = DPZBLModel.deserialize(md_dp.serialize()).to(self.device) + md_pt.eval() + + coord_np = self.coord.detach().cpu().numpy() + atype_np = self.atype.detach().cpu().numpy() + cell_np = self.cell.reshape(1, 9).detach().cpu().numpy() + ret_dp = md_dp(coord_np.reshape(1, -1), atype_np, cell_np) + + coord = self.coord.clone().requires_grad_(True) + ret_pt = md_pt(coord, self.atype, self.cell.reshape(1, 9)) + + np.testing.assert_allclose( + ret_dp["energy"], + ret_pt["energy"].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + ret_dp["atom_energy"], + ret_pt["atom_energy"].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + ) + + def test_output_keys(self) -> None: + md_dp = self._make_dp_model() + md_pt = DPZBLModel.deserialize(md_dp.serialize()).to(self.device) + md_pt.eval() + coord = self.coord.clone().requires_grad_(True) + ret = md_pt(coord, self.atype, self.cell.reshape(1, 9)) + self.assertIn("energy", ret) + self.assertIn("atom_energy", ret) + self.assertIn("force", ret) + self.assertIn("virial", ret) + + def test_forward_lower_exportable(self) -> None: + md_dp = self._make_dp_model() + md_pt = DPZBLModel.deserialize(md_dp.serialize()).to(self.device) + md_pt.eval() + + ext_coord, ext_atype, nlist_t, mapping_t = self._prepare_lower_inputs() + fparam = None + aparam = None + + ret_eager = md_pt.forward_lower( + ext_coord.requires_grad_(True), + ext_atype, + nlist_t, + mapping_t, + fparam=fparam, + aparam=aparam, + ) + + traced = md_pt.forward_lower_exportable( + ext_coord, + ext_atype, + nlist_t, + mapping_t, + fparam=fparam, + aparam=aparam, + ) + self.assertIsInstance(traced, torch.nn.Module) + + exported = torch.export.export( + traced, + (ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam), + strict=False, + ) + self.assertIsNotNone(exported) + + ret_traced = traced(ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam) + ret_exported = exported.module()( + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam + ) + + for key in ("atom_energy", "energy"): + np.testing.assert_allclose( + ret_eager[key].detach().cpu().numpy(), + ret_traced[key].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=f"traced vs eager: {key}", + ) + np.testing.assert_allclose( + ret_eager[key].detach().cpu().numpy(), + ret_exported[key].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=f"exported vs eager: {key}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt_expt/model/test_ener_model.py b/source/tests/pt_expt/model/test_ener_model.py index 588b0ec2a9..f4fd9106e8 100644 --- a/source/tests/pt_expt/model/test_ener_model.py +++ b/source/tests/pt_expt/model/test_ener_model.py @@ -184,7 +184,7 @@ def test_forward_lower_exportable(self) -> None: ) # --- eager reference with zero params --- - ret_eager_zero = md._forward_lower( + ret_eager_zero = md.forward_lower( ext_coord.requires_grad_(True), ext_atype, nlist_t, @@ -262,7 +262,7 @@ def test_forward_lower_exportable(self) -> None: dtype=torch.float64, device=self.device, ) - ret_eager_nz = md._forward_lower( + ret_eager_nz = md.forward_lower( ext_coord.requires_grad_(True), ext_atype, nlist_t, diff --git a/source/tests/pt_expt/model/test_polar_model.py b/source/tests/pt_expt/model/test_polar_model.py new file mode 100644 index 0000000000..acfa929db2 --- /dev/null +++ b/source/tests/pt_expt/model/test_polar_model.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np +import torch + +from deepmd.dpmodel.descriptor import DescrptSeA as DPDescrptSeA +from deepmd.dpmodel.fitting import PolarFitting as DPPolarFitting +from deepmd.dpmodel.model.polar_model import PolarModel as DPPolarModel +from deepmd.dpmodel.utils.nlist import ( + build_neighbor_list, + extend_coord_with_ghosts, +) +from deepmd.dpmodel.utils.region import ( + normalize_coord, +) +from deepmd.pt_expt.model import ( + PolarModel, +) +from deepmd.pt_expt.utils import ( + env, +) + +from ...seed import ( + GLOBAL_SEED, +) + + +class TestPolarModel(unittest.TestCase): + def setUp(self) -> None: + self.device = env.DEVICE + self.natoms = 5 + self.rcut = 4.0 + self.rcut_smth = 0.5 + self.sel = [8, 6] + self.nt = 2 + self.type_map = ["foo", "bar"] + + generator = torch.Generator(device=self.device).manual_seed(GLOBAL_SEED) + cell = torch.rand( + [3, 3], dtype=torch.float64, device=self.device, generator=generator + ) + cell = (cell + cell.T) + 5.0 * torch.eye(3, device=self.device) + self.cell = cell.unsqueeze(0) + coord = torch.rand( + [self.natoms, 3], + dtype=torch.float64, + device=self.device, + generator=generator, + ) + coord = torch.matmul(coord, cell) + self.coord = coord.unsqueeze(0).to(self.device) + self.atype = torch.tensor( + [[0, 0, 0, 1, 1]], dtype=torch.int64, device=self.device + ) + + def _make_dp_model(self): + ds = DPDescrptSeA(self.rcut, self.rcut_smth, self.sel) + ft = DPPolarFitting( + self.nt, + ds.get_dim_out(), + embedding_width=ds.get_dim_emb(), + seed=GLOBAL_SEED, + ) + return DPPolarModel(ds, ft, type_map=self.type_map) + + def _prepare_lower_inputs(self): + coord_np = self.coord.detach().cpu().numpy() + atype_np = self.atype.detach().cpu().numpy() + cell_np = self.cell.reshape(1, 9).detach().cpu().numpy() + coord_normalized = normalize_coord( + coord_np.reshape(1, self.natoms, 3), + cell_np.reshape(1, 3, 3), + ) + extended_coord, extended_atype, mapping = extend_coord_with_ghosts( + coord_normalized, atype_np, cell_np, self.rcut + ) + nlist = build_neighbor_list( + extended_coord, + extended_atype, + self.natoms, + self.rcut, + self.sel, + distinguish_types=True, + ) + extended_coord = extended_coord.reshape(1, -1, 3) + return ( + torch.tensor(extended_coord, dtype=torch.float64, device=self.device), + torch.tensor(extended_atype, dtype=torch.int64, device=self.device), + torch.tensor(nlist, dtype=torch.int64, device=self.device), + torch.tensor(mapping, dtype=torch.int64, device=self.device), + ) + + def test_dp_consistency(self) -> None: + md_dp = self._make_dp_model() + md_pt = PolarModel.deserialize(md_dp.serialize()).to(self.device) + md_pt.eval() + + coord_np = self.coord.detach().cpu().numpy() + atype_np = self.atype.detach().cpu().numpy() + cell_np = self.cell.reshape(1, 9).detach().cpu().numpy() + ret_dp = md_dp(coord_np.reshape(1, -1), atype_np, cell_np) + + coord = self.coord.clone().requires_grad_(True) + ret_pt = md_pt(coord, self.atype, self.cell.reshape(1, 9)) + + np.testing.assert_allclose( + ret_dp["global_polar"], + ret_pt["global_polar"].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + ret_dp["polar"], + ret_pt["polar"].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + ) + + def test_output_keys(self) -> None: + md_dp = self._make_dp_model() + md_pt = PolarModel.deserialize(md_dp.serialize()).to(self.device) + md_pt.eval() + coord = self.coord.clone().requires_grad_(True) + ret = md_pt(coord, self.atype, self.cell.reshape(1, 9)) + self.assertIn("polar", ret) + self.assertIn("global_polar", ret) + + def test_forward_lower_exportable(self) -> None: + md_dp = self._make_dp_model() + md_pt = PolarModel.deserialize(md_dp.serialize()).to(self.device) + md_pt.eval() + + ext_coord, ext_atype, nlist_t, mapping_t = self._prepare_lower_inputs() + fparam = None + aparam = None + + ret_eager = md_pt.forward_lower( + ext_coord.requires_grad_(True), + ext_atype, + nlist_t, + mapping_t, + fparam=fparam, + aparam=aparam, + ) + + traced = md_pt.forward_lower_exportable( + ext_coord, + ext_atype, + nlist_t, + mapping_t, + fparam=fparam, + aparam=aparam, + ) + self.assertIsInstance(traced, torch.nn.Module) + + exported = torch.export.export( + traced, + (ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam), + strict=False, + ) + self.assertIsNotNone(exported) + + ret_traced = traced(ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam) + ret_exported = exported.module()( + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam + ) + + for key in ("polar", "global_polar"): + np.testing.assert_allclose( + ret_eager[key].detach().cpu().numpy(), + ret_traced[key].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=f"traced vs eager: {key}", + ) + np.testing.assert_allclose( + ret_eager[key].detach().cpu().numpy(), + ret_exported[key].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=f"exported vs eager: {key}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt_expt/model/test_property_model.py b/source/tests/pt_expt/model/test_property_model.py new file mode 100644 index 0000000000..12b12afea1 --- /dev/null +++ b/source/tests/pt_expt/model/test_property_model.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np +import torch + +from deepmd.dpmodel.descriptor import DescrptSeA as DPDescrptSeA +from deepmd.dpmodel.fitting import PropertyFittingNet as DPPropertyFittingNet +from deepmd.dpmodel.model.property_model import PropertyModel as DPPropertyModel +from deepmd.dpmodel.utils.nlist import ( + build_neighbor_list, + extend_coord_with_ghosts, +) +from deepmd.dpmodel.utils.region import ( + normalize_coord, +) +from deepmd.pt_expt.model import ( + PropertyModel, +) +from deepmd.pt_expt.utils import ( + env, +) + +from ...seed import ( + GLOBAL_SEED, +) + + +class TestPropertyModel(unittest.TestCase): + def setUp(self) -> None: + self.device = env.DEVICE + self.natoms = 5 + self.rcut = 4.0 + self.rcut_smth = 0.5 + self.sel = [8, 6] + self.nt = 2 + self.type_map = ["foo", "bar"] + + generator = torch.Generator(device=self.device).manual_seed(GLOBAL_SEED) + cell = torch.rand( + [3, 3], dtype=torch.float64, device=self.device, generator=generator + ) + cell = (cell + cell.T) + 5.0 * torch.eye(3, device=self.device) + self.cell = cell.unsqueeze(0) + coord = torch.rand( + [self.natoms, 3], + dtype=torch.float64, + device=self.device, + generator=generator, + ) + coord = torch.matmul(coord, cell) + self.coord = coord.unsqueeze(0).to(self.device) + self.atype = torch.tensor( + [[0, 0, 0, 1, 1]], dtype=torch.int64, device=self.device + ) + + def _make_dp_model(self): + ds = DPDescrptSeA(self.rcut, self.rcut_smth, self.sel) + ft = DPPropertyFittingNet( + self.nt, + ds.get_dim_out(), + task_dim=3, + seed=GLOBAL_SEED, + ) + return DPPropertyModel(ds, ft, type_map=self.type_map) + + def _prepare_lower_inputs(self): + coord_np = self.coord.detach().cpu().numpy() + atype_np = self.atype.detach().cpu().numpy() + cell_np = self.cell.reshape(1, 9).detach().cpu().numpy() + coord_normalized = normalize_coord( + coord_np.reshape(1, self.natoms, 3), + cell_np.reshape(1, 3, 3), + ) + extended_coord, extended_atype, mapping = extend_coord_with_ghosts( + coord_normalized, atype_np, cell_np, self.rcut + ) + nlist = build_neighbor_list( + extended_coord, + extended_atype, + self.natoms, + self.rcut, + self.sel, + distinguish_types=True, + ) + extended_coord = extended_coord.reshape(1, -1, 3) + return ( + torch.tensor(extended_coord, dtype=torch.float64, device=self.device), + torch.tensor(extended_atype, dtype=torch.int64, device=self.device), + torch.tensor(nlist, dtype=torch.int64, device=self.device), + torch.tensor(mapping, dtype=torch.int64, device=self.device), + ) + + def test_dp_consistency(self) -> None: + md_dp = self._make_dp_model() + md_pt = PropertyModel.deserialize(md_dp.serialize()).to(self.device) + md_pt.eval() + + coord_np = self.coord.detach().cpu().numpy() + atype_np = self.atype.detach().cpu().numpy() + cell_np = self.cell.reshape(1, 9).detach().cpu().numpy() + ret_dp = md_dp(coord_np.reshape(1, -1), atype_np, cell_np) + + coord = self.coord.clone().requires_grad_(True) + ret_pt = md_pt(coord, self.atype, self.cell.reshape(1, 9)) + + var_name = md_pt.get_var_name() + np.testing.assert_allclose( + ret_dp[var_name], + ret_pt[var_name].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + ret_dp[f"atom_{var_name}"], + ret_pt[f"atom_{var_name}"].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + ) + + def test_output_keys(self) -> None: + md_dp = self._make_dp_model() + md_pt = PropertyModel.deserialize(md_dp.serialize()).to(self.device) + md_pt.eval() + coord = self.coord.clone().requires_grad_(True) + ret = md_pt(coord, self.atype, self.cell.reshape(1, 9)) + var_name = md_pt.get_var_name() + self.assertIn(var_name, ret) + self.assertIn(f"atom_{var_name}", ret) + + def test_forward_lower_exportable(self) -> None: + md_dp = self._make_dp_model() + md_pt = PropertyModel.deserialize(md_dp.serialize()).to(self.device) + md_pt.eval() + + ext_coord, ext_atype, nlist_t, mapping_t = self._prepare_lower_inputs() + fparam = None + aparam = None + + ret_eager = md_pt.forward_lower( + ext_coord.requires_grad_(True), + ext_atype, + nlist_t, + mapping_t, + fparam=fparam, + aparam=aparam, + ) + + traced = md_pt.forward_lower_exportable( + ext_coord, + ext_atype, + nlist_t, + mapping_t, + fparam=fparam, + aparam=aparam, + ) + self.assertIsInstance(traced, torch.nn.Module) + + exported = torch.export.export( + traced, + (ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam), + strict=False, + ) + self.assertIsNotNone(exported) + + ret_traced = traced(ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam) + ret_exported = exported.module()( + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam + ) + + var_name = md_pt.get_var_name() + for key in (f"atom_{var_name}", var_name): + np.testing.assert_allclose( + ret_eager[key].detach().cpu().numpy(), + ret_traced[key].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=f"traced vs eager: {key}", + ) + np.testing.assert_allclose( + ret_eager[key].detach().cpu().numpy(), + ret_exported[key].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=f"exported vs eager: {key}", + ) + + +if __name__ == "__main__": + unittest.main()