Skip to content

fix(data): allow unused elements in LMDB type maps #6100

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
OutisLi merged 2 commits into from
Oct 23, 2026
+201 −66
Merged

fix(data): allow unused elements in LMDB type maps #6100

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f3a9ff4
fix(data): allow unused elements in LMDB type maps
OutisLi Oct 22, 2026
6dad114
fix(data): avoid decoding unselected LMDB frames
OutisLi Oct 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
136 changes: 74 additions & 62 deletions deepmd/dpmodel/utils/lmdb_data.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -708,6 +723,7 @@ class LmdbDecodeConfig:
dataset: str = "<unknown LMDB>"
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:
Expand Down Expand Up @@ -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"]
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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,
)

Expand Down Expand Up @@ -2490,20 +2511,26 @@ 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
-------
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)
)
Expand Down Expand Up @@ -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__(
Expand All @@ -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
Expand All @@ -3954,18 +3963,21 @@ 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.
self.pbc = True
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

Expand Down
Loading
Loading