From 7558f3797ec577b5a76de9f5a8ddfcb0a93a391f Mon Sep 17 00:00:00 2001 From: ShryukGrandhi <85529990+ShryukGrandhi@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:50:26 -0700 Subject: [PATCH] fix(stat): cap data-stat batches at the available batches per system The backend-agnostic collect_batches/make_stat_input used by TF, JAX and pt_expt drew data_stat_nbatch batches from every system, wrapping around and resampling systems with fewer batches. The PT backend stops at the dataloader length. Cap the count at the system's batch count so small systems are not overweighted and the backends collect the same frames. Signed-off-by: ShryukGrandhi <85529990+ShryukGrandhi@users.noreply.github.com> --- deepmd/utils/model_stat.py | 14 +++-- .../tests/common/dpmodel/test_dist_check.py | 3 +- .../tests/consistent/test_make_stat_input.py | 62 +++++++++++++++++++ 3 files changed, 74 insertions(+), 5 deletions(-) diff --git a/deepmd/utils/model_stat.py b/deepmd/utils/model_stat.py index 109de19fd5..74ab9c7914 100644 --- a/deepmd/utils/model_stat.py +++ b/deepmd/utils/model_stat.py @@ -29,11 +29,17 @@ def _get_stat_nsystems(data: Any) -> int: def _get_stat_numb_batches(data: Any, sys_idx: int, nbatches: int) -> int: - """Limit sampling to the available batches of one statistical system.""" - get_stat_numb_batches = getattr(data, "get_stat_numb_batches", None) - if get_stat_numb_batches is None: + """Limit sampling to the available batches of one statistical system. + + A system with fewer batches than ``nbatches`` contributes each of its + batches once instead of being resampled, matching the PyTorch backend. + """ + if ( + getattr(data, "get_stat_numb_batches", None) is None + and getattr(data, "get_nbatches", None) is None + ): return nbatches - return min(nbatches, int(get_stat_numb_batches(sys_idx))) + return min(nbatches, _get_stat_pass_length(data, sys_idx)) def _get_stat_batch(data: Any, sys_idx: int) -> dict[str, Any]: diff --git a/source/tests/common/dpmodel/test_dist_check.py b/source/tests/common/dpmodel/test_dist_check.py index 7f8fb9b929..d498a2b6e1 100644 --- a/source/tests/common/dpmodel/test_dist_check.py +++ b/source/tests/common/dpmodel/test_dist_check.py @@ -605,7 +605,8 @@ def test_disabled_filter_takes_the_next_batches(self) -> None: data = _CyclicDataSystem([[0.2, 0.3, 1.5], [0.1, 0.2]]) packed = make_stat_input(data, nbatches=3) np.testing.assert_array_equal(packed[0]["coord"][:, 0], [0.0, 1.0, 2.0]) - np.testing.assert_array_equal(packed[1]["coord"][:, 0], [0.0, 1.0, 0.0]) + # A system with fewer batches than requested is not resampled. + np.testing.assert_array_equal(packed[1]["coord"][:, 0], [0.0, 1.0]) def test_dataset_without_valid_frames_is_reported(self) -> None: data = _CyclicDataSystem([[0.2, 0.3], [0.1]]) diff --git a/source/tests/consistent/test_make_stat_input.py b/source/tests/consistent/test_make_stat_input.py index 5a927f5541..88920249e3 100644 --- a/source/tests/consistent/test_make_stat_input.py +++ b/source/tests/consistent/test_make_stat_input.py @@ -8,6 +8,8 @@ """ import os +import shutil +import tempfile import unittest import numpy as np @@ -367,5 +369,65 @@ def test_consistency(self) -> None: ) +def _make_truncated_system(src: str, dst: str, nframes: int) -> str: + """Copy the first set of ``src`` to ``dst`` keeping ``nframes`` frames.""" + os.makedirs(os.path.join(dst, "set.000")) + for name in ("type.raw", "type_map.raw"): + shutil.copy(os.path.join(src, name), os.path.join(dst, name)) + src_set = os.path.join(src, "set.000") + for name in os.listdir(src_set): + arr = np.load(os.path.join(src_set, name)) + np.save(os.path.join(dst, "set.000", name), arr[:nframes]) + return dst + + +@unittest.skipUnless(INSTALLED_PT, "PyTorch backend not installed") +class TestMakeStatInputSmallSystems(unittest.TestCase): + """Systems with fewer batches than data_stat_nbatch must not be oversampled.""" + + def setUp(self) -> None: + self.tmpdir = tempfile.mkdtemp() + + def tearDown(self) -> None: + shutil.rmtree(self.tmpdir) + + def test_consistency(self) -> None: + src = os.path.join(EXAMPLE_DIR, "data", "data_0") + if not os.path.isdir(src): + self.skipTest(f"Example data not found: {src}") + nframes = [3, 7] + systems = [ + _make_truncated_system(src, os.path.join(self.tmpdir, f"sys_{nf}"), nf) + for nf in nframes + ] + config = _build_config(systems, ["O", "H"], data_stat_nbatch=10) + + universal_stat = _get_universal_stat(config, _ENER_DATA_REQ) + pt_stat = _get_pt_stat(config, _ENER_DATA_REQ) + + _compare_stat(self, universal_stat, pt_stat, _COMMON_CHECK_KEYS) + for sys_idx, nf in enumerate(nframes): + self.assertEqual(universal_stat[sys_idx]["energy"].shape[0], nf) + self.assertEqual(pt_stat[sys_idx]["energy"].shape[0], nf) + + # the TF backend collects batches through collect_batches + from deepmd.utils.data_system import ( + DeepmdDataSystem, + ) + from deepmd.utils.model_stat import ( + collect_batches, + ) + + data = DeepmdDataSystem( + systems=systems, + batch_size=1, + test_size=1, + rcut=config["model"]["descriptor"]["rcut"], + type_map=["O", "H"], + ) + all_stat = collect_batches(data, 10, merge_sys=False) + self.assertEqual([len(sys_box) for sys_box in all_stat["box"]], nframes) + + if __name__ == "__main__": unittest.main()