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()