diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index d4e4c65b23..1f3353a397 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -149,6 +149,18 @@ def _remap_keys(frame: dict[str, Any]) -> dict[str, Any]: return out +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 is_lmdb(systems: str) -> bool: """Check if systems points to an LMDB dataset.""" return systems.endswith(".lmdb") or Path(systems, "data.mdb").is_file() @@ -466,10 +478,16 @@ def __init__( def _compute_natoms_vec(self, atype: np.ndarray) -> np.ndarray: """Compute natoms_vec from a frame's atype array. + Negative virtual types are excluded from the per-type counts, matching + mixed-type NPY data handling. This function also excludes positive + indices outside the configured type map. The leading nloc entries still + include every atom slot. + Returns [nloc, nloc, count_type0, count_type1, ...] with length ntypes+2. """ nloc = len(atype) - counts = np.bincount(atype, minlength=self._ntypes)[: self._ntypes] + real_atype = atype[(atype >= 0) & (atype < self._ntypes)] + counts = np.bincount(real_atype, minlength=self._ntypes) vec = np.empty(self._ntypes + 2, dtype=np.int64) vec[0] = nloc vec[1] = nloc @@ -579,7 +597,7 @@ def __getitem__(self, index: int) -> dict[str, Any]: frame["atype"] = frame["atype"].reshape(-1).astype(np.int64) # Remap atom types from LMDB's type_map to model's type_map if self._type_remap is not None: - frame["atype"] = self._type_remap[frame["atype"]].astype(np.int64) + frame["atype"] = _remap_atom_types(frame["atype"], self._type_remap) if "virial" in frame and isinstance(frame["virial"], np.ndarray): frame["virial"] = ( frame["virial"].reshape(9).astype(self._resolve_dtype("virial")) @@ -1459,9 +1477,9 @@ def __init__( and "atype" in frame and isinstance(frame["atype"], np.ndarray) ): - frame["atype"] = self._type_remap[ - frame["atype"].reshape(-1) - ].astype(np.int64) + frame["atype"] = _remap_atom_types( + frame["atype"].reshape(-1), self._type_remap + ) self._frames.append(frame) # Shuffle if requested diff --git a/deepmd/dpmodel/utils/stat.py b/deepmd/dpmodel/utils/stat.py index ca00f6c064..4e418cae8a 100644 --- a/deepmd/dpmodel/utils/stat.py +++ b/deepmd/dpmodel/utils/stat.py @@ -43,7 +43,8 @@ def collect_observed_types(sampled: list[dict], type_map: list[str]) -> list[str Returns ------- list[str] - Sorted list of observed element symbols. + Sorted list of observed element symbols. Negative virtual atom types + and indices outside ``type_map`` are ignored. """ from deepmd.utils.econf_embd import ( sort_element_type, @@ -54,7 +55,7 @@ def collect_observed_types(sampled: list[dict], type_map: list[str]) -> list[str atype = to_numpy_array(system["atype"]) # shape: [nframes, natoms] observed_indices.update(np.unique(atype).tolist()) observed_types = [ - type_map[i] for i in sorted(observed_indices) if i < len(type_map) + type_map[i] for i in sorted(observed_indices) if 0 <= i < len(type_map) ] return sort_element_type(observed_types) diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index 0838cca0ab..59db455a76 100644 --- a/source/tests/common/dpmodel/test_lmdb_data.py +++ b/source/tests/common/dpmodel/test_lmdb_data.py @@ -17,6 +17,7 @@ LmdbTestDataNlocView, SameNlocBatchSampler, _expand_indices_by_blocks, + _remap_atom_types, compute_block_targets, is_lmdb, make_neighbor_stat_data, @@ -165,6 +166,44 @@ def _create_lmdb_with_type_map( return path +def _create_lmdb_with_virtual_type(path: str) -> str: + """Create one LMDB frame containing a negative virtual atom type.""" + atype = np.array([0, -1, 1], dtype=np.int64) + frame = _make_frame(natoms=len(atype)) + frame["atom_types"] = { + "type": " None: result = collect_observed_types(sampled, type_map) self.assertEqual(result, ["O"]) + def test_virtual_type_ignored(self) -> None: + """Negative virtual types must not alias the end of the type map.""" + sampled = [ + {"atype": np.array([[0, -1, 1]])}, + ] + type_map = ["O", "H", "Au"] + result = collect_observed_types(sampled, type_map) + self.assertEqual(result, ["H", "O"]) + class TestObservedTypeStatFile(unittest.TestCase): """Test stat file save/load round-trip for observed_type (dpmodel)."""