diff --git a/deepmd/tf/descriptor/se_a.py b/deepmd/tf/descriptor/se_a.py index 95330513ba..32be75d54e 100644 --- a/deepmd/tf/descriptor/se_a.py +++ b/deepmd/tf/descriptor/se_a.py @@ -77,6 +77,9 @@ from .se import ( DescrptSe, ) +from .stat import ( + load_or_compute_se_input_stats, +) @Descriptor.register("se_e2_a") @@ -374,7 +377,8 @@ def compute_input_stats( **kwargs Additional keyword arguments. """ - if True: + + def compute_stats() -> dict[str, Any]: sumr = [] suma = [] sumn = [] @@ -398,7 +402,16 @@ def compute_input_stats( "sumr2": sumr2, "suma2": suma2, } - self.merge_input_stats(stat_dict) + return stat_dict + + stat_dict = load_or_compute_se_input_stats( + self, + kwargs.get("stat_file_path"), + last_dim=4, + compute=compute_stats, + mixed_types=False, + ) + self.merge_input_stats(stat_dict) def merge_input_stats(self, stat_dict: dict[str, Any]) -> None: """Merge the statistics computed from compute_input_stats to obtain the self.davg and self.dstd. diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py index 0058763466..1bbb0a5595 100644 --- a/deepmd/tf/descriptor/se_atten.py +++ b/deepmd/tf/descriptor/se_atten.py @@ -90,6 +90,9 @@ from .se_a import ( DescrptSeA, ) +from .stat import ( + load_or_compute_se_input_stats, +) log = logging.getLogger(__name__) @@ -373,7 +376,8 @@ def compute_input_stats( **kwargs Additional keyword arguments. """ - if True: + + def compute_stats() -> dict[str, Any]: sumr = [] suma = [] sumn = [] @@ -418,7 +422,16 @@ def compute_input_stats( "sumr2": sumr2, "suma2": suma2, } - self.merge_input_stats(stat_dict) + return stat_dict + + stat_dict = load_or_compute_se_input_stats( + self, + kwargs.get("stat_file_path"), + last_dim=4, + compute=compute_stats, + mixed_types=True, + ) + self.merge_input_stats(stat_dict) def enable_compression( self, diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py index 3508da98a3..0acab2b94b 100644 --- a/deepmd/tf/descriptor/se_r.py +++ b/deepmd/tf/descriptor/se_r.py @@ -50,6 +50,9 @@ from .se import ( DescrptSe, ) +from .stat import ( + load_or_compute_se_input_stats, +) @Descriptor.register("se_e2_r") @@ -274,17 +277,27 @@ def compute_input_stats( **kwargs Additional keyword arguments. """ - sumr = [] - sumn = [] - sumr2 = [] - for cc, bb, tt, nn, mm in zip( - data_coord, data_box, data_atype, natoms_vec, mesh, strict=True - ): - sysr, sysr2, sysn = self._compute_dstats_sys_se_r(cc, bb, tt, nn, mm) - sumr.append(sysr) - sumn.append(sysn) - sumr2.append(sysr2) - stat_dict = {"sumr": sumr, "sumn": sumn, "sumr2": sumr2} + + def compute_stats() -> dict[str, Any]: + sumr = [] + sumn = [] + sumr2 = [] + for cc, bb, tt, nn, mm in zip( + data_coord, data_box, data_atype, natoms_vec, mesh, strict=True + ): + sysr, sysr2, sysn = self._compute_dstats_sys_se_r(cc, bb, tt, nn, mm) + sumr.append(sysr) + sumn.append(sysn) + sumr2.append(sysr2) + return {"sumr": sumr, "sumn": sumn, "sumr2": sumr2} + + stat_dict = load_or_compute_se_input_stats( + self, + kwargs.get("stat_file_path"), + last_dim=1, + compute=compute_stats, + mixed_types=False, + ) self.merge_input_stats(stat_dict) def merge_input_stats(self, stat_dict: dict[str, Any]) -> None: diff --git a/deepmd/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py index 0cf81b7c5f..16bec59bf0 100644 --- a/deepmd/tf/descriptor/se_t.py +++ b/deepmd/tf/descriptor/se_t.py @@ -52,6 +52,9 @@ from .se import ( DescrptSe, ) +from .stat import ( + load_or_compute_se_input_stats, +) @Descriptor.register("se_e3") @@ -257,7 +260,8 @@ def compute_input_stats( **kwargs Additional keyword arguments. """ - if True: + + def compute_stats() -> dict[str, Any]: sumr = [] suma = [] sumn = [] @@ -281,7 +285,16 @@ def compute_input_stats( "sumr2": sumr2, "suma2": suma2, } - self.merge_input_stats(stat_dict) + return stat_dict + + stat_dict = load_or_compute_se_input_stats( + self, + kwargs.get("stat_file_path"), + last_dim=4, + compute=compute_stats, + mixed_types=False, + ) + self.merge_input_stats(stat_dict) def merge_input_stats(self, stat_dict: dict[str, Any]) -> None: """Merge the statistics computed from compute_input_stats to obtain the self.davg and self.dstd. diff --git a/deepmd/tf/descriptor/stat.py b/deepmd/tf/descriptor/stat.py new file mode 100644 index 0000000000..bf91d5db92 --- /dev/null +++ b/deepmd/tf/descriptor/stat.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) +from typing import ( + Any, +) + +import numpy as np + +from deepmd.common import ( + get_hash, +) +from deepmd.utils.path import ( + DPPath, +) + + +def _descriptor_rcut_smth(descrpt: Any) -> float: + if hasattr(descrpt, "rcut_smth"): + return descrpt.rcut_smth + return descrpt.rcut_r_smth + + +def _descriptor_sel(descrpt: Any, last_dim: int) -> list[int]: + if hasattr(descrpt, "get_sel"): + sel = descrpt.get_sel() + elif last_dim == 1: + sel = descrpt.sel_r + else: + sel = descrpt.sel_a + if isinstance(sel, np.ndarray): + sel = sel.tolist() + elif isinstance(sel, int): + sel = [sel] + return [int(ii) for ii in sel] + + +def _descriptor_stat_path( + descrpt: Any, + stat_file_path: DPPath | None, + last_dim: int, + mixed_types: bool, +) -> DPPath | None: + if stat_file_path is None: + return None + sel = _descriptor_sel(descrpt, last_dim) + stat_hash = get_hash( + { + "type": "se_a" if last_dim == 4 else "se_r", + "ntypes": descrpt.get_ntypes(), + "rcut": round(descrpt.get_rcut(), 2), + "rcut_smth": round(_descriptor_rcut_smth(descrpt), 2), + "nsel": sum(sel), + "sel": sel, + "mixed_types": mixed_types, + } + ) + return stat_file_path / stat_hash + + +def _stat_keys(ntypes: int, angular: bool) -> list[str]: + keys = [f"r_{ii}" for ii in range(ntypes)] + if angular: + keys.extend(f"a_{ii}" for ii in range(ntypes)) + return keys + + +def _load_se_input_stats( + path: DPPath | None, + ntypes: int, + angular: bool, +) -> dict[str, list[list[float]]] | None: + if path is None or not path.is_dir(): + return None + if any(not (path / kk).is_file() for kk in _stat_keys(ntypes, angular)): + return None + + sumr = [] + sumn = [] + sumr2 = [] + suma = [] + suma2 = [] + for type_i in range(ntypes): + r_stat = (path / f"r_{type_i}").load_numpy() + sumn.append(float(r_stat[0])) + sumr.append(float(r_stat[1])) + sumr2.append(float(r_stat[2])) + if angular: + a_stat = (path / f"a_{type_i}").load_numpy() + suma.append(float(a_stat[1]) / 3.0) + suma2.append(float(a_stat[2]) / 3.0) + + ret = { + "sumr": [sumr], + "sumn": [sumn], + "sumr2": [sumr2], + } + if angular: + ret["suma"] = [suma] + ret["suma2"] = [suma2] + return ret + + +def _save_se_input_stats( + path: DPPath | None, + stat_dict: dict[str, Any], + ntypes: int, + angular: bool, +) -> None: + if path is None: + return + path.mkdir(parents=True, exist_ok=True) + + sumr = np.sum(stat_dict["sumr"], axis=0) + sumn = np.sum(stat_dict["sumn"], axis=0) + sumr2 = np.sum(stat_dict["sumr2"], axis=0) + if angular: + suma = np.sum(stat_dict["suma"], axis=0) + suma2 = np.sum(stat_dict["suma2"], axis=0) + + for type_i in range(ntypes): + (path / f"r_{type_i}").save_numpy( + np.array([sumn[type_i], sumr[type_i], sumr2[type_i]]) + ) + if angular: + (path / f"a_{type_i}").save_numpy( + np.array([3.0 * sumn[type_i], 3.0 * suma[type_i], 3.0 * suma2[type_i]]) + ) + + +def load_or_compute_se_input_stats( + descrpt: Any, + stat_file_path: DPPath | None, + last_dim: int, + compute: Callable[[], dict[str, Any]], + mixed_types: bool = False, +) -> dict[str, Any]: + """Load or compute SE descriptor input statistics using EnvMatStatSe format.""" + angular = last_dim == 4 + stat_path = _descriptor_stat_path(descrpt, stat_file_path, last_dim, mixed_types) + stat_dict = _load_se_input_stats(stat_path, descrpt.get_ntypes(), angular) + if stat_dict is not None: + return stat_dict + + stat_dict = compute() + _save_se_input_stats(stat_path, stat_dict, descrpt.get_ntypes(), angular) + return stat_dict diff --git a/deepmd/tf/entrypoints/train.py b/deepmd/tf/entrypoints/train.py index c0031a9def..d7b40ebfee 100755 --- a/deepmd/tf/entrypoints/train.py +++ b/deepmd/tf/entrypoints/train.py @@ -8,10 +8,14 @@ import json import logging import time +from pathlib import ( + Path, +) from typing import ( Any, ) +import h5py import numpy as np from deepmd.common import ( @@ -50,6 +54,9 @@ from deepmd.utils.data_system import ( get_data, ) +from deepmd.utils.path import ( + DPPath, +) __all__ = ["train"] @@ -232,6 +239,21 @@ def _do_work( # setup data modifier modifier = get_modifier(jdata["model"].get("modifier", None)) + # extract stat_file from training parameters + stat_file_path = None + if not is_compress: + stat_file_raw = jdata["training"].get("stat_file", None) + if stat_file_raw is not None and run_opt.is_chief: + stat_file_target = Path(stat_file_raw) + stat_file_target.parent.mkdir(parents=True, exist_ok=True) + if not stat_file_target.exists(): + if stat_file_raw.endswith((".h5", ".hdf5")): + with h5py.File(stat_file_raw, "w") as f: + pass + else: + stat_file_target.mkdir(parents=True, exist_ok=True) + stat_file_path = DPPath(stat_file_raw, "a") + # decouple the training data from the model compress process train_data = None valid_data = None @@ -289,7 +311,12 @@ def _do_work( origin_type_map = get_data( jdata["training"]["training_data"], rcut, None, modifier ).get_type_map() - model.build(train_data, stop_batch, origin_type_map=origin_type_map) + model.build( + train_data, + stop_batch, + origin_type_map=origin_type_map, + stat_file_path=stat_file_path, + ) if not is_compress: # train the model with the provided systems in a cyclic way diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py index ee81288197..166ed2e355 100644 --- a/deepmd/tf/fit/dos.py +++ b/deepmd/tf/fit/dos.py @@ -19,6 +19,13 @@ from deepmd.tf.fit.fitting import ( Fitting, ) +from deepmd.tf.fit.stat import ( + load_param_stats, + make_aparam_stats, + make_fparam_stats, + save_param_stats, + stats_avg_std, +) from deepmd.tf.loss.dos import ( DOSLoss, ) @@ -48,6 +55,9 @@ from deepmd.utils.out_stat import ( compute_stats_from_redu, ) +from deepmd.utils.path import ( + DPPath, +) if TYPE_CHECKING: from deepmd.tf.utils.learning_rate import ( @@ -265,7 +275,12 @@ def _compute_output_stats( return dos_shift - def compute_input_stats(self, all_stat: dict, protection: float = 1e-2) -> None: + def compute_input_stats( + self, + all_stat: dict, + protection: float = 1e-2, + stat_file_path: DPPath | None = None, + ) -> None: """Compute the input statistics. Parameters @@ -276,35 +291,24 @@ def compute_input_stats(self, all_stat: dict, protection: float = 1e-2) -> None: can be prepared by model.make_stat_input protection Divided-by-zero protection + stat_file_path + The path to the stat file. """ # stat fparam if self.numb_fparam > 0: - cat_data = np.concatenate(all_stat["fparam"], axis=0) - cat_data = np.reshape(cat_data, [-1, self.numb_fparam]) - self.fparam_avg = np.average(cat_data, axis=0) - self.fparam_std = np.std(cat_data, axis=0) - for ii in range(self.fparam_std.size): - if self.fparam_std[ii] < protection: - self.fparam_std[ii] = protection + fparam_stats = load_param_stats(stat_file_path, "fparam", self.numb_fparam) + if fparam_stats is None: + fparam_stats = make_fparam_stats(all_stat, self.numb_fparam) + save_param_stats(stat_file_path, "fparam", fparam_stats) + self.fparam_avg, self.fparam_std = stats_avg_std(fparam_stats, protection) self.fparam_inv_std = 1.0 / self.fparam_std # stat aparam if self.numb_aparam > 0: - sys_sumv = [] - sys_sumv2 = [] - sys_sumn = [] - for ss_ in all_stat["aparam"]: - ss = np.reshape(ss_, [-1, self.numb_aparam]) - sys_sumv.append(np.sum(ss, axis=0)) - sys_sumv2.append(np.sum(np.multiply(ss, ss), axis=0)) - sys_sumn.append(ss.shape[0]) - sumv = np.sum(sys_sumv, axis=0) - sumv2 = np.sum(sys_sumv2, axis=0) - sumn = np.sum(sys_sumn) - self.aparam_avg = (sumv) / sumn - self.aparam_std = self._compute_std(sumv2, sumv, sumn) - for ii in range(self.aparam_std.size): - if self.aparam_std[ii] < protection: - self.aparam_std[ii] = protection + aparam_stats = load_param_stats(stat_file_path, "aparam", self.numb_aparam) + if aparam_stats is None: + aparam_stats = make_aparam_stats(all_stat, self.numb_aparam) + save_param_stats(stat_file_path, "aparam", aparam_stats) + self.aparam_avg, self.aparam_std = stats_avg_std(aparam_stats, protection) self.aparam_inv_std = 1.0 / self.aparam_std def _compute_std( diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index 2b4027d464..f3a45430a7 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -22,6 +22,13 @@ from deepmd.tf.fit.fitting import ( Fitting, ) +from deepmd.tf.fit.stat import ( + load_param_stats, + make_aparam_stats, + make_fparam_stats, + save_param_stats, + stats_avg_std, +) from deepmd.tf.infer import ( DeepPotential, ) @@ -65,6 +72,9 @@ from deepmd.utils.out_stat import ( compute_stats_from_redu, ) +from deepmd.utils.path import ( + DPPath, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -334,7 +344,12 @@ def _compute_output_stats( ) return energy_shift.ravel() - def compute_input_stats(self, all_stat: dict, protection: float = 1e-2) -> None: + def compute_input_stats( + self, + all_stat: dict, + protection: float = 1e-2, + stat_file_path: DPPath | None = None, + ) -> None: """Compute the input statistics. Parameters @@ -345,35 +360,24 @@ def compute_input_stats(self, all_stat: dict, protection: float = 1e-2) -> None: can be prepared by model.make_stat_input protection Divided-by-zero protection + stat_file_path + The path to the stat file. """ # stat fparam if self.numb_fparam > 0: - cat_data = np.concatenate(all_stat["fparam"], axis=0) - cat_data = np.reshape(cat_data, [-1, self.numb_fparam]) - self.fparam_avg = np.average(cat_data, axis=0) - self.fparam_std = np.std(cat_data, axis=0) - for ii in range(self.fparam_std.size): - if self.fparam_std[ii] < protection: - self.fparam_std[ii] = protection + fparam_stats = load_param_stats(stat_file_path, "fparam", self.numb_fparam) + if fparam_stats is None: + fparam_stats = make_fparam_stats(all_stat, self.numb_fparam) + save_param_stats(stat_file_path, "fparam", fparam_stats) + self.fparam_avg, self.fparam_std = stats_avg_std(fparam_stats, protection) self.fparam_inv_std = 1.0 / self.fparam_std # stat aparam if self.numb_aparam > 0: - sys_sumv = [] - sys_sumv2 = [] - sys_sumn = [] - for ss_ in all_stat["aparam"]: - ss = np.reshape(ss_, [-1, self.numb_aparam]) - sys_sumv.append(np.sum(ss, axis=0)) - sys_sumv2.append(np.sum(np.multiply(ss, ss), axis=0)) - sys_sumn.append(ss.shape[0]) - sumv = np.sum(sys_sumv, axis=0) - sumv2 = np.sum(sys_sumv2, axis=0) - sumn = np.sum(sys_sumn) - self.aparam_avg = (sumv) / sumn - self.aparam_std = self._compute_std(sumv2, sumv, sumn) - for ii in range(self.aparam_std.size): - if self.aparam_std[ii] < protection: - self.aparam_std[ii] = protection + aparam_stats = load_param_stats(stat_file_path, "aparam", self.numb_aparam) + if aparam_stats is None: + aparam_stats = make_aparam_stats(all_stat, self.numb_aparam) + save_param_stats(stat_file_path, "aparam", aparam_stats) + self.aparam_avg, self.aparam_std = stats_avg_std(aparam_stats, protection) self.aparam_inv_std = 1.0 / self.aparam_std def _compute_std(self, sumv2: float, sumv: float, sumn: int) -> float: diff --git a/deepmd/tf/fit/stat.py b/deepmd/tf/fit/stat.py new file mode 100644 index 0000000000..f0605fcaf1 --- /dev/null +++ b/deepmd/tf/fit/stat.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import numpy as np + +from deepmd.utils.env_mat_stat import ( + StatItem, +) +from deepmd.utils.path import ( + DPPath, +) + + +def load_param_stats( + stat_file_path: DPPath | None, + name: str, + dim: int, +) -> list[StatItem] | None: + if ( + stat_file_path is None + or not stat_file_path.is_dir() + or not (stat_file_path / name).is_file() + ): + return None + arr = (stat_file_path / name).load_numpy() + if arr.shape != (dim, 3): + raise ValueError(f"Invalid {name} stat shape {arr.shape}; expected ({dim}, 3).") + return [ + StatItem(number=arr[ii, 0], sum=arr[ii, 1], squared_sum=arr[ii, 2]) + for ii in range(dim) + ] + + +def save_param_stats( + stat_file_path: DPPath | None, + name: str, + stats: list[StatItem], +) -> None: + if stat_file_path is None: + return + stat_file_path.mkdir(parents=True, exist_ok=True) + arr = np.array([[ss.number, ss.sum, ss.squared_sum] for ss in stats]) + (stat_file_path / name).save_numpy(arr) + + +def make_fparam_stats(all_stat: dict[str, Any], dim: int) -> list[StatItem]: + cat_data = np.concatenate(all_stat["fparam"], axis=0) + cat_data = np.reshape(cat_data, [-1, dim]) + sumv = np.sum(cat_data, axis=0) + sumv2 = np.sum(cat_data * cat_data, axis=0) + sumn = cat_data.shape[0] + return [ + StatItem(number=sumn, sum=sumv[ii], squared_sum=sumv2[ii]) for ii in range(dim) + ] + + +def make_aparam_stats(all_stat: dict[str, Any], dim: int) -> list[StatItem]: + sys_sumv = [] + sys_sumv2 = [] + sys_sumn = [] + for ss_ in all_stat["aparam"]: + ss = np.reshape(ss_, [-1, dim]) + 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(sys_sumv, axis=0) + sumv2 = np.sum(sys_sumv2, axis=0) + sumn = np.sum(sys_sumn) + return [ + StatItem(number=sumn, sum=sumv[ii], squared_sum=sumv2[ii]) for ii in range(dim) + ] + + +def stats_avg_std( + stats: list[StatItem], + protection: float, +) -> tuple[np.ndarray, np.ndarray]: + avg = np.array([ss.compute_avg() for ss in stats]) + std = np.array([ss.compute_std(protection=protection) for ss in stats]) + return avg, std diff --git a/deepmd/tf/model/dos.py b/deepmd/tf/model/dos.py index dc3c54c42a..5d53abecf9 100644 --- a/deepmd/tf/model/dos.py +++ b/deepmd/tf/model/dos.py @@ -14,6 +14,9 @@ from deepmd.utils.data_system import ( DeepmdDataSystem, ) +from deepmd.utils.path import ( + DPPath, +) from .model import ( StandardModel, @@ -22,6 +25,9 @@ make_stat_input, merge_sys_stat, ) +from .stat_file import ( + add_type_map_to_stat_path, +) @StandardModel.register("dos") @@ -92,17 +98,27 @@ def get_numb_aparam(self) -> int: """Get the number of atomic parameters.""" return self.numb_aparam - def data_stat(self, data: DeepmdDataSystem) -> None: + def data_stat( + self, data: DeepmdDataSystem, stat_file_path: DPPath | None = None + ) -> None: all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys=False) m_all_stat = merge_sys_stat(all_stat) + stat_file_path = add_type_map_to_stat_path(stat_file_path, self.type_map) self._compute_input_stat( - m_all_stat, protection=self.data_stat_protect, mixed_type=data.mixed_type + m_all_stat, + protection=self.data_stat_protect, + mixed_type=data.mixed_type, + stat_file_path=stat_file_path, ) # self._compute_output_stat(all_stat, mixed_type=data.mixed_type) # self.bias_atom_e = data.compute_energy_shift(self.rcond) def _compute_input_stat( - self, all_stat: dict, protection: float = 1e-2, mixed_type: bool = False + self, + all_stat: dict, + protection: float = 1e-2, + mixed_type: bool = False, + stat_file_path: DPPath | None = None, ) -> None: if mixed_type: self.descrpt.compute_input_stats( @@ -114,6 +130,7 @@ def _compute_input_stat( all_stat, mixed_type, all_stat["real_natoms_vec"], + stat_file_path=stat_file_path, ) else: self.descrpt.compute_input_stats( @@ -123,8 +140,11 @@ def _compute_input_stat( all_stat["natoms_vec"], all_stat["default_mesh"], all_stat, + stat_file_path=stat_file_path, ) - self.fitting.compute_input_stats(all_stat, protection=protection) + self.fitting.compute_input_stats( + all_stat, protection=protection, stat_file_path=stat_file_path + ) def _compute_output_stat(self, all_stat: dict, mixed_type: bool = False) -> None: if mixed_type: diff --git a/deepmd/tf/model/ener.py b/deepmd/tf/model/ener.py index ab664b9404..b42be718bc 100644 --- a/deepmd/tf/model/ener.py +++ b/deepmd/tf/model/ener.py @@ -5,6 +5,15 @@ import numpy as np +from deepmd.dpmodel.utils.batch import ( + normalize_batch, +) +from deepmd.dpmodel.utils.stat import ( + _restore_observed_type_from_file, + _save_observed_type_to_file, + collect_observed_types, + compute_output_stats, +) from deepmd.tf.env import ( MODEL_VERSION, global_cvt_2_ener_float, @@ -23,6 +32,9 @@ from deepmd.tf.utils.type_embed import ( TypeEmbedNet, ) +from deepmd.utils.path import ( + DPPath, +) from .model import ( StandardModel, @@ -31,6 +43,45 @@ make_stat_input, merge_sys_stat, ) +from .stat_file import ( + add_type_map_to_stat_path, +) + + +def _pack_stat_batches(all_stat: dict) -> list[dict[str, Any]]: + """Pack TensorFlow statistics batches into backend-agnostic samples.""" + first_key = next(iter(all_stat.keys())) + nsystems = len(all_stat[first_key]) + sampled = [] + for sys_idx in range(nsystems): + merged = {} + for key, values in all_stat.items(): + sys_values = values[sys_idx] + if isinstance(sys_values[0], np.ndarray): + if sys_values[0].ndim >= 2: + merged[key] = np.concatenate(sys_values, axis=0) + else: + # 1-D arrays such as natoms_vec are per-system constants. + merged[key] = sys_values[0] + else: + # Scalar flags such as find_*. + merged[key] = sys_values[0] + sampled.append(normalize_batch(merged)) + return sampled + + +def _save_observed_types_to_file( + stat_file_path: DPPath | None, + sampled: list[dict[str, Any]], + type_map: list[str] | None, +) -> None: + """Save observed atom types using the backend-agnostic dpmodel helpers.""" + if stat_file_path is None or type_map is None: + return + observed = _restore_observed_type_from_file(stat_file_path) + if observed is None: + observed = collect_observed_types(sampled, type_map) + _save_observed_type_to_file(stat_file_path, observed) @StandardModel.register("ener") @@ -134,17 +185,29 @@ def get_numb_aparam(self) -> int: """Get the number of atomic parameters.""" return self.numb_aparam - def data_stat(self, data: DeepmdDataSystem) -> None: + def data_stat( + self, data: DeepmdDataSystem, stat_file_path: DPPath | None = None + ) -> None: all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys=False) m_all_stat = merge_sys_stat(all_stat) + stat_file_path = add_type_map_to_stat_path(stat_file_path, self.type_map) self._compute_input_stat( - m_all_stat, protection=self.data_stat_protect, mixed_type=data.mixed_type + m_all_stat, + protection=self.data_stat_protect, + mixed_type=data.mixed_type, + stat_file_path=stat_file_path, + ) + self._compute_output_stat( + all_stat, mixed_type=data.mixed_type, stat_file_path=stat_file_path ) - self._compute_output_stat(all_stat, mixed_type=data.mixed_type) # self.bias_atom_e = data.compute_energy_shift(self.rcond) def _compute_input_stat( - self, all_stat: dict, protection: float = 1e-2, mixed_type: bool = False + self, + all_stat: dict, + protection: float = 1e-2, + mixed_type: bool = False, + stat_file_path: DPPath | None = None, ) -> None: if mixed_type: self.descrpt.compute_input_stats( @@ -156,6 +219,7 @@ def _compute_input_stat( all_stat, mixed_type, all_stat["real_natoms_vec"], + stat_file_path=stat_file_path, ) else: self.descrpt.compute_input_stats( @@ -165,14 +229,41 @@ def _compute_input_stat( all_stat["natoms_vec"], all_stat["default_mesh"], all_stat, + stat_file_path=stat_file_path, ) - self.fitting.compute_input_stats(all_stat, protection=protection) + self.fitting.compute_input_stats( + all_stat, protection=protection, stat_file_path=stat_file_path + ) - def _compute_output_stat(self, all_stat: dict, mixed_type: bool = False) -> None: - if mixed_type: - self.fitting.compute_output_stats(all_stat, mixed_type=mixed_type) - else: - self.fitting.compute_output_stats(all_stat) + def _compute_output_stat( + self, + all_stat: dict, + mixed_type: bool = False, + stat_file_path: DPPath | None = None, + ) -> None: + # Reuse the backend-agnostic dpmodel stat implementation instead of + # maintaining a TensorFlow copy of the same save/load/stat logic. This + # intentionally uses the dpmodel/PT per-frame energy-bias regression even + # when no stat file is provided, so TF computes the same initial bias it + # would later restore from a cross-backend stat file. + sampled = _pack_stat_batches(all_stat) + _save_observed_types_to_file(stat_file_path, sampled, self.type_map) + preset_bias = None + if len(self.fitting.atom_ener) > 0: + preset_bias = {"energy": self.fitting.atom_ener_v} + bias_out, _ = compute_output_stats( + sampled, + self.ntypes, + keys=["energy"], + stat_file_path=stat_file_path, + rcond=getattr(self.fitting, "rcond", None), + preset_bias=preset_bias, + ) + + if "energy" in bias_out: + # TensorFlow fitting code historically stores a 1-D bias vector, + # while stat files use the PyTorch-compatible (ntypes, 1) shape. + self.fitting.bias_atom_e = bias_out["energy"].ravel() def build( self, diff --git a/deepmd/tf/model/frozen.py b/deepmd/tf/model/frozen.py index 045a73b076..48fbb802d9 100644 --- a/deepmd/tf/model/frozen.py +++ b/deepmd/tf/model/frozen.py @@ -43,6 +43,9 @@ from deepmd.utils.data_system import ( DeepmdDataSystem, ) +from deepmd.utils.path import ( + DPPath, +) from .model import ( Model, @@ -202,7 +205,9 @@ def get_rcut(self) -> float: def get_ntypes(self) -> int: return self.model.get_ntypes() - def data_stat(self, data: DeepmdDataSystem) -> None: + def data_stat( + self, data: DeepmdDataSystem, stat_file_path: DPPath | None = None + ) -> None: pass def init_variables( diff --git a/deepmd/tf/model/linear.py b/deepmd/tf/model/linear.py index a9545cab5d..90f559f73f 100644 --- a/deepmd/tf/model/linear.py +++ b/deepmd/tf/model/linear.py @@ -31,6 +31,9 @@ from deepmd.utils.data_system import ( DeepmdDataSystem, ) +from deepmd.utils.path import ( + DPPath, +) from .model import ( Model, @@ -92,9 +95,14 @@ def get_ntypes(self) -> int: raise ValueError("Models have different ntypes") return self.models[0].get_ntypes() - def data_stat(self, data: DeepmdDataSystem) -> None: - for model in self.models: - model.data_stat(data) + def data_stat( + self, data: DeepmdDataSystem, stat_file_path: DPPath | None = None + ) -> None: + for ii, model in enumerate(self.models): + model_stat_path = ( + None if stat_file_path is None else stat_file_path / f"model{ii}" + ) + model.data_stat(data, stat_file_path=model_stat_path) def init_variables( self, diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index ccd541299b..b793f9145e 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -76,6 +76,9 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.path import ( + DPPath, +) from deepmd.utils.plugin import ( make_plugin_registry, ) @@ -473,7 +476,7 @@ def get_ntypes(self) -> int: """Get the number of types.""" @abstractmethod - def data_stat(self, data: dict) -> None: + def data_stat(self, data: dict, stat_file_path: DPPath | None = None) -> None: """Data staticis.""" def get_feed_dict( diff --git a/deepmd/tf/model/pairwise_dprc.py b/deepmd/tf/model/pairwise_dprc.py index 91ae67446c..7c68ba68ad 100644 --- a/deepmd/tf/model/pairwise_dprc.py +++ b/deepmd/tf/model/pairwise_dprc.py @@ -39,6 +39,9 @@ from deepmd.utils.data_system import ( DeepmdDataSystem, ) +from deepmd.utils.path import ( + DPPath, +) @Model.register("pairwise_dprc") @@ -319,9 +322,11 @@ def get_rcut(self) -> float: def get_ntypes(self) -> int: return self.ntypes - def data_stat(self, data: dict) -> None: - self.qm_model.data_stat(data) - self.qmmm_model.data_stat(data) + def data_stat(self, data: dict, stat_file_path: DPPath | None = None) -> None: + qm_stat_path = None if stat_file_path is None else stat_file_path / "qm" + qmmm_stat_path = None if stat_file_path is None else stat_file_path / "qmmm" + self.qm_model.data_stat(data, stat_file_path=qm_stat_path) + self.qmmm_model.data_stat(data, stat_file_path=qmmm_stat_path) def init_variables( self, diff --git a/deepmd/tf/model/stat_file.py b/deepmd/tf/model/stat_file.py new file mode 100644 index 0000000000..2c39caa62b --- /dev/null +++ b/deepmd/tf/model/stat_file.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.utils.path import ( + DPPath, +) + + +def add_type_map_to_stat_path( + stat_file_path: DPPath | None, + type_map: list[str] | None, +) -> DPPath | None: + if stat_file_path is not None and type_map is not None: + return stat_file_path / " ".join(type_map) + return stat_file_path diff --git a/deepmd/tf/model/tensor.py b/deepmd/tf/model/tensor.py index d8c8994cd8..bcde66c51f 100644 --- a/deepmd/tf/model/tensor.py +++ b/deepmd/tf/model/tensor.py @@ -20,6 +20,9 @@ from deepmd.utils.data_system import ( DeepmdDataSystem, ) +from deepmd.utils.path import ( + DPPath, +) from .model import ( StandardModel, @@ -28,6 +31,9 @@ make_stat_input, merge_sys_stat, ) +from .stat_file import ( + add_type_map_to_stat_path, +) class TensorModel(StandardModel): @@ -90,13 +96,25 @@ def get_sel_type(self) -> list[int]: def get_out_size(self) -> int: return self.fitting.get_out_size() - def data_stat(self, data: DeepmdDataSystem) -> None: + def data_stat( + self, data: DeepmdDataSystem, stat_file_path: DPPath | None = None + ) -> None: all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys=False) m_all_stat = merge_sys_stat(all_stat) - self._compute_input_stat(m_all_stat, protection=self.data_stat_protect) + stat_file_path = add_type_map_to_stat_path(stat_file_path, self.type_map) + self._compute_input_stat( + m_all_stat, + protection=self.data_stat_protect, + stat_file_path=stat_file_path, + ) self._compute_output_stat(m_all_stat) - def _compute_input_stat(self, all_stat: dict, protection: float = 1e-2) -> None: + def _compute_input_stat( + self, + all_stat: dict, + protection: float = 1e-2, + stat_file_path: DPPath | None = None, + ) -> None: self.descrpt.compute_input_stats( all_stat["coord"], all_stat["box"], @@ -104,6 +122,7 @@ def _compute_input_stat(self, all_stat: dict, protection: float = 1e-2) -> None: all_stat["natoms_vec"], all_stat["default_mesh"], all_stat, + stat_file_path=stat_file_path, ) if hasattr(self.fitting, "compute_input_stats"): self.fitting.compute_input_stats(all_stat, protection=protection) diff --git a/deepmd/tf/train/trainer.py b/deepmd/tf/train/trainer.py index 2b2b767c6f..e1d7deb04b 100644 --- a/deepmd/tf/train/trainer.py +++ b/deepmd/tf/train/trainer.py @@ -68,6 +68,9 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.path import ( + DPPath, +) if TYPE_CHECKING: from collections.abc import ( @@ -209,6 +212,7 @@ def build( stop_batch: int = 0, origin_type_map: list[str] | None = None, suffix: str = "", + stat_file_path: DPPath | None = None, ) -> None: self.ntypes = self.model.get_ntypes() self.stop_batch = stop_batch @@ -248,7 +252,7 @@ def build( # self.saver.restore (in self._init_session) will restore avg and std variables, so data_stat is useless # init_from_frz_model will restore data_stat variables in `init_variables` method log.info("data stating... (this step may take long time)") - self.model.data_stat(data) + self.model.data_stat(data, stat_file_path=stat_file_path) # config the init_frz_model command if self.run_opt.init_mode == "init_from_frz_model": diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 19fbe8cebd..5b7ee37818 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5042,9 +5042,7 @@ def training_args( data_args = [ arg_training_data, arg_validation_data, - Argument( - "stat_file", str, optional=True, doc=doc_only_pt_supported + doc_stat_file - ), + Argument("stat_file", str, optional=True, doc=doc_stat_file), ] args = ( data_args diff --git a/source/tests/consistent/test_stat_file.py b/source/tests/consistent/test_stat_file.py new file mode 100644 index 0000000000..8380b3ad96 --- /dev/null +++ b/source/tests/consistent/test_stat_file.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Test consistency of stat file generation between TensorFlow and PyTorch backends.""" + +import json +import os +import subprocess +import sys +import tempfile +import unittest +from copy import ( + deepcopy, +) +from pathlib import ( + Path, +) + +import numpy as np + +from .common import ( + INSTALLED_PT, + INSTALLED_TF, +) + + +class TestStatFileConsistency(unittest.TestCase): + """Test that TensorFlow and PyTorch produce identical stat files.""" + + def setUp(self) -> None: + """Set up test data and configuration.""" + # Use a minimal but realistic configuration + self.config_base = { + "model": { + "type_map": ["O", "H"], + "data_stat_nbatch": 80, # Cover the whole test set for deterministic stats + "descriptor": { + "type": "se_e2_a", + "sel": [2, 4], + "rcut_smth": 0.50, + "rcut": 1.00, + "neuron": [4, 8], + "resnet_dt": False, + "axis_neuron": 4, + "seed": 42, + }, + "fitting_net": { + "neuron": [8, 8], + "resnet_dt": True, + "seed": 42, + }, + }, + "learning_rate": { + "type": "exp", + "decay_steps": 10, + "start_lr": 0.001, + "stop_lr": 1e-8, + }, + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + }, + "training": { + "training_data": { + "systems": [], # Will be filled with test data + "batch_size": 1, + }, + "seed": 42, + "numb_steps": 1, # Minimal training to just generate stat files + "disp_freq": 1, + "save_freq": 1, + }, + } + + # Find the test data directory + examples_path = Path(__file__).parent.parent.parent.parent / "examples" + self.test_data_path = examples_path / "water" / "data" / "data_0" + self.unequal_frame_data_paths = [ + examples_path / "water" / "data" / "data_0", + examples_path / "water" / "data" / "data_1", + ] + + # Skip if test data not available + if not self.test_data_path.exists(): + self.skipTest("Test data not available") + if any(not path.exists() for path in self.unequal_frame_data_paths): + self.skipTest("Unequal-frame test data not available") + + def _run_training_with_stat_file( + self, backend: str, config: dict, temp_dir: str, stat_dir: str + ) -> None: + """Run training with specified backend to generate stat files. + + Parameters + ---------- + backend : str + Backend to use ('tf' or 'pt') + config : dict + Training configuration + temp_dir : str + Temporary directory for output + stat_dir : str + Directory for stat files + """ + config_copy = deepcopy(config) + config_copy["training"]["stat_file"] = stat_dir + if not config_copy["training"]["training_data"]["systems"]: + config_copy["training"]["training_data"]["systems"] = [ + str(self.test_data_path) + ] + + config_file = os.path.join(temp_dir, f"input_{backend}.json") + + with open(config_file, "w") as f: + json.dump(config_copy, f, indent=2) + + # Run training with specified backend using subprocess + env = os.environ.copy() + dp_cmd = Path(sys.executable).with_name("dp") + base_cmd = ( + [str(dp_cmd)] + if dp_cmd.exists() + else [sys.executable, "-c", "from deepmd.main import main; main()"] + ) + cmd = [*base_cmd, "train", config_file] + if backend == "pt": + cmd = [*base_cmd, "--pt", "train", config_file] + + cmd.extend(["--log-level", "WARNING"]) + + result = subprocess.run( + cmd, + cwd=temp_dir, + capture_output=True, + text=True, + env=env, + timeout=120, + ) + + if result.returncode != 0: + self.fail( + f"Training failed for {backend} backend:\n" + f"stdout: {result.stdout}\n" + f"stderr: {result.stderr}" + ) + + def _compare_stat_directories( + self, + tf_stat_dir: str, + pt_stat_dir: str, + selected_names: set[str] | None = None, + rtol: float = 1e-10, + atol: float = 1e-12, + ) -> None: + """Compare stat file directories between TensorFlow and PyTorch. + + Parameters + ---------- + tf_stat_dir : str + TensorFlow stat file directory + pt_stat_dir : str + PyTorch stat file directory + selected_names : set[str], optional + Basenames of stat files to compare. When omitted, compare every file. + rtol : float + Relative tolerance for numeric stat file comparisons. + atol : float + Absolute tolerance for numeric stat file comparisons. + """ + tf_path = Path(tf_stat_dir) + pt_path = Path(pt_stat_dir) + + # Both directories should exist + self.assertTrue(tf_path.exists(), "TensorFlow stat directory should exist") + self.assertTrue(pt_path.exists(), "PyTorch stat directory should exist") + + # Both should be directories + self.assertTrue(tf_path.is_dir(), "TensorFlow stat path should be a directory") + self.assertTrue(pt_path.is_dir(), "PyTorch stat path should be a directory") + + tf_files = sorted( + ff.relative_to(tf_path) for ff in tf_path.rglob("*") if ff.is_file() + ) + pt_files = sorted( + ff.relative_to(pt_path) for ff in pt_path.rglob("*") if ff.is_file() + ) + if selected_names is not None: + tf_files = [ff for ff in tf_files if ff.name in selected_names] + pt_files = [ff for ff in pt_files if ff.name in selected_names] + self.assertEqual( + {ff.name for ff in tf_files}, + selected_names, + "TensorFlow should create the selected stat files", + ) + self.assertEqual( + {ff.name for ff in pt_files}, + selected_names, + "PyTorch should create the selected stat files", + ) + + self.assertEqual(tf_files, pt_files, "Both backends should create same files") + if selected_names is None: + self.assertTrue( + any(len(ff.parts) > 2 for ff in tf_files), + "Descriptor stat files should be saved under their hash directory", + ) + + for filename in tf_files: + tf_file = tf_path / filename + pt_file = pt_path / filename + + tf_data = np.load(tf_file) + pt_data = np.load(pt_file) + + self.assertEqual( + tf_data.shape, + pt_data.shape, + f"Shape mismatch in {filename}", + ) + + if np.issubdtype(tf_data.dtype, np.number): + np.testing.assert_allclose( + tf_data, + pt_data, + rtol=rtol, + atol=atol, + err_msg=f"Values differ in {filename}", + ) + else: + np.testing.assert_array_equal( + tf_data, + pt_data, + err_msg=f"Values differ in {filename}", + ) + + @unittest.skipUnless( + INSTALLED_TF and INSTALLED_PT, "TensorFlow and PyTorch required" + ) + def test_stat_file_consistency_basic(self) -> None: + """Test basic stat file consistency between TensorFlow and PyTorch backends.""" + with tempfile.TemporaryDirectory() as temp_dir: + tf_stat_dir = os.path.join(temp_dir, "tf_stat") + pt_stat_dir = os.path.join(temp_dir, "pt_stat") + + # Run TensorFlow training + self._run_training_with_stat_file( + "tf", self.config_base, temp_dir, tf_stat_dir + ) + + # Run PyTorch training + self._run_training_with_stat_file( + "pt", self.config_base, temp_dir, pt_stat_dir + ) + + # Compare the generated stat files with tight fp64 tolerances. + self._compare_stat_directories(tf_stat_dir, pt_stat_dir) + + @unittest.skipUnless( + INSTALLED_TF and INSTALLED_PT, "TensorFlow and PyTorch required" + ) + def test_output_stat_file_consistency_unequal_frame_systems(self) -> None: + """Test TF/PT output-stat consistency with unequal frame counts.""" + config = deepcopy(self.config_base) + config["training"]["training_data"]["systems"] = [ + str(path) for path in self.unequal_frame_data_paths + ] + config["training"]["training_data"]["batch_size"] = [1, 2] + + with tempfile.TemporaryDirectory() as temp_dir: + tf_stat_dir = os.path.join(temp_dir, "tf_stat") + pt_stat_dir = os.path.join(temp_dir, "pt_stat") + + # This case catches the per-system vs per-frame output-bias regression + # distinction: the legacy TF path weighted each system equally, whereas + # the shared stat implementation used by stat files weights frames + # consistently with the PyTorch backend. The per-system batch sizes + # make TF and PT collect the same 80-frame and 160-frame samples. + # Compare only the shared output-stat file that determines the + # restored energy bias. + self._run_training_with_stat_file("tf", config, temp_dir, tf_stat_dir) + self._run_training_with_stat_file("pt", config, temp_dir, pt_stat_dir) + + self._compare_stat_directories( + tf_stat_dir, + pt_stat_dir, + selected_names={"bias_atom_energy"}, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/tf/test_fitting_stat.py b/source/tests/tf/test_fitting_stat.py index eb28ef2bfa..b44031f803 100644 --- a/source/tests/tf/test_fitting_stat.py +++ b/source/tests/tf/test_fitting_stat.py @@ -3,6 +3,15 @@ from collections import ( defaultdict, ) +from pathlib import ( + Path, +) +from tempfile import ( + TemporaryDirectory, +) +from typing import ( + NoReturn, +) import numpy as np @@ -12,6 +21,9 @@ from deepmd.tf.fit import ( EnerFitting, ) +from deepmd.utils.path import ( + DPPath, +) from .common import ( j_loader, @@ -101,3 +113,52 @@ def test(self) -> None: np.testing.assert_almost_equal(frefs, fitting.fparam_std) np.testing.assert_almost_equal(arefa, fitting.aparam_avg) np.testing.assert_almost_equal(arefs, fitting.aparam_std) + + def test_stat_file(self) -> None: + descrpt = DescrptSeA(6.0, 5.8, [46, 92], neuron=[25, 50, 100], axis_neuron=16) + avgs = [0, 10] + stds = [2, 0.4] + sys_natoms = [10, 100] + sys_nframes = [5, 2] + all_data = _make_fake_data(sys_natoms, sys_nframes, avgs, stds) + frefa, frefs = _brute_fparam(all_data, len(avgs)) + arefa, arefs = _brute_aparam(all_data, len(avgs)) + + with TemporaryDirectory() as tempdir: + stat_path = DPPath(str(Path(tempdir)), "a") + fitting = EnerFitting( + descrpt.get_ntypes(), + descrpt.get_dim_out(), + neuron=[240, 240, 240], + resnet_dt=True, + numb_fparam=2, + numb_aparam=2, + ) + fitting.compute_input_stats( + all_data, protection=1e-2, stat_file_path=stat_path + ) + self.assertTrue((stat_path / "fparam").is_file()) + self.assertTrue((stat_path / "aparam").is_file()) + np.testing.assert_almost_equal(frefa, fitting.fparam_avg) + np.testing.assert_almost_equal(frefs, fitting.fparam_std) + np.testing.assert_almost_equal(arefa, fitting.aparam_avg) + np.testing.assert_almost_equal(arefs, fitting.aparam_std) + + def raise_error() -> NoReturn: + raise RuntimeError + + fitting = EnerFitting( + descrpt.get_ntypes(), + descrpt.get_dim_out(), + neuron=[240, 240, 240], + resnet_dt=True, + numb_fparam=2, + numb_aparam=2, + ) + fitting.compute_input_stats( + raise_error, protection=1e-2, stat_file_path=stat_path + ) + np.testing.assert_almost_equal(frefa, fitting.fparam_avg) + np.testing.assert_almost_equal(frefs, fitting.fparam_std) + np.testing.assert_almost_equal(arefa, fitting.aparam_avg) + np.testing.assert_almost_equal(arefs, fitting.aparam_std) diff --git a/source/tests/tf/test_gen_stat_data.py b/source/tests/tf/test_gen_stat_data.py index a49fe72f11..9acbe811a9 100644 --- a/source/tests/tf/test_gen_stat_data.py +++ b/source/tests/tf/test_gen_stat_data.py @@ -11,6 +11,9 @@ from deepmd.tf.fit import ( EnerFitting, ) +from deepmd.tf.model import ( + EnerModel, +) from deepmd.tf.model.model_stat import ( _make_all_stat_ref, make_stat_input, @@ -155,3 +158,58 @@ def test_ener_shift_assigned(self) -> None: tot0 = np.dot(data.compute_energy_shift(rcond=1), natoms) tot1 = np.dot(ener_shift1, natoms) np.testing.assert_almost_equal(tot0, tot1) + + def test_model_output_stat_matches_fitting(self) -> None: + dp_random.seed(0) + data = DeepmdDataSystem(["system_0", "system_1"], 5, 10, 1.0) + data.add("energy", 1, must=True) + all_stat = make_stat_input(data, 6, merge_sys=False) + descrpt = DescrptSeA( + 6.0, 5.8, [46, 92, 92], neuron=[25, 50, 100], axis_neuron=16 + ) + fitting = EnerFitting( + descrpt.get_ntypes(), + descrpt.get_dim_out(), + neuron=[240, 240, 240], + resnet_dt=True, + ) + ener_shift = fitting._compute_output_stats(all_stat, rcond=fitting.rcond) + + fitting = EnerFitting( + descrpt.get_ntypes(), + descrpt.get_dim_out(), + neuron=[240, 240, 240], + resnet_dt=True, + ) + model = EnerModel(descrpt, fitting) + model._compute_output_stat(all_stat) + np.testing.assert_almost_equal(model.fitting.bias_atom_e, ener_shift) + + def test_model_output_stat_assigned_matches_fitting(self) -> None: + dp_random.seed(0) + ae0 = dp_random.random() + data = DeepmdDataSystem(["system_0"], 5, 10, 1.0) + data.add("energy", 1, must=True) + all_stat = make_stat_input(data, 6, merge_sys=False) + descrpt = DescrptSeA( + 6.0, 5.8, [46, 92, 92], neuron=[25, 50, 100], axis_neuron=16 + ) + fitting = EnerFitting( + descrpt.get_ntypes(), + descrpt.get_dim_out(), + neuron=[240, 240, 240], + resnet_dt=True, + atom_ener=[ae0, None, None], + ) + ener_shift = fitting._compute_output_stats(all_stat, rcond=fitting.rcond) + + fitting = EnerFitting( + descrpt.get_ntypes(), + descrpt.get_dim_out(), + neuron=[240, 240, 240], + resnet_dt=True, + atom_ener=[ae0, None, None], + ) + model = EnerModel(descrpt, fitting) + model._compute_output_stat(all_stat) + np.testing.assert_almost_equal(model.fitting.bias_atom_e, ener_shift) diff --git a/source/tests/tf/test_model_se_a.py b/source/tests/tf/test_model_se_a.py index da20291087..8147918dcf 100644 --- a/source/tests/tf/test_model_se_a.py +++ b/source/tests/tf/test_model_se_a.py @@ -1,4 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from pathlib import ( + Path, +) +from tempfile import ( + TemporaryDirectory, +) + import dpdata import numpy as np @@ -17,6 +24,9 @@ from deepmd.tf.utils.type_embed import ( TypeEmbedNet, ) +from deepmd.utils.path import ( + DPPath, +) from .common import ( DataSystem, @@ -37,6 +47,41 @@ def setUp(self) -> None: def tearDown(self) -> None: del_data() + def test_descriptor_stat_file(self) -> None: + jdata = j_loader("water_se_a.json") + data = DataSystem( + jdata["systems"], + "set", + 1, + 1, + jdata["model"]["descriptor"]["rcut"], + run_opt=None, + ) + test_data = data.get_test() + + jdata["model"]["descriptor"].pop("type", None) + descrpt = DescrptSeA(**jdata["model"]["descriptor"], uniform_seed=True) + jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() + jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() + fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) + model = EnerModel(descrpt, fitting) + + input_data = { + "coord": [test_data["coord"]], + "box": [test_data["box"]], + "type": [test_data["type"]], + "natoms_vec": [test_data["natoms_vec"]], + "default_mesh": [test_data["default_mesh"]], + } + with TemporaryDirectory() as tempdir: + stat_path = DPPath(str(Path(tempdir)), "a") + model._compute_input_stat(input_data, stat_file_path=stat_path) + self.assertTrue( + any(child.is_dir() for child in Path(tempdir).iterdir()), + "Descriptor stat hash directory should be created", + ) + def test_model_atom_ener(self) -> None: jfile = "water_se_a.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_stat_file.py b/source/tests/tf/test_stat_file.py new file mode 100644 index 0000000000..2f5b8bbe98 --- /dev/null +++ b/source/tests/tf/test_stat_file.py @@ -0,0 +1,206 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import os +import tempfile +import unittest +from pathlib import ( + Path, +) + +from deepmd.common import ( + j_loader, +) +from deepmd.tf.descriptor.stat import ( + load_or_compute_se_input_stats, +) +from deepmd.tf.entrypoints.train import ( + _do_work, +) +from deepmd.tf.env import ( + tf, +) +from deepmd.tf.train.run_options import ( + RunOptions, +) +from deepmd.tf.utils.argcheck import ( + normalize, +) +from deepmd.tf.utils.compat import ( + update_deepmd_input, +) +from deepmd.utils.path import ( + DPPath, +) + +from .common import ( + tests_path, +) + + +class _FakeSeADescriptor: + rcut_smth = 0.5 + + def get_sel(self) -> list[int]: + return [2, 4] + + def get_ntypes(self) -> int: + return 2 + + def get_rcut(self) -> float: + return 4.0 + + +class _FakeSeRDescriptor: + rcut_r_smth = 0.5 + + def __init__(self) -> None: + self.sel_r = [2, 4] + + def get_ntypes(self) -> int: + return 2 + + def get_rcut(self) -> float: + return 4.0 + + +class _FakeSeAttenDescriptor(_FakeSeADescriptor): + pass + + +class TestStatFile(unittest.TestCase): + def setUp(self) -> None: + tf.reset_default_graph() + # Use a minimal config for testing + self.config_file = str(tests_path / "model_compression" / "input.json") + self.jdata = j_loader(self.config_file) + # Add missing type field for fitting_net + self.jdata["model"]["fitting_net"]["type"] = "ener" + # Move data_stat_nbatch to model section + self.jdata["model"]["data_stat_nbatch"] = 1 + # Fix the data path to be absolute + data_path = str(tests_path / "model_compression" / "data") + self.jdata["training"]["training_data"]["systems"] = [data_path] + self.jdata["training"]["validation_data"]["systems"] = [data_path] + # Reduce number of steps and data for faster testing + self.jdata["training"]["numb_steps"] = 10 + self.jdata["training"]["disp_freq"] = 1 + self.jdata["training"]["save_freq"] = 5 + self.jdata = normalize(update_deepmd_input(self.jdata, warning=False)) + + def tearDown(self) -> None: + tf.reset_default_graph() + + def test_stat_file_tf(self) -> None: + """Test that stat_file parameter works in TensorFlow training.""" + with tempfile.TemporaryDirectory() as temp_dir: + stat_file_path = os.path.join(temp_dir, "stat_files") + + # Add stat_file to training config + self.jdata["training"]["stat_file"] = stat_file_path + self.jdata["training"]["disp_file"] = os.path.join(temp_dir, "lcurve.out") + self.jdata["training"]["save_ckpt"] = os.path.join(temp_dir, "model.ckpt") + + # Create run options + run_opt = RunOptions( + init_model=None, + restart=None, + init_frz_model=None, + finetune=None, + log_path=None, + log_level=20, # INFO level + mpi_log="master", + ) + + # Run training - this should create the stat file + _do_work(self.jdata, run_opt, is_compress=False) + + # Check if stat files were created + stat_path = Path(stat_file_path) + self.assertTrue(stat_path.exists(), "Stat file directory should be created") + self.assertTrue(stat_path.is_dir(), "Stat file path should be a directory") + type_path = stat_path / "O H" + self.assertTrue(type_path.is_dir(), "Type-map stat directory should exist") + self.assertTrue((type_path / "bias_atom_energy").is_file()) + self.assertTrue((type_path / "std_atom_energy").is_file()) + self.assertTrue( + any(child.is_dir() for child in type_path.iterdir()), + "Descriptor stat hash directory should be created", + ) + + +class TestDescriptorStatFile(unittest.TestCase): + def _assert_load_round_trip( + self, + descrpt, + stat_dict: dict, + last_dim: int, + mixed_types: bool = False, + ) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + stat_path = DPPath(str(Path(temp_dir)), "a") + + saved = load_or_compute_se_input_stats( + descrpt, + stat_path, + last_dim=last_dim, + compute=lambda: stat_dict, + mixed_types=mixed_types, + ) + + def fail_if_recomputed() -> dict: + self.fail("Descriptor statistics should have been loaded from file") + + loaded = load_or_compute_se_input_stats( + descrpt, + stat_path, + last_dim=last_dim, + compute=fail_if_recomputed, + mixed_types=mixed_types, + ) + + self.assertEqual(saved, stat_dict) + self.assertEqual(loaded, stat_dict) + + def test_se_a_descriptor_stats_reload_from_file(self) -> None: + """Exercise the angular descriptor-stat save/load path.""" + self._assert_load_round_trip( + _FakeSeADescriptor(), + { + "sumr": [[1.0, 2.0]], + "sumn": [[3.0, 4.0]], + "sumr2": [[5.0, 6.0]], + "suma": [[7.0, 8.0]], + "suma2": [[9.0, 10.0]], + }, + last_dim=4, + ) + + def test_se_r_descriptor_stats_reload_from_file(self) -> None: + """Exercise the radial-only descriptor-stat save/load path.""" + self._assert_load_round_trip( + _FakeSeRDescriptor(), + { + "sumr": [[1.0, 2.0]], + "sumn": [[3.0, 4.0]], + "sumr2": [[5.0, 6.0]], + }, + last_dim=1, + ) + + def test_se_atten_descriptor_stats_reload_from_mixed_type_hash(self) -> None: + """Exercise the mixed-type descriptor-stat hash branch used by se_atten.""" + self._assert_load_round_trip( + _FakeSeAttenDescriptor(), + { + "sumr": [[1.0, 2.0]], + "sumn": [[3.0, 4.0]], + "sumr2": [[5.0, 6.0]], + "suma": [[7.0, 8.0]], + "suma2": [[9.0, 10.0]], + }, + last_dim=4, + mixed_types=True, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/tf/test_stat_file_integration.py b/source/tests/tf/test_stat_file_integration.py new file mode 100644 index 0000000000..4cb7536ef8 --- /dev/null +++ b/source/tests/tf/test_stat_file_integration.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Integration test to validate stat_file functionality end-to-end.""" + +import json +import os +import tempfile +import unittest +from pathlib import ( + Path, +) + +from deepmd.tf.entrypoints.train import ( + train, +) +from deepmd.tf.env import ( + tf, +) + +# Get the test data directory +tests_path = Path(__file__).parent.parent.parent.parent / "examples" + + +class TestStatFileIntegration(unittest.TestCase): + def setUp(self) -> None: + tf.reset_default_graph() + + def tearDown(self) -> None: + tf.reset_default_graph() + + def test_stat_file_path_is_accepted_and_created(self) -> None: + """Test that TF training accepts training.stat_file and creates its directory.""" + # Create a minimal training configuration + config = { + "model": { + "type_map": ["O", "H"], + "data_stat_nbatch": 1, + "descriptor": { + "type": "se_e2_a", + "sel": [2, 4], + "rcut_smth": 0.50, + "rcut": 1.00, + "neuron": [4, 8], + "resnet_dt": False, + "axis_neuron": 4, + "seed": 1, + }, + "fitting_net": {"neuron": [8, 8], "resnet_dt": True, "seed": 1}, + }, + "learning_rate": { + "type": "exp", + "decay_steps": 100, + "start_lr": 0.001, + "stop_lr": 1e-8, + }, + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + }, + "training": { + "training_data": { + "systems": [ + str(tests_path / "water" / "data" / "data_0") + ], # Use actual test data + "batch_size": 1, + }, + "numb_steps": 2, # Very short training + "disp_freq": 1, + "save_freq": 1, + }, + } + + with tempfile.TemporaryDirectory() as temp_dir: + # Create config file + config_file = os.path.join(temp_dir, "input.json") + stat_file_path = os.path.join(temp_dir, "stat_files") + + # Add stat_file to config + config["training"]["stat_file"] = stat_file_path + config["training"]["disp_file"] = os.path.join(temp_dir, "lcurve.out") + config["training"]["save_ckpt"] = os.path.join(temp_dir, "model.ckpt") + + # Write config + with open(config_file, "w") as f: + json.dump(config, f, indent=2) + + # Run a short training and verify stat_file is accepted by the TF pipeline. + train( + INPUT=config_file, + init_model=None, + restart=None, + output=os.path.join(temp_dir, "output.json"), + init_frz_model=None, + mpi_log="master", + log_level=20, + log_path=None, + is_compress=False, + skip_neighbor_stat=True, + finetune=None, + use_pretrain_script=False, + ) + + # The main validation is that the code didn't crash with an unrecognized parameter + # and that the stat file directory was created. + stat_path = Path(stat_file_path) + self.assertTrue(stat_path.exists(), "Stat file path should be created") + self.assertTrue(stat_path.is_dir(), "Stat file path should be a directory") + type_path = stat_path / "O H" + self.assertTrue(type_path.is_dir(), "Type-map stat directory should exist") + self.assertTrue( + any(child.is_dir() for child in type_path.iterdir()), + "Descriptor stat hash directory should be created", + ) + + +if __name__ == "__main__": + unittest.main()