feat(pt): full validation support lmdb format (#5419)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Full validation now supports LMDB-backed validation datasets with
snapshot caching and group-by-atom-count evaluation.
* Datasets can register and expose data-requirement metadata for
downstream tooling.
* Added a view adapter to evaluate mixed atom-count validation frames
group-by-group.

* **Refactor**
  * Consolidated per-atom-count test-data wrapper into a shared utility.

* **Tests**
* Added unit and integration tests for mixed-atom-count views,
data-requirement registration, and validation iteration.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: OutisLi

deepmd/dpmodel/utils/__init__.py

@@ -10,6 +10,7 @@
     DistributedSameNlocBatchSampler,
     LmdbDataReader,
     LmdbTestData,
+    LmdbTestDataNlocView,
     SameNlocBatchSampler,
     is_lmdb,
     make_neighbor_stat_data,
@@ -58,6 +59,7 @@
     "FittingNet",
     "LmdbDataReader",
     "LmdbTestData",
+    "LmdbTestDataNlocView",
     "NativeLayer",
     "NativeNet",
     "NetworkCollection",

deepmd/dpmodel/utils/lmdb_data.py

@@ -535,6 +535,11 @@ def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> N
         for item in data_requirement:
             self._data_requirements[item["key"]] = item
 
+    @property
+    def data_requirements(self) -> list[DataRequirementItem]:
+        """Registered data requirements in insertion order."""
+        return list(self._data_requirements.values())
+
     def print_summary(self, name: str, prob: Any) -> None:
         """Print basic dataset info."""
         n_groups = len(self._nloc_groups)
@@ -1286,6 +1291,25 @@ def add(
             "dtype": dtype,
         }
 
+    def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> None:
+        """Register expected keys from ``DataRequirementItem`` objects.
+
+        Mirrors :meth:`LmdbDataReader.add_data_requirement` so the same
+        requirement list can be forwarded to both the training reader and
+        the full-validation test data.
+        """
+        for item in data_requirement:
+            self.add(
+                item["key"],
+                ndof=item["ndof"],
+                atomic=item["atomic"],
+                must=item["must"],
+                high_prec=item["high_prec"],
+                repeat=item["repeat"],
+                default=item["default"],
+                dtype=item["dtype"],
+            )
+
     def _resolve_dtype(self, key: str) -> np.dtype:
         """Resolve target dtype for a key using registered requirements."""
         if key in self._requirements:
@@ -1444,6 +1468,28 @@ def _stack_frames(
         return result
 
 
+class LmdbTestDataNlocView:
+    """Thin wrapper exposing a fixed-``nloc`` view of :class:`LmdbTestData`.
+
+    The underlying :class:`LmdbTestData` groups frames by atom count. This
+    view fixes one ``nloc`` group, so ``get_test()`` returns only the frames
+    with that atom count and all other attributes (``pbc``, ``mixed_type``,
+    …) are forwarded to the underlying object. It lets downstream consumers
+    that expect a ``DeepmdData``-style system (one fixed natoms per
+    ``get_test()``) work on mixed-nloc LMDB datasets without changes.
+    """
+
+    def __init__(self, lmdb_test_data: "LmdbTestData", nloc: int) -> None:
+        self._inner = lmdb_test_data
+        self._nloc = nloc
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._inner, name)
+
+    def get_test(self) -> dict[str, Any]:
+        return self._inner.get_test(nloc=self._nloc)
+
+
 def merge_lmdb(
     src_paths: list[str],
     dst_path: str,

deepmd/entrypoints/test.py

@@ -17,6 +17,7 @@
 )
 from deepmd.dpmodel.utils.lmdb_data import (
     LmdbTestData,
+    LmdbTestDataNlocView,
     is_lmdb,
 )
 from deepmd.infer.deep_dipole import (
@@ -77,24 +78,6 @@
 log = logging.getLogger(__name__)
 
 
-class _LmdbTestDataNlocView:
-    """Thin wrapper that makes LmdbTestData.get_test() return a specific nloc group.
-
-    Delegates all attributes to the underlying LmdbTestData, but get_test()
-    returns only frames with the specified nloc.
-    """
-
-    def __init__(self, lmdb_test_data: LmdbTestData, nloc: int) -> None:
-        self._inner = lmdb_test_data
-        self._nloc = nloc
-
-    def __getattr__(self, name: str) -> Any:
-        return getattr(self._inner, name)
-
-    def get_test(self) -> dict:
-        return self._inner.get_test(nloc=self._nloc)
-
-
 def test(
     *,
     model: str,
@@ -221,7 +204,7 @@ def test(
             for nloc_val in nloc_keys:
                 label = f"{system} [nloc={nloc_val}]" if len(nloc_keys) > 1 else system
                 # Create a thin wrapper that returns only this nloc group
-                data_items.append((_LmdbTestDataNlocView(lmdb_data, nloc_val), label))
+                data_items.append((LmdbTestDataNlocView(lmdb_data, nloc_val), label))
         else:
             data = DeepmdData(
                 system,

deepmd/pt/train/validation.py

@@ -24,6 +24,10 @@
 import torch.distributed as dist
 
 from deepmd.dpmodel.common import PRECISION_DICT as NP_PRECISION_DICT
+from deepmd.dpmodel.utils.lmdb_data import (
+    LmdbTestData,
+    LmdbTestDataNlocView,
+)
 from deepmd.pt.utils.auto_batch_size import (
     AutoBatchSize,
 )
@@ -35,6 +39,9 @@
     GLOBAL_PT_FLOAT_PRECISION,
     RESERVED_PRECISION_DICT,
 )
+from deepmd.pt.utils.lmdb_dataset import (
+    LmdbDataset,
+)
 from deepmd.pt.utils.utils import (
     to_torch_tensor,
 )
@@ -55,6 +62,10 @@
 log = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
+    from collections.abc import (
+        Iterator,
+    )
+
     from deepmd.utils.data import (
         DeepmdData,
     )
@@ -229,6 +240,12 @@ def __init__(
         if self.rank == 0:
             self._initialize_best_checkpoints(restart_training=restart_training)
 
+        # Lazily-populated full test snapshot for LMDB validation. Mixed-nloc
+        # LMDB datasets cannot be stacked as a single (nframes, natoms*3)
+        # tensor, so we materialize frames grouped by nloc the first time
+        # full validation runs and reuse the snapshot on subsequent calls.
+        self._lmdb_test_data: LmdbTestData | None = None
+
     def should_run(self, display_step: int) -> bool:
         """Check whether the current step should trigger full validation."""
         if not self.enabled or self.start_step is None:
@@ -348,14 +365,10 @@ def evaluate_all_systems(self) -> dict[str, float]:
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
 
-        system_metrics = []
-        for dataset in self.validation_data.systems:
-            if not isinstance(dataset, DeepmdDataSetForLoader):
-                raise TypeError(
-                    "Full validation expects each dataset in validation_data.systems "
-                    f"to be DeepmdDataSetForLoader, got {type(dataset)!r}."
-                )
-            system_metrics.append(self._evaluate_system(dataset.data_system))
+        system_metrics = [
+            self._evaluate_system(data_system)
+            for data_system in self._iter_validation_data_systems()
+        ]
 
         aggregated = weighted_average([metric for metric in system_metrics if metric])
         return {
@@ -364,6 +377,53 @@ def evaluate_all_systems(self) -> dict[str, float]:
             if metric_key in aggregated
         }
 
+    def _iter_validation_data_systems(self) -> Iterator[Any]:
+        """Yield ``DeepmdData``-like systems to evaluate in this run.
+
+        - For ``DpLoaderSet``-style validation data, each entry in
+          ``validation_data.systems`` is a :class:`DeepmdDataSetForLoader`,
+          and we forward its underlying ``DeepmdData`` instance.
+        - For ``LmdbDataset`` validation data, we lazily materialize a
+          :class:`LmdbTestData` snapshot (cached across calls) and yield one
+          :class:`LmdbTestDataNlocView` per ``nloc`` group, so mixed-nloc
+          frames can be stacked and evaluated group by group.
+        """
+        validation_data = self.validation_data
+        if isinstance(validation_data, LmdbDataset):
+            lmdb_test_data = self._get_lmdb_test_data_snapshot(validation_data)
+            for nloc in sorted(lmdb_test_data.nloc_groups.keys()):
+                yield LmdbTestDataNlocView(lmdb_test_data, nloc)
+            return
+
+        for dataset in validation_data.systems:
+            if not isinstance(dataset, DeepmdDataSetForLoader):
+                raise TypeError(
+                    "Full validation expects each dataset in validation_data.systems "
+                    f"to be DeepmdDataSetForLoader, got {type(dataset)!r}."
+                )
+            yield dataset.data_system
+
+    def _get_lmdb_test_data_snapshot(self, lmdb_dataset: LmdbDataset) -> LmdbTestData:
+        """Build (once) and return the cached LMDB test snapshot.
+
+        Reuses the ``type_map`` and previously-registered
+        ``DataRequirementItem`` entries from the validation dataset so that
+        the full-validation snapshot sees exactly the same fields and
+        dtypes as training batches.
+        """
+        if self._lmdb_test_data is not None:
+            return self._lmdb_test_data
+
+        self._lmdb_test_data = LmdbTestData(
+            lmdb_dataset.lmdb_path,
+            type_map=list(lmdb_dataset.type_map),
+            shuffle_test=False,
+        )
+        data_requirements = lmdb_dataset.data_requirements
+        if data_requirements:
+            self._lmdb_test_data.add_data_requirement(data_requirements)
+        return self._lmdb_test_data
+
     def _evaluate_system(
         self, data_system: DeepmdData
     ) -> dict[str, tuple[float, float]]:

deepmd/pt/utils/lmdb_dataset.py

@@ -212,6 +212,14 @@ def mixed_type(self) -> bool:
     def batch_size(self) -> int:
         return self._reader.batch_size
 
+    @property
+    def type_map(self) -> list[str]:
+        return self._reader.type_map
+
+    @property
+    def data_requirements(self) -> list[DataRequirementItem]:
+        return self._reader.data_requirements
+
     def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> None:
         self._reader.add_data_requirement(data_requirement)
 

source/tests/common/dpmodel/test_lmdb_data.py

@@ -14,6 +14,7 @@
 from deepmd.dpmodel.utils.lmdb_data import (
     LmdbDataReader,
     LmdbTestData,
+    LmdbTestDataNlocView,
     SameNlocBatchSampler,
     _expand_indices_by_blocks,
     compute_block_targets,
@@ -415,6 +416,27 @@ def test_test_data_get_test_specific_nloc(self):
         r12 = td.get_test(nloc=12)
         self.assertEqual(r12["coord"].shape, (2, 12 * 3))
 
+    def test_test_data_nloc_view(self):
+        """LmdbTestDataNlocView delegates attributes and fixes nloc."""
+        td = LmdbTestData(self._lmdb_path, type_map=self._type_map, shuffle_test=False)
+        td.add("energy", 1, atomic=False, must=False, high_prec=True)
+        view = LmdbTestDataNlocView(td, 9)
+
+        self.assertEqual(view.pbc, td.pbc)
+        self.assertIs(view.nloc_groups, td.nloc_groups)
+
+        expected = td.get_test(nloc=9)
+        actual = view.get_test()
+        self.assertEqual(actual["coord"].shape, (4, 9 * 3))
+        self.assertEqual(actual["type"].shape, (4, 9))
+        self.assertEqual(actual.keys(), expected.keys())
+        for key, expected_value in expected.items():
+            actual_value = actual[key]
+            if isinstance(expected_value, np.ndarray):
+                np.testing.assert_array_equal(actual_value, expected_value)
+            else:
+                self.assertEqual(actual_value, expected_value)
+
     def test_test_data_get_test_default_mixed(self):
         td = LmdbTestData(self._lmdb_path, type_map=self._type_map, shuffle_test=False)
         td.add("energy", 1, atomic=False, must=False, high_prec=True)
@@ -851,6 +873,66 @@ def test_testdata_repeat_applied(self):
             (self._nframes, self._natoms * 3),
         )
 
+    def test_testdata_add_data_requirement_matches_manual_add(self):
+        """DataRequirementItem forwarding matches manual requirement registration."""
+        from deepmd.utils.data import (
+            DataRequirementItem,
+        )
+
+        requirements = [
+            DataRequirementItem(
+                "drdq",
+                ndof=6,
+                atomic=True,
+                must=False,
+                high_prec=False,
+                repeat=2,
+                default=1.25,
+                dtype=np.float64,
+            ),
+            DataRequirementItem(
+                "aux",
+                ndof=2,
+                atomic=False,
+                must=False,
+                high_prec=False,
+                repeat=3,
+                default=-2.0,
+                dtype=np.float32,
+            ),
+        ]
+        manual = LmdbTestData(
+            self._lmdb_path,
+            type_map=self._type_map,
+            shuffle_test=False,
+        )
+        forwarded = LmdbTestData(
+            self._lmdb_path,
+            type_map=self._type_map,
+            shuffle_test=False,
+        )
+        for item in requirements:
+            manual.add(
+                item["key"],
+                ndof=item["ndof"],
+                atomic=item["atomic"],
+                must=item["must"],
+                high_prec=item["high_prec"],
+                repeat=item["repeat"],
+                default=item["default"],
+                dtype=item["dtype"],
+            )
+        forwarded.add_data_requirement(requirements)
+
+        manual_result = manual.get_test()
+        forwarded_result = forwarded.get_test()
+        for item in requirements:
+            key = item["key"]
+            self.assertEqual(forwarded_result[f"find_{key}"], 0.0)
+            self.assertEqual(forwarded_result[key].shape, manual_result[key].shape)
+            self.assertEqual(forwarded_result[key].dtype, manual_result[key].dtype)
+            np.testing.assert_array_equal(forwarded_result[key], manual_result[key])
+
     def test_testdata_missing_key_not_found(self):
         """Keys absent from LMDB frames get find_*=0.0 in get_test()."""
         tmpdir = tempfile.TemporaryDirectory()