diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index 0d7ea42a89..404795a0fa 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -657,16 +657,27 @@ def _scan_lmdb_path_in_worker( executor.shutdown(wait=True, cancel_futures=True) -def _remap_atom_types(atype: np.ndarray, type_remap: np.ndarray) -> np.ndarray: - """Remap real atom types while preserving negative virtual sentinels. - - Positive indices retain NumPy's normal bounds checking, so malformed LMDB - data cannot be silently reinterpreted as a different species. - """ - remapped_atype = atype.astype(np.int64, copy=True) - real_atom_mask = remapped_atype >= 0 - remapped_atype[real_atom_mask] = type_remap[remapped_atype[real_atom_mask]] - return remapped_atype +def _make_type_remap( + lmdb_type_map: list[str] | None, type_map: list[str] | None +) -> np.ndarray | None: + """Map declared source types to model indices, marking unsupported entries.""" + if ( + lmdb_type_map is None + or type_map is None + or list(lmdb_type_map) == list(type_map) + ): + return None + remap = np.array( + [type_map.index(name) if name in type_map else -1 for name in lmdb_type_map], + dtype=np.int32, + ) + log.info( + "Type remapping: LMDB %s -> model %s, remap=%s", + lmdb_type_map, + type_map, + remap, + ) + return remap @dataclass @@ -685,7 +696,8 @@ class LmdbDecodeConfig: natoms Fallback atom count for records without ``atom_types``. type_remap - Optional LMDB-type to model-type lookup table. + Optional LMDB-type to model-type lookup table. Negative entries denote + unsupported elements, rejected only when referenced by a real atom. data_requirements Registered data requirements keyed by field name. dataset @@ -699,6 +711,9 @@ class LmdbDecodeConfig: their inputs and derive their labels here, before the model runs, which is the only place that works for a backend whose loss never sees the model. ``None``, the default, leaves decoding unchanged. + lmdb_type_map + Source element symbols used to identify unsupported atom types in + frame-level diagnostics. """ ntypes: int @@ -708,6 +723,7 @@ class LmdbDecodeConfig: dataset: str = "" type_map: list[str] | None = None frame_transform: Callable[[dict[str, Any], int], dict[str, Any]] | None = None + lmdb_type_map: list[str] | None = None def _requirement_dtype(requirement: Any) -> np.dtype: @@ -1041,7 +1057,26 @@ def decode_lmdb_frame( if "atype" in frame and isinstance(frame["atype"], np.ndarray): frame["atype"] = frame["atype"].reshape(-1).astype(np.int64, copy=False) if config.type_remap is not None: - frame["atype"] = _remap_atom_types(frame["atype"], config.type_remap) + atype = frame["atype"] + real_atom_mask = atype >= 0 + real_types = atype[real_atom_mask] + # Negative input sentinels are not lookup indices. Positive indices + # retain NumPy's bounds checks for malformed source type indices. + mapped = config.type_remap[real_types] + unsupported = mapped < 0 + if np.any(unsupported): + indices = np.unique(real_types[unsupported]) + missing = ( + [config.lmdb_type_map[index] for index in indices] + if config.lmdb_type_map is not None + else indices.tolist() + ) + raise ValueError( + f"LMDB frame {original_key} in '{config.dataset}' uses atom " + f"types {missing} not found in model type_map {config.type_map}" + ) + frame["atype"] = atype.copy() + frame["atype"][real_atom_mask] = mapped if "virial" in frame and isinstance(frame["virial"], np.ndarray): frame["virial"] = ( frame["virial"] @@ -2055,6 +2090,8 @@ class LmdbDataReader: Reads LMDB frames and returns dicts of numpy arrays. Backend-specific Dataset classes (PyTorch, JAX, etc.) wrap this. + Element compatibility is checked on decoded real atoms, not on unused + entries in the dataset's element table. An LMDB typically holds frames of many different atom counts. The ``batch_size`` rule decides how those frames are grouped: @@ -2121,25 +2158,8 @@ def __init__( self._natoms = sum(self._natoms_per_type) self._ntypes = len(type_map) - # Build type remapping if LMDB's type_map differs from model's type_map - lmdb_type_map = meta.get("type_map") - self._lmdb_type_map = lmdb_type_map - self._type_remap: np.ndarray | None = None - if lmdb_type_map is not None and list(lmdb_type_map) != list(type_map): - # Build remap: lmdb_type_idx -> model_type_idx - remap = np.empty(len(lmdb_type_map), dtype=np.int32) - for i, name in enumerate(lmdb_type_map): - if name not in type_map: - raise ValueError( - f"Element '{name}' in LMDB type_map {lmdb_type_map} " - f"not found in model type_map {type_map}" - ) - remap[i] = type_map.index(name) - self._type_remap = remap - log.info( - f"Type remapping: LMDB {lmdb_type_map} -> model {type_map}, " - f"remap={remap}" - ) + self._lmdb_type_map = meta.get("type_map") + self._type_remap = _make_type_remap(self._lmdb_type_map, type_map) # The parent transaction serves synchronous reads. Decoder workers open # independent process-local environments and transactions. @@ -2303,6 +2323,7 @@ def __init__( data_requirements=self._data_requirements, dataset=self.lmdb_path, type_map=self._type_map, + lmdb_type_map=self._lmdb_type_map, ) # Which fields carry an atom axis follows from the registered # requirements, so this cache is invalidated when they change. @@ -2457,7 +2478,7 @@ def batch_layout( # workers a chunk of it is shipped to. return BatchLayout.over( self._frame_nlocs[np.asarray(indices, dtype=np.int64)].astype(np.int64), - self.per_atom_strides(), + self.per_atom_strides(indices[0]), ragged=self._ragged_batches if ragged is None else ragged, ) @@ -2490,12 +2511,18 @@ def use_ragged_batches(self, ragged: bool) -> None: """ self._ragged_batches = ragged and self.mixed_nloc - def per_atom_strides(self) -> dict[str, int]: + def per_atom_strides(self, index: int = 0) -> dict[str, int]: """Return the leading-axis entries per atom of each per-atom field. Every frame of one LMDB exposes the same fields, so the classification is a property of the dataset and its registered requirements rather - than of a batch, and is resolved from the first frame once. + than of a batch, and is resolved from a selected frame once. + + Parameters + ---------- + index : int, optional + Dataset index of the frame used to initialize the cached strides. + Batch consumers supply a member of their selected indices. Returns ------- @@ -2503,7 +2530,7 @@ def per_atom_strides(self) -> dict[str, int]: Entries per atom, keyed by field name. """ if self._per_atom_strides is None: - frame = self[0] + frame = self[index] self._per_atom_strides = per_atom_strides( frame, resolve_per_atom_keys(frame, self._decode_config) ) @@ -3889,6 +3916,8 @@ class LmdbTestData: For mixed-nloc datasets, frames are grouped by nloc. get_test(nloc=...) returns data for a specific group. + Only types used by decoded frames must be supported by the model; + unused element-table entries do not restrict the model's type map. """ def __init__( @@ -3912,28 +3941,8 @@ def __init__( len(self._type_map) if self._type_map else len(self._natoms_per_type) ) - # Build type remapping if LMDB's type_map differs from model's type_map - lmdb_type_map = meta.get("type_map") - self._lmdb_type_map = lmdb_type_map - self._type_remap: np.ndarray | None = None - if ( - lmdb_type_map is not None - and self._type_map - and list(lmdb_type_map) != list(self._type_map) - ): - remap = np.empty(len(lmdb_type_map), dtype=np.int32) - for i, name in enumerate(lmdb_type_map): - if name not in self._type_map: - raise ValueError( - f"Element '{name}' in LMDB type_map {lmdb_type_map} " - f"not found in model type_map {self._type_map}" - ) - remap[i] = self._type_map.index(name) - self._type_remap = remap - log.info( - f"LmdbTestData type remapping: LMDB {lmdb_type_map} -> " - f"model {self._type_map}, remap={remap}" - ) + self._lmdb_type_map = meta.get("type_map") + self._type_remap = _make_type_remap(self._lmdb_type_map, self._type_map or None) # Select the frames to test, without reading any of them. Atom counts # come from the metadata, so grouping, shuffling and truncation are @@ -3954,6 +3963,7 @@ def __init__( data_requirements=self._requirements, dataset=self.lmdb_path, type_map=self._type_map, + lmdb_type_map=self._lmdb_type_map, ) # Detect PBC from the first retained frame. @@ -3961,11 +3971,13 @@ def __init__( first = next( (indices[0] for indices in self._nloc_groups.values() if len(indices)), None ) - probe = self._read_frames([first]) if first is not None else [] - if probe: - box = probe[0].get("box") - if not isinstance(box, np.ndarray) or np.allclose(box, 0.0): - self.pbc = False + if first is not None: + with self._env.begin() as transaction: + raw = transaction.get(format(first, self._frame_fmt).encode()) + if raw is not None: + frame = _remap_keys(msgpack.unpackb(raw, raw=False)) + box = _decode_value(frame.get("box"), copy_arrays=False) + self.pbc = isinstance(box, np.ndarray) and not np.allclose(box, 0.0) self.mixed_type = True diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index 304a4ba506..30a49f515d 100644 --- a/source/tests/common/dpmodel/test_lmdb_data.py +++ b/source/tests/common/dpmodel/test_lmdb_data.py @@ -48,7 +48,6 @@ _chop_mixed_nloc, _expand_indices_by_blocks, _merge_lmdb_chunks, - _remap_atom_types, collate_lmdb_frames, collect_lmdb_sampling_groups, compute_block_targets, @@ -1888,6 +1887,124 @@ def setUpClass(cls): def tearDownClass(cls): cls._tmpdir.cleanup() + def _subset_type_map_lmdb( + self, name: str, second_type: int = 2, *, first_type: int = 2 + ) -> str: + """Build two frames using a sparse source type map and a virtual atom.""" + path = _create_lmdb_with_virtual_type(f"{self._tmpdir.name}/{name}.lmdb") + with lmdb.open(path) as environment, environment.begin(write=True) as txn: + meta = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) + meta.update( + nframes=2, + type_map=["O", "He", "H"], + frame_nlocs=[3, 3], + system_info={"natoms": [1, 0, 1]}, + ) + txn.put(b"__metadata__", msgpack.packb(meta, use_bin_type=True)) + frame = msgpack.unpackb(txn.get(b"000000000000"), raw=False) + for index, last_type in enumerate((first_type, second_type)): + frame["atom_types"] = { + "type": " None: + path = self._subset_type_map_lmdb("unused_elements") + reader = LmdbDataReader(path, ["H", "O"], batch_size=2) + self.addCleanup(reader.close) + for index in range(2): + np.testing.assert_array_equal(reader[index]["atype"], [1, -1, 0]) + np.testing.assert_array_equal(reader[index]["natoms"], [3, 3, 1, 1]) + batch = reader.decode_batch([0, 1]) + np.testing.assert_array_equal(batch["atype"], [[1, -1, 0], [1, -1, 0]]) + test_data = LmdbTestData(path, type_map=["H", "O"], shuffle_test=False) + np.testing.assert_array_equal(test_data.get_test()["type"], batch["atype"]) + original_types = LmdbTestData(path, shuffle_test=False) + np.testing.assert_array_equal( + original_types.get_test()["type"], [[0, -1, 2], [0, -1, 2]] + ) + + def test_unsupported_element_in_a_later_frame_is_rejected(self) -> None: + path = self._subset_type_map_lmdb("used_unsupported_element", second_type=1) + reader = LmdbDataReader(path, ["H", "O"], batch_size=2) + self.addCleanup(reader.close) + np.testing.assert_array_equal(reader[0]["atype"], [1, -1, 0]) + with self.assertRaisesRegex(ValueError, "frame 1.*He.*model type_map"): + reader.__getitem__(1) + with self.assertRaisesRegex(ValueError, "frame 1.*He.*model type_map"): + reader.decode_batch([0, 1]) + iterator = LmdbBatchIterator( + reader, LmdbBatchSampler(reader, shuffle=False), num_workers=2 + ) + self.addCleanup(iterator.close) + with self.assertRaisesRegex(ValueError, "frame 1.*He.*model type_map"): + next(iterator) + test_data = LmdbTestData(path, type_map=["H", "O"], shuffle_test=False) + with self.assertRaisesRegex(ValueError, "frame 1.*He.*model type_map"): + test_data.get_test() + selected = LmdbTestData( + path, type_map=["H", "O"], shuffle_test=False, max_frames=1 + ) + np.testing.assert_array_equal(selected.get_test()["type"], [[1, -1, 0]]) + + def test_batch_selection_ignores_unsupported_first_frame(self) -> None: + """Layout discovery reads a frame belonging to the requested batch.""" + path = self._subset_type_map_lmdb("unsupported_first_batch", first_type=1) + for num_workers in (0, 2): + with self.subTest(num_workers=num_workers): + reader = LmdbDataReader(path, ["H", "O"], batch_size=2) + self.addCleanup(reader.close) + iterator = LmdbBatchIterator(reader, [[1, 1]], num_workers=num_workers) + self.addCleanup(iterator.close) + np.testing.assert_array_equal( + next(iterator)["atype"], [[1, -1, 0], [1, -1, 0]] + ) + with self.assertRaisesRegex(ValueError, "frame 0.*He.*model type_map"): + reader.decode_batch([0]) + + def test_testdata_probe_ignores_unsupported_first_frame(self) -> None: + """PBC discovery is independent of element validation and frame selection.""" + for box_kind in ("cells", "box", "zero", "missing"): + with self.subTest(box_kind=box_kind): + path = self._subset_type_map_lmdb( + f"unsupported_first_test_{box_kind}", first_type=1 + ) + with ( + lmdb.open(path) as environment, + environment.begin(write=True) as txn, + ): + frame = msgpack.unpackb(txn.get(b"000000000000"), raw=False) + if box_kind == "box": + frame["box"] = frame.pop("cells") + elif box_kind == "zero": + frame["cells"]["data"] = np.zeros((3, 3)).tobytes() + elif box_kind == "missing": + del frame["cells"] + txn.put(b"000000000000", msgpack.packb(frame, use_bin_type=True)) + test_data = LmdbTestData(path, type_map=["H", "O"], shuffle_test=False) + self.assertEqual(test_data.pbc, box_kind in ("cells", "box")) + np.testing.assert_array_equal( + test_data.get_test_by_indices([1])["type"], [[1, -1, 0]] + ) + with self.assertRaisesRegex(ValueError, "frame 0.*He.*model type_map"): + test_data.get_test_by_indices([0]) + + def test_subset_remapping_reaches_parallel_decoder_workers(self) -> None: + path = self._subset_type_map_lmdb("parallel_subset") + reader = LmdbDataReader(path, ["H", "O"], batch_size=2) + self.addCleanup(reader.close) + iterator = LmdbBatchIterator( + reader, LmdbBatchSampler(reader, shuffle=False), num_workers=2 + ) + self.addCleanup(iterator.close) + np.testing.assert_array_equal(next(iterator)["atype"], [[1, -1, 0], [1, -1, 0]]) + def test_reader_no_remap_when_match(self): reader = LmdbDataReader(self._lmdb_path, ["O", "H"]) self.assertIsNone(reader._type_remap) @@ -1914,8 +2031,10 @@ def test_reader_natoms_vec_after_remap(self): self.assertEqual(natoms[3], 2) # O count def test_reader_missing_element_raises(self): + reader = LmdbDataReader(self._lmdb_path, ["O"]) + self.addCleanup(reader.close) with self.assertRaises(ValueError): - LmdbDataReader(self._lmdb_path, ["O"]) + reader.__getitem__(0) def test_reader_no_type_map_in_metadata(self): tmpdir = tempfile.TemporaryDirectory() @@ -1945,8 +2064,9 @@ def test_testdata_remap_superset(self): self.assertIsNotNone(td._type_remap) def test_testdata_missing_element_raises(self): + td = LmdbTestData(self._lmdb_path, type_map=["O"], shuffle_test=False) with self.assertRaises(ValueError): - LmdbTestData(self._lmdb_path, type_map=["O"], shuffle_test=False) + td.get_test() def test_testdata_no_type_map_in_metadata(self): tmpdir = tempfile.TemporaryDirectory() @@ -1971,8 +2091,11 @@ def test_virtual_type_preserved_during_remap(self): def test_positive_out_of_range_type_still_raises(self): """A malformed real type must not be mistaken for a virtual sentinel.""" + path = self._subset_type_map_lmdb("invalid_type_index", second_type=3) + reader = LmdbDataReader(path, ["H", "O"]) + self.addCleanup(reader.close) with self.assertRaises(IndexError): - _remap_atom_types(np.array([0, 2]), np.array([1, 0])) + reader.__getitem__(1) # ============================================================