fixup

OutisLi · OutisLi · commit a27d0ae5c89b · 2026-04-26T16:29:55.000+08:00
diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/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)
diff --git a/deepmd/pt/train/validation.py b/deepmd/pt/train/validation.py
@@ -407,20 +407,19 @@ def _get_lmdb_test_data_snapshot(self, lmdb_dataset: LmdbDataset) -> LmdbTestDat
         """Build (once) and return the cached LMDB test snapshot.
 
         Reuses the ``type_map`` and previously-registered
-        ``DataRequirementItem`` entries on ``LmdbDataset._reader`` so that
+        ``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
 
-        reader = lmdb_dataset._reader
         self._lmdb_test_data = LmdbTestData(
             lmdb_dataset.lmdb_path,
-            type_map=list(reader._type_map),
+            type_map=list(lmdb_dataset.type_map),
             shuffle_test=False,
         )
-        data_requirements = list(reader._data_requirements.values())
+        data_requirements = lmdb_dataset.data_requirements
         if data_requirements:
             self._lmdb_test_data.add_data_requirement(data_requirements)
         return self._lmdb_test_data
diff --git a/deepmd/pt/utils/lmdb_dataset.py b/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)
 
diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/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()
diff --git a/source/tests/pt/test_validation.py b/source/tests/pt/test_validation.py
@@ -8,7 +8,13 @@
 from pathlib import (
     Path,
 )
+from unittest.mock import (
+    patch,
+)
 
+import lmdb
+import msgpack
+import numpy as np
 import torch
 from dargs.dargs import (
     ArgumentValueError,
@@ -20,6 +26,9 @@
     FullValidator,
     resolve_full_validation_start_step,
 )
+from deepmd.pt.utils.lmdb_dataset import (
+    LmdbDataset,
+)
 from deepmd.utils.argcheck import (
     normalize,
 )
@@ -45,6 +54,85 @@ def get_dim_aparam(self) -> int:
         return 0
 
 
+def _make_lmdb_frame(natoms: int, seed: int) -> dict:
+    """Create one synthetic LMDB frame for full-validation tests."""
+    rng = np.random.RandomState(seed)
+    n_type0 = max(1, natoms // 3)
+    n_type1 = natoms - n_type0
+    atype = np.array([0] * n_type0 + [1] * n_type1, dtype=np.int64)
+    return {
+        "atom_names": ["O", "H"],
+        "atom_numbs": [
+            {
+                "type": "<i8",
+                "shape": (1,),
+                "data": np.array([n_type0], dtype=np.int64).tobytes(),
+            },
+            {
+                "type": "<i8",
+                "shape": (1,),
+                "data": np.array([n_type1], dtype=np.int64).tobytes(),
+            },
+        ],
+        "atom_types": {
+            "type": "<i8",
+            "shape": (natoms,),
+            "data": atype.tobytes(),
+        },
+        "coords": {
+            "type": "<f8",
+            "shape": (natoms, 3),
+            "data": rng.randn(natoms, 3).astype(np.float64).tobytes(),
+        },
+        "cells": {
+            "type": "<f8",
+            "shape": (3, 3),
+            "data": (np.eye(3) * 10.0).astype(np.float64).tobytes(),
+        },
+        "energies": {
+            "type": "<f8",
+            "shape": (1,),
+            "data": rng.randn(1).astype(np.float64).tobytes(),
+        },
+        "forces": {
+            "type": "<f8",
+            "shape": (natoms, 3),
+            "data": rng.randn(natoms, 3).astype(np.float64).tobytes(),
+        },
+    }
+
+
+def _create_mixed_nloc_lmdb(path: str) -> str:
+    """Create a mixed-nloc LMDB dataset with 6, 9, and 12-atom frames."""
+    frame_specs = [(6, 4), (9, 4), (12, 2)]
+    total_frames = sum(count for _, count in frame_specs)
+    env = lmdb.open(path, map_size=10 * 1024 * 1024)
+    with env.begin(write=True) as txn:
+        metadata = {
+            "nframes": total_frames,
+            "frame_idx_fmt": "012d",
+            "type_map": ["O", "H"],
+            "system_info": {
+                "natoms": [2, 4],
+                "formula": "mixed",
+            },
+        }
+        txn.put(b"__metadata__", msgpack.packb(metadata, use_bin_type=True))
+        frame_idx = 0
+        for natoms, count in frame_specs:
+            for _ in range(count):
+                txn.put(
+                    format(frame_idx, "012d").encode(),
+                    msgpack.packb(
+                        _make_lmdb_frame(natoms=natoms, seed=frame_idx),
+                        use_bin_type=True,
+                    ),
+                )
+                frame_idx += 1
+    env.close()
+    return path
+
+
 def _make_single_task_config() -> dict:
     return {
         "model": deepcopy(model_se_e2_a),
@@ -192,6 +280,56 @@ def test_full_validator_restores_top_k_checkpoints(self) -> None:
                 ["best.ckpt-10.t-2.pt", "best.ckpt-20.t-1.pt"],
             )
 
+    def test_full_validator_lmdb_full_validation_iterates_nloc_groups(self) -> None:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            lmdb_path = _create_mixed_nloc_lmdb(f"{tmpdir}/mixed.lmdb")
+            validation_data = LmdbDataset(
+                lmdb_path,
+                type_map=["O", "H"],
+                batch_size=2,
+            )
+            validator = FullValidator(
+                validating_params={
+                    "full_validation": True,
+                    "validation_freq": 1,
+                    "save_best": False,
+                    "max_best_ckpt": 1,
+                    "validation_metric": "E:MAE",
+                    "full_val_file": "val.log",
+                    "full_val_start": 0.0,
+                },
+                validation_data=validation_data,
+                model=_DummyModel(),
+                train_infos={},
+                num_steps=10,
+                rank=0,
+                zero_stage=0,
+                restart_training=False,
+            )
+            observed_natoms = []
+
+            def fake_evaluate_system(data_system):
+                test_data = data_system.get_test()
+                natoms = int(test_data["type"].shape[1])
+                nframes = int(test_data["coord"].shape[0])
+                observed_natoms.append(natoms)
+                return {
+                    "mae_e_per_atom": (float(natoms), nframes),
+                    "rmse_e_per_atom": (float(natoms), nframes),
+                }
+
+            with patch.object(
+                validator,
+                "_evaluate_system",
+                side_effect=fake_evaluate_system,
+            ) as evaluate_system:
+                metrics = validator.evaluate_all_systems()
+
+        self.assertEqual(observed_natoms, [6, 9, 12])
+        self.assertEqual(evaluate_system.call_count, 3)
+        self.assertAlmostEqual(metrics["mae_e_per_atom"], 8.4)
+        self.assertAlmostEqual(metrics["rmse_e_per_atom"], np.sqrt(75.6))
+
 
 class TestValidationArgcheck(unittest.TestCase):
     def test_normalize_rejects_missing_validation_data(self) -> None:
