diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py index 6c433ba2ab..588c8ea1ae 100644 --- a/deepmd/dpmodel/utils/__init__.py +++ b/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", diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index c2ce362a50..60ed5fcfb6 100644 --- 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) @@ -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, diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index ab4cbdd963..604857c837 100644 --- a/deepmd/entrypoints/test.py +++ b/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, diff --git a/deepmd/pt/train/validation.py b/deepmd/pt/train/validation.py index c38e3b584a..fa3d2b364c 100644 --- a/deepmd/pt/train/validation.py +++ b/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]]: diff --git a/deepmd/pt/utils/lmdb_dataset.py b/deepmd/pt/utils/lmdb_dataset.py index ac7b1f0e3d..d196ad522e 100644 --- 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 index ac096633c2..b676d755a9 100644 --- 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 index f2e024aed2..eb4e32b065 100644 --- 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": " 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: