fix(stat): cap data-stat batches at the available batches per system

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]:

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

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