fix(pt-expt): group LMDB statistics by atom count

deepmd/pt_expt/utils/lmdb_dataset.py

@@ -32,7 +32,10 @@ class LmdbDataSystem:
     Exposes the small surface that pt_expt's trainer touches:
     ``get_batch(sys_idx=None)``, ``add_data_requirements(list)``, and
     ``get_nsystems()``. Internally uses :class:`LmdbDataReader` for I/O and
-    :class:`SameNlocBatchSampler` to draw same-nloc batches.
+    :class:`SameNlocBatchSampler` to draw same-nloc batches. Statistics use a
+    separate logical-system view in which every ``nloc`` group is sampled
+    independently, matching the PyTorch DataLoader adapter without changing
+    the identity of the LMDB as one training dataset.
 
     Parameters
     ----------
@@ -76,6 +79,8 @@ def __init__(
             block_targets=block_targets,
         )
         self._iter = iter(self._sampler)
+        self._stat_nlocs = tuple(sorted(self._reader.nloc_groups))
+        self._stat_offsets = [0] * len(self._stat_nlocs)
 
     # ------------------------------------------------------------------
     # pt_expt trainer surface
@@ -93,6 +98,60 @@ def get_batch(self, sys_idx: int | None = None) -> dict[str, Any]:
         except StopIteration:
             self._iter = iter(self._sampler)
             indices = next(self._iter)
+        return self._collate_indices(indices)
+
+    def get_stat_batch(self, sys_idx: int) -> dict[str, Any]:
+        """Return one batch from a fixed-``nloc`` statistical system.
+
+        Parameters
+        ----------
+        sys_idx : int
+            Index into the sorted ``nloc`` groups.
+
+        Returns
+        -------
+        dict[str, Any]
+            A collated NumPy batch whose frames have one atom count.
+
+        Raises
+        ------
+        IndexError
+            If ``sys_idx`` does not identify an available ``nloc`` group.
+        """
+        if not 0 <= sys_idx < len(self._stat_nlocs):
+            raise IndexError(
+                f"Statistical system index {sys_idx} is out of range for "
+                f"{len(self._stat_nlocs)} nloc groups."
+            )
+
+        nloc = self._stat_nlocs[sys_idx]
+        group_indices = self._reader.nloc_groups[nloc]
+        batch_size = self._reader.get_batch_size_for_nloc(nloc)
+        start = self._stat_offsets[sys_idx]
+        if start >= len(group_indices):
+            start = 0
+        stop = min(start + batch_size, len(group_indices))
+        self._stat_offsets[sys_idx] = stop
+        return self._collate_indices(group_indices[start:stop])
+
+    def get_stat_nsystems(self) -> int:
+        """Return the number of fixed-``nloc`` statistical systems."""
+        return len(self._stat_nlocs)
+
+    def get_stat_numb_batches(self, sys_idx: int) -> int:
+        """Return the available batch count for one statistical system."""
+        if not 0 <= sys_idx < len(self._stat_nlocs):
+            raise IndexError(
+                f"Statistical system index {sys_idx} is out of range for "
+                f"{len(self._stat_nlocs)} nloc groups."
+            )
+        nloc = self._stat_nlocs[sys_idx]
+        nframes = len(self._reader.nloc_groups[nloc])
+        batch_size = self._reader.get_batch_size_for_nloc(nloc)
+        return (nframes + batch_size - 1) // batch_size
+
+    def _collate_indices(self, indices: list[int]) -> dict[str, Any]:
+        """Load and collate the requested dataset indices."""
         frames = [self._reader[int(i)] for i in indices]
         return collate_lmdb_frames(frames)
 
@@ -102,11 +161,7 @@ def add_data_requirements(
         self._reader.add_data_requirement(data_requirement)
 
     def get_nsystems(self) -> int:
-        """Return 1: pt_expt's stat collection treats LMDB as a single system.
-
-        Per-system sampling within the LMDB is handled by
-        ``SameNlocBatchSampler`` + ``block_targets``.
-        """
+        """Return one logical LMDB training dataset."""
         return 1
 
     # ------------------------------------------------------------------

deepmd/utils/model_stat.py

@@ -16,11 +16,35 @@
 log = logging.getLogger(__name__)
 
 
+def _get_stat_nsystems(data: Any) -> int:
+    """Return the number of shape-compatible systems used for statistics."""
+    get_stat_nsystems = getattr(data, "get_stat_nsystems", None)
+    if get_stat_nsystems is not None:
+        return int(get_stat_nsystems())
+    return int(data.get_nsystems())
+
+
+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:
+        return nbatches
+    return min(nbatches, int(get_stat_numb_batches(sys_idx)))
+
+
+def _get_stat_batch(data: Any, sys_idx: int) -> dict[str, Any]:
+    """Return one batch from a shape-compatible statistical system."""
+    get_stat_batch = getattr(data, "get_stat_batch", None)
+    if get_stat_batch is not None:
+        return get_stat_batch(sys_idx)
+    return data.get_batch(sys_idx=sys_idx)
+
+
 def _make_all_stat_ref(data: Any, nbatches: int) -> dict[str, list[Any]]:
     all_stat = defaultdict(list)
-    for ii in range(data.get_nsystems()):
-        for jj in range(nbatches):
-            stat_data = data.get_batch(sys_idx=ii)
+    for ii in range(_get_stat_nsystems(data)):
+        for jj in range(_get_stat_numb_batches(data, ii, nbatches)):
+            stat_data = _get_stat_batch(data, ii)
             for dd in stat_data:
                 if dd == "natoms_vec":
                     stat_data[dd] = stat_data[dd].astype(np.int32)
@@ -39,7 +63,11 @@ def collect_batches(
     Parameters
     ----------
     data
-        The data (must support ``get_nsystems()`` and ``get_batch(sys_idx=)``)
+        The data source. It must support ``get_nsystems()`` and
+        ``get_batch(sys_idx=)``. Optional ``get_stat_nsystems()``,
+        ``get_stat_numb_batches(sys_idx)``, and ``get_stat_batch(sys_idx)``
+        hooks may expose shape-compatible logical systems and their available
+        batches specifically for statistics.
     nbatches : int
         The number of batches per system
     merge_sys : bool (True)
@@ -55,10 +83,10 @@ def collect_batches(
             all_stat[key][batch_idx][frame_idx]
     """
     all_stat = defaultdict(list)
-    for ii in range(data.get_nsystems()):
+    for ii in range(_get_stat_nsystems(data)):
         sys_stat = defaultdict(list)
-        for jj in range(nbatches):
-            stat_data = data.get_batch(sys_idx=ii)
+        for jj in range(_get_stat_numb_batches(data, ii, nbatches)):
+            stat_data = _get_stat_batch(data, ii)
             for dd in stat_data:
                 if dd == "natoms_vec":
                     stat_data[dd] = stat_data[dd].astype(np.int32)
@@ -78,16 +106,21 @@ def make_stat_input(
 ) -> list[dict[str, np.ndarray]]:
     """Pack data for statistics using DeepmdDataSystem.
 
-    Collects *nbatches* batches from each system and concatenates them
-    into a single dict per system.  The returned format
+    Collects up to *nbatches* batches from each shape-compatible statistical
+    system and concatenates them into one dictionary per system. Data sources
+    with variable atom counts may expose dedicated statistical-system methods
+    so incompatible ``nloc`` groups remain separate. The returned format
     (``list[dict[str, np.ndarray]]``) is backend-agnostic and can be
     consumed by ``compute_or_load_stat`` in dpmodel, pt_expt, and jax.
 
     Parameters
     ----------
     data
-        The multi-system data manager
-        (must support ``get_nsystems()`` and ``get_batch(sys_idx=)``).
+        The multi-system data manager. It must support ``get_nsystems()`` and
+        ``get_batch(sys_idx=)``. Optional ``get_stat_nsystems()``,
+        ``get_stat_numb_batches(sys_idx)``, and ``get_stat_batch(sys_idx)``
+        hooks may expose shape-compatible logical systems and their available
+        batches specifically for statistics.
     nbatches : int
         Number of batches to collect per system.
 
@@ -98,7 +131,7 @@ def make_stat_input(
     """
     all_stat = collect_batches(data, nbatches, merge_sys=False)
 
-    nsystems = data.get_nsystems()
+    nsystems = _get_stat_nsystems(data)
     log.info(f"Packing data for statistics from {nsystems} systems")
 
     keys = list(all_stat.keys())

source/tests/pt_expt/test_lmdb_training.py

@@ -28,6 +28,9 @@
 from deepmd.pt_expt.utils.lmdb_dataset import (
     LmdbDataSystem,
 )
+from deepmd.pt_expt.utils.stat import (
+    make_stat_input,
+)
 from deepmd.utils.argcheck import (
     normalize,
 )
@@ -85,13 +88,42 @@ def _create_test_lmdb(path: str, nframes: int, natoms: int) -> None:
     env.close()
 
 
+def _create_mixed_nloc_test_lmdb(path: str) -> None:
+    """Write an LMDB with five six-atom and five nine-atom frames."""
+    frame_nlocs = [6] * 5 + [9] * 5
+    env = lmdb.open(path, map_size=10 * 1024 * 1024)
+    fmt = "012d"
+    metadata = {
+        "nframes": len(frame_nlocs),
+        "frame_idx_fmt": fmt,
+        "frame_nlocs": frame_nlocs,
+        "type_map": ["O", "H"],
+        "system_info": {
+            "formula": "mixed",
+            "natoms": [3, 3],
+            "nframes": len(frame_nlocs),
+        },
+    }
+    with env.begin(write=True) as txn:
+        txn.put(b"__metadata__", msgpack.packb(metadata, use_bin_type=True))
+        for index, nloc in enumerate(frame_nlocs):
+            key = format(index, fmt).encode()
+            txn.put(
+                key,
+                msgpack.packb(_make_frame(nloc, index), use_bin_type=True),
+            )
+    env.close()
+
+
 class TestLmdbDataSystemGetBatch(unittest.TestCase):
     """LmdbDataSystem.get_batch produces a numpy dict that normalize_batch accepts."""
 
     def setUp(self) -> None:
         self.tmpdir = tempfile.mkdtemp()
         self.lmdb_path = os.path.join(self.tmpdir, "test.lmdb")
+        self.mixed_lmdb_path = os.path.join(self.tmpdir, "mixed.lmdb")
         _create_test_lmdb(self.lmdb_path, nframes=8, natoms=6)
+        _create_mixed_nloc_test_lmdb(self.mixed_lmdb_path)
 
     def tearDown(self) -> None:
         shutil.rmtree(self.tmpdir, ignore_errors=True)
@@ -159,6 +191,58 @@ def test_add_data_requirements_passthrough(self) -> None:
         self.assertIn("energy", batch)
         self.assertIn("find_energy", batch)
 
+    def test_stat_input_partitions_mixed_nloc_batches(self) -> None:
+        """Statistics expose each atom-count group as one logical system."""
+        ds = LmdbDataSystem(
+            lmdb_path=self.mixed_lmdb_path,
+            type_map=["O", "H"],
+            batch_size=2,
+            seed=0,
+        )
+
+        self.assertEqual(ds.get_nsystems(), 1)
+        self.assertEqual(ds.get_stat_nsystems(), 2)
+        sampled = make_stat_input(ds, nbatches=10)
+
+        self.assertEqual(len(sampled), 2)
+        self.assertEqual(
+            sorted(sample["atype"].shape for sample in sampled),
+            [(5, 6), (5, 9)],
+        )
+
+    def test_stat_input_limits_batches_per_nloc_group(self) -> None:
+        """Statistics honor the requested batch cap for every nloc group."""
+        ds = LmdbDataSystem(
+            lmdb_path=self.mixed_lmdb_path,
+            type_map=["O", "H"],
+            batch_size=2,
+            seed=0,
+        )
+
+        sampled = make_stat_input(ds, nbatches=2)
+
+        self.assertEqual(
+            sorted(sample["atype"].shape for sample in sampled),
+            [(4, 6), (4, 9)],
+        )
+
+    def test_stat_batches_restart_after_non_divisible_pass(self) -> None:
+        """Statistical batches restart after including a partial tail batch."""
+        ds = LmdbDataSystem(
+            lmdb_path=self.mixed_lmdb_path,
+            type_map=["O", "H"],
+            batch_size=2,
+            seed=0,
+        )
+
+        for sys_idx, nloc in enumerate((6, 9)):
+            self.assertEqual(ds.get_stat_numb_batches(sys_idx), 3)
+            shapes = [ds.get_stat_batch(sys_idx)["atype"].shape for _ in range(4)]
+            self.assertEqual(
+                shapes,
+                [(2, nloc), (2, nloc), (1, nloc), (2, nloc)],
+            )
+
 
 class TestLmdbTrainingLoop(unittest.TestCase):
     """End-to-end: get_trainer routes an LMDB path and runs training steps."""
@@ -167,8 +251,10 @@ class TestLmdbTrainingLoop(unittest.TestCase):
     def setUpClass(cls) -> None:
         cls.tmpdir = tempfile.mkdtemp()
         cls.lmdb_path = os.path.join(cls.tmpdir, "train.lmdb")
+        cls.mixed_lmdb_path = os.path.join(cls.tmpdir, "mixed.lmdb")
         cls.val_lmdb_path = os.path.join(cls.tmpdir, "val.lmdb")
         _create_test_lmdb(cls.lmdb_path, nframes=8, natoms=6)
+        _create_mixed_nloc_test_lmdb(cls.mixed_lmdb_path)
         _create_test_lmdb(cls.val_lmdb_path, nframes=4, natoms=6)
 
     @classmethod
@@ -244,6 +330,25 @@ def test_get_trainer_routes_lmdb(self) -> None:
         finally:
             os.chdir(cwd)
 
+    def test_mixed_nloc_statistics_and_training(self) -> None:
+        """Trainer computes statistics and trains across fixed-nloc batches."""
+        config = self._make_lmdb_config(numb_steps=2)
+        config["model"]["data_stat_nbatch"] = 10
+        config["training"]["training_data"]["systems"] = self.mixed_lmdb_path
+        config = update_deepmd_input(config, warning=False)
+        config = normalize(config)
+
+        with tempfile.TemporaryDirectory(dir=self.tmpdir) as run_dir:
+            cwd = os.getcwd()
+            os.chdir(run_dir)
+            try:
+                trainer = get_trainer(config)
+                self.assertEqual(trainer.training_data.get_nsystems(), 1)
+                self.assertEqual(trainer.training_data.get_stat_nsystems(), 2)
+                trainer.run()
+            finally:
+                os.chdir(cwd)
+
 
 if __name__ == "__main__":
     unittest.main()