From 543cb00e88136958fb1061c1f7a20d05e211d81a Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 12 Jul 2026 19:27:29 +0800 Subject: [PATCH 1/6] feat(lmdb): align backend with DeePMD-kit Replace the legacy LMDB schema with a safe, mixed-type format that preserves system partitions and interoperates with DeePMD-kit training datasets. --- docs/systems/index.rst | 1 + docs/systems/lmdb.md | 99 ++ dpdata/format.py | 10 +- dpdata/formats/lmdb/__init__.py | 4 +- dpdata/formats/lmdb/format.py | 2239 ++++++++++++++++++++++++++++--- pyproject.toml | 2 +- tests/test_lmdb.py | 1459 +++++++++++++++++++- tests/test_lmdb_custom_dtype.py | 366 +++-- 8 files changed, 3831 insertions(+), 349 deletions(-) create mode 100644 docs/systems/lmdb.md diff --git a/docs/systems/index.rst b/docs/systems/index.rst index dfcbe698e..ab27479a5 100644 --- a/docs/systems/index.rst +++ b/docs/systems/index.rst @@ -9,3 +9,4 @@ Systems multi bond_order_system mixed + lmdb diff --git a/docs/systems/lmdb.md b/docs/systems/lmdb.md new file mode 100644 index 000000000..3a9399970 --- /dev/null +++ b/docs/systems/lmdb.md @@ -0,0 +1,99 @@ +# LMDB Format + +The format `lmdb` stores the frames of one or more systems in a single [LMDB](http://www.lmdb.tech/doc/) database, and can be loaded or dumped through {class}`dpdata.System`, {class}`dpdata.LabeledSystem`, and {class}`dpdata.MultiSystems`. The on-disk layout coincides with the LMDB datasets read by the DeePMD-kit data loader. Core fields and registered additional fields are therefore available to DeePMD-kit under their `deepmd_name`. + +In contrast to the directory-based `deepmd/npy` format, every frame is stored as an independent record indexed by a global frame number. Frames within one database may therefore differ in the number of atoms and in chemical composition, which is suited to data sets in which the number of frames per system is small. + +## Layout + +A database contains one metadata record together with one record per frame. + +The metadata record is stored under the key `__metadata__` and holds the following entries. + +| key | description | +| ------------------ | ---------------------------------------------------- | +| `nframes` | total number of frames | +| `frame_idx_fmt` | format string of the integer frame key (default `012d`) | +| `type_map` | the global element table | +| `frame_nlocs` | the number of atoms of each frame | +| `frame_system_ids` | the index of the source system of each frame | +| `dp_data_shapes` | optional symbolic shapes of additional dpdata fields | +| `dp_data_names` | optional mapping from DeePMD-kit field names to dpdata field names | + +Each frame is stored under its zero-padded global index (`000000000000`, `000000000001`, ...). The per-frame arrays `coords`, `cells`, `energies`, `forces`, `virials`, and `atom_types` are stored together with their data type and shape; `atom_types` records global indices into `type_map`, and `atom_numbs` is a list of per-element counts over `type_map`. The numerical dtype of each source array is retained, except that `atom_types` is stored as `int32`. + +## Dump data + +A single system or a collection of systems is written to one database. + +```python +import dpdata + +dpdata.LabeledSystem("OUTCAR", fmt="vasp/outcar").to("lmdb", "data.lmdb") + +dpdata.MultiSystems(*systems).to("lmdb", "data.lmdb") +``` + +The element table recorded in `type_map` defaults to the union of the elements present in the data. An explicit table may be supplied through the `type_map` argument, for example the full periodic table. + +```python +from dpdata.periodic_table import ELEMENTS + +dpdata.MultiSystems(*systems).to("lmdb", "data.lmdb", type_map=list(ELEMENTS)) +``` + +Frames are committed in batches of 1,000 by default. The `write_batch_size` argument changes the transaction size. If a transaction exceeds `map_size`, the map is enlarged and the same encoded batch is retried before frame counters are advanced. The destination must not exist unless `overwrite=True` is supplied. + +Data are first written to a temporary sibling database. After closing it, dpdata reopens the staged database and validates the metadata, entry count, and continuous frame-key sequence. A new destination is published by one directory rename. For `overwrite=True`, the staged `data.mdb` replaces the existing `data.mdb` in one file rename, while the LMDB runtime lock file remains in place. A failed conversion or validation leaves the previous destination unchanged. All dpdata readers of the destination must be closed before replacement; readers owned by other libraries must likewise be closed because their environments cannot be coordinated through dpdata's process-local cache. + +## Preserving the system partition + +The `frame_system_ids` entry of the metadata records, for every frame, the index of the source system it belongs to. DeePMD-kit uses this partition for system-wise sampling, for example through `prob_sys_size`. + +When a {class}`dpdata.MultiSystems` is dumped through `to("lmdb", ...)`, its frames are first grouped by chemical formula, and systems that share a formula are merged into a single entry. The resulting `frame_system_ids` therefore reflect the formula grouping rather than the original sources, and the number of systems may be smaller than the number of inputs. + +When the original partition must be retained, the function {func}`dpdata.formats.lmdb.dump_systems` writes an ordered sequence of systems without formula merging; each input becomes one system, numbered in iteration order. + +```python +import dpdata +from dpdata.formats.lmdb import dump_systems + +systems = [dpdata.LabeledSystem(d, fmt="deepmd/npy") for d in directories] +dump_systems(systems, "data.lmdb", type_map=["H", "C", "N", "O"]) +``` + +If `type_map` is supplied, the systems may be provided as a generator, so that the sequence need not be held in memory. + +```python +def load(): + for d in directories: + yield dpdata.LabeledSystem(d, fmt="deepmd/npy") + +dump_systems(load(), "data.lmdb", type_map=["H", "C", "N", "O"]) +``` + +Reading an LMDB through dpdata subsequently groups frames by composition and does not reconstruct the original `frame_system_ids` partition. The partition remains available to the DeePMD-kit LMDB data loader. + +## Load data + +```python +import dpdata + +ms = dpdata.MultiSystems.from_file("data.lmdb", fmt="lmdb") +``` + +Frames are grouped by composition, and each composition is returned as one system. Atom order is canonicalized by a stable sort on the global atom type; coordinates and all registered atomic fields are permuted consistently. Frames of one composition must have identical field sets and a consistent periodic-boundary condition. By default (`mixed_type=False`) the element table of each resulting system is restricted to the elements that the system contains. When `mixed_type=True`, every system retains the complete element set from the database. Loading through {class}`dpdata.MultiSystems` may normalize the order of `atom_names`, so callers should use element names rather than assume that stored numerical indices are retained. + +The general {class}`dpdata.System` constructor applies its `type_map` argument after format parsing. Consequently, a direct single-system load with an explicit `type_map` retains that complete requested table even when `mixed_type=False`; this is standard `System` behavior. The compact/full distinction above describes the dictionaries yielded to {class}`dpdata.MultiSystems`. + +```python +ms = dpdata.MultiSystems.from_file("data.lmdb", fmt="lmdb", mixed_type=True) +``` + +A database that holds a single composition may also be read into a {class}`dpdata.LabeledSystem`. If the database contains several compositions, only the first can be represented by a single system and a warning is issued. + +```python +ls = dpdata.LabeledSystem("data.lmdb", fmt="lmdb") +``` + +The reader loads all selected frames into memory. The default `max_frames=100000` guard rejects larger data sets before decoding; set it to `None` only when sufficient memory is available. Large training data sets should normally be consumed directly by the DeePMD-kit data loader. diff --git a/dpdata/format.py b/dpdata/format.py index ade83c21c..5d3b6b139 100644 --- a/dpdata/format.py +++ b/dpdata/format.py @@ -2,7 +2,9 @@ from __future__ import annotations +import collections.abc # noqa: TC003 - required by typing.get_type_hints import os +import typing from abc import ABC from .plugin import Plugin @@ -291,7 +293,9 @@ class MultiModes: MultiMode = MultiModes.NotImplemented - def from_multi_systems(self, directory, **kwargs): + def from_multi_systems( + self, directory, **kwargs + ) -> collections.abc.Iterable[typing.Any]: """Implement MultiSystems.from that converts from this format to MultiSystems. By default, this method follows MultiMode to implement the conversion. @@ -318,7 +322,9 @@ def from_multi_systems(self, directory, **kwargs): f"{self.__class__.__name__} doesn't support MultiSystems.from" ) - def to_multi_systems(self, formulas, directory, **kwargs): + def to_multi_systems( + self, formulas, directory, **kwargs + ) -> collections.abc.Iterable[typing.Any]: """Implement MultiSystems.to that converts from MultiSystems to this format. By default, this method follows MultiMode to implement the conversion. diff --git a/dpdata/formats/lmdb/__init__.py b/dpdata/formats/lmdb/__init__.py index 53a3e8f0e..20e6d265c 100644 --- a/dpdata/formats/lmdb/__init__.py +++ b/dpdata/formats/lmdb/__init__.py @@ -1,5 +1,5 @@ from __future__ import annotations -from .format import LMDBFormat +from .format import LMDBFormat, dump_systems -__all__ = ["LMDBFormat"] +__all__ = ["LMDBFormat", "dump_systems"] diff --git a/dpdata/formats/lmdb/format.py b/dpdata/formats/lmdb/format.py index 9b518be6b..fe109b6cf 100644 --- a/dpdata/formats/lmdb/format.py +++ b/dpdata/formats/lmdb/format.py @@ -1,15 +1,263 @@ +"""LMDB format that is fully interoperable with DeePMD-kit. + +The on-disk layout is a *flat* sequence of frames, identical to what +DeePMD-kit's ``LmdbDataReader`` consumes and what the community +``npy_to_lmdb`` / ``json_to_lmdb`` converters produce: + +- ``__metadata__`` (msgpack dict, string keys):: + + { + "nframes": int, + "frame_idx_fmt": "012d", + "type_map": [str, ...], # global element table + "frame_nlocs": [int, ...], # atoms per frame + "frame_system_ids": [int, ...], # source-system index per frame + } + +- one entry per frame, keyed by the zero-padded global frame index + (``"000000000000"`` by default), whose value is a msgpack dict with + string keys. Array values use the manual encoding + ``{"type": str(dtype), "shape": [...], "data": }`` and + ``atom_numbs`` is a plain list of per-type counts over ``type_map``. + + Core per-frame keys use the plural names consumed by DeePMD-kit + (``coords``, ``cells``, ``energies``, ``forces``, ``virials``, + ``atom_types``). Registered additional fields use + :attr:`dpdata.data_type.DataType.deepmd_name`. + +``atom_types`` is stored as ``int32`` global indices into ``type_map``. + +Notes +----- +This reader loads frames into memory, mirroring ``dpdata``'s in-memory +``System`` model. A configurable frame-count guard prevents accidental +decoding of very large databases; training datasets should normally be +consumed directly by DeePMD-kit's streaming dataloader. +""" + from __future__ import annotations +import itertools +import numbers import os +import shutil +import tempfile +import threading +import warnings +from collections import OrderedDict +from dataclasses import dataclass +from pathlib import Path +from typing import Any import lmdb import msgpack -import msgpack_numpy as m import numpy as np +from dpdata.data_type import Axis, DataType from dpdata.format import Format -m.patch() +DEFAULT_MAP_SIZE = 1024**4 # 1 TiB sparse, matches the reference converters +DEFAULT_FRAME_FMT = "012d" +DEFAULT_MAX_FRAMES = 100_000 +DEFAULT_WRITE_BATCH_SIZE = 1_000 +METADATA_KEY = b"__metadata__" + +_RESERVED_KEYS = frozenset( + { + "atom_numbs", + "atom_names", + "atom_types", + "orig", + "real_atom_types", + "real_atom_names", + "nopbc", + } +) +_CORE_DISK_NAMES = { + "cells": "cells", + "coords": "coords", + "energies": "energies", + "forces": "forces", + "virials": "virials", +} +_CORE_FRAME_KEYS = frozenset(_CORE_DISK_NAMES.values()) +_DEEPMD_CORE_REMAP = { + "coords": "coord", + "cells": "box", + "energies": "energy", + "forces": "force", + "atom_types": "atype", + "virials": "virial", +} +_FORBIDDEN_CUSTOM_DISK_NAMES = frozenset( + { + METADATA_KEY.decode(), + "atom_types", + "atom_numbs", + "atom_names", + "orig", + "type", + "natoms", + "real_natoms_vec", + "fid", + "sid", + *_CORE_FRAME_KEYS, + *_DEEPMD_CORE_REMAP.values(), + } +) + +_TOKEN_TO_AXIS = { + "nframes": Axis.NFRAMES, + "natoms": Axis.NATOMS, + "ntypes": Axis.NTYPES, + "nbonds": Axis.NBONDS, +} + + +@dataclass(frozen=True) +class _FieldSpec: + """Describe one dpdata field and its LMDB representation.""" + + dp_name: str + disk_name: str + shape: tuple[int | Axis, ...] + deepmd_name: str + + @property + def frame_axis(self) -> int | None: + """Return the frame axis in the in-memory array.""" + return ( + self.shape.index(Axis.NFRAMES) + if Axis.NFRAMES in self.shape + else None + ) + + @property + def atom_axes(self) -> tuple[int, ...]: + """Return atom axes in the in-memory array.""" + return tuple(i for i, dim in enumerate(self.shape) if dim is Axis.NATOMS) + + @property + def payload_atom_axes(self) -> tuple[int, ...]: + """Return atom axes after the frame axis has been removed.""" + frame_axis = self.frame_axis + if frame_axis is None: + return self.atom_axes + return tuple(axis - (axis > frame_axis) for axis in self.atom_axes) + + +@dataclass(frozen=True) +class _AtomSource: + """Validated atom-type data for one source system.""" + + names: tuple[str, ...] + types: np.ndarray + masks: np.ndarray | None + + +_PROTOCOL_DTYPES = ( + DataType( + "spins", + np.ndarray, + (Axis.NFRAMES, Axis.NATOMS, 3), + required=False, + deepmd_name="spin", + ), + DataType( + "force_mags", + np.ndarray, + (Axis.NFRAMES, Axis.NATOMS, 3), + required=False, + deepmd_name="force_mag", + ), + DataType( + "charges", + np.ndarray, + (Axis.NFRAMES, Axis.NATOMS), + required=False, + deepmd_name="charge", + ), + DataType( + "move", + np.ndarray, + (Axis.NFRAMES, Axis.NATOMS, 3), + required=False, + deepmd_name="move", + ), + DataType( + "hessian", + np.ndarray, + (Axis.NFRAMES, Axis.NATOMS, 3, Axis.NATOMS, 3), + required=False, + deepmd_name="hessian", + ), + DataType( + "fparam", + np.ndarray, + (Axis.NFRAMES, -1), + required=False, + deepmd_name="fparam", + ), + DataType( + "aparam", + np.ndarray, + (Axis.NFRAMES, Axis.NATOMS, -1), + required=False, + deepmd_name="aparam", + ), + DataType( + "atom_ener", + np.ndarray, + (Axis.NFRAMES, Axis.NATOMS), + required=False, + deepmd_name="atom_ener", + ), + DataType( + "charge_spin", + np.ndarray, + (Axis.NFRAMES, -1), + required=False, + deepmd_name="charge_spin", + ), +) +_PROTOCOL_FIELD_NAMES = frozenset( + { + *_CORE_DISK_NAMES, + "atom_pref", + *(dtype.name for dtype in _PROTOCOL_DTYPES), + } +) +_PROTOCOL_DISK_NAMES = { + dtype.deepmd_name: dtype.name for dtype in _PROTOCOL_DTYPES +} + +_READ_ENV_CACHE: dict[str, tuple[lmdb.Environment, int]] = {} +_READ_ENV_LOCK = threading.Lock() +_PUBLISHING_PATHS: set[str] = set() +_READ_ENV_PID = os.getpid() + + +def _shape_tokens(shape: tuple[int | Axis, ...]) -> list[int | str]: + """Serialize a (possibly symbolic) DataType shape to msgpack-able tokens.""" + tokens: list[int | str] = [] + for s in shape: + if isinstance(s, Axis): + tokens.append(s.value) + else: + tokens.append(int(s)) + return tokens + + +def _tokens_to_shape(tokens: list[int | str | bytes]) -> tuple[int | Axis, ...]: + """Inverse of :func:`_shape_tokens`.""" + out: list[int | Axis] = [] + for tok in tokens: + t = tok.decode() if isinstance(tok, bytes) else tok + if isinstance(t, str) and t in _TOKEN_TO_AXIS: + out.append(_TOKEN_TO_AXIS[t]) + else: + out.append(int(t)) + return tuple(out) class LMDBError(Exception): @@ -17,270 +265,1857 @@ class LMDBError(Exception): class LMDBMetadataError(LMDBError): - """Metadata not found in LMDB.""" + """Metadata not found or malformed in LMDB.""" class LMDBFrameError(LMDBError): """Frame data not found in LMDB.""" -class LMDBFormat(Format): +def _encode_array(arr: np.ndarray) -> dict[str, Any]: + """Encode a numpy array as ``{"type", "shape", "data"}`` (string keys). + + ``shape`` is taken from the original array (so 0-d scalars stay 0-d); + ``data`` is the C-ordered byte buffer. + """ + a = np.asarray(arr) + if a.dtype.hasobject or a.dtype.fields is not None: + raise LMDBError( + f"LMDB array dtype {a.dtype} is not a portable raw-byte dtype." + ) + try: + restored_dtype = np.dtype(str(a.dtype)) + except TypeError as exc: + raise LMDBError( + f"LMDB array dtype {a.dtype} cannot be reconstructed from its " + "serialized name." + ) from exc + if restored_dtype != a.dtype: + raise LMDBError( + f"LMDB array dtype {a.dtype} is not preserved by its serialized " + f"name {str(a.dtype)!r}." + ) + return { + "type": str(a.dtype), + "shape": list(a.shape), + "data": np.ascontiguousarray(a).tobytes(), + } + + +def _is_encoded_array(val) -> bool: + """Whether ``val`` is a manually encoded array dict (str or byte keys).""" + if not isinstance(val, dict): + return False + return ("type" in val and "data" in val) or (b"type" in val and b"data" in val) + + +def _decode_array(d: dict[str | bytes, Any]) -> np.ndarray: + """Reconstruct a numpy array from the manual encoding. + + Tolerates both string and byte keys so files written by any compliant + producer can be read back. """ - Class for handling the LMDB format, which stores atomic configurations in a - Lightning Memory-Mapped Database (LMDB). + type_key = "type" if "type" in d else b"type" + shape_key = "shape" if "shape" in d else b"shape" + data_key = "data" if "data" in d else b"data" + raw_type = d[type_key] + dtype = np.dtype(raw_type.decode() if isinstance(raw_type, bytes) else raw_type) + if dtype.hasobject or dtype.fields is not None: + raise ValueError(f"dtype {dtype} is not a portable raw-byte dtype") + buf = d[data_key] + if shape_key in d: + shape = tuple(d[shape_key]) + else: + shape = (len(buf) // dtype.itemsize,) + # frombuffer returns a read-only view; copy so downstream code can mutate. + return np.frombuffer(buf, dtype=dtype).reshape(shape).copy() + + +def _decode_frame(raw: bytes) -> dict[str, Any]: + """Decode a msgpack frame into a dict of numpy arrays / python scalars.""" + try: + frame = msgpack.unpackb(raw, raw=False) + except ( + msgpack.exceptions.ExtraData, + msgpack.exceptions.FormatError, + msgpack.exceptions.StackError, + ValueError, + ) as exc: + raise LMDBFrameError(f"Cannot decode frame: {exc}") from exc + if not isinstance(frame, dict): + raise LMDBFrameError("Each frame must contain a mapping.") + out = {} + for key, val in frame.items(): + if isinstance(key, bytes): + name = key.decode() + elif isinstance(key, str): + name = key + else: + raise LMDBFrameError("Frame keys must be strings or bytes.") + try: + out[name] = _decode_array(val) if _is_encoded_array(val) else val + except (TypeError, ValueError) as exc: + raise LMDBFrameError( + f"Cannot decode array field '{name}': {exc}" + ) from exc + return out + - This format is optimized for machine learning workflows where fast, random - access to a large number of frames is required. All frames from multiple - systems (with potentially different numbers of atoms) are stored in a - single LMDB database file. +def _meta_get(meta: dict[str | bytes, Any], key: str) -> Any: + """Read a metadata value, tolerating string or byte keys.""" + if key in meta: + return meta[key] + bkey = key.encode() + if bkey in meta: + return meta[bkey] + return None - Both single systems and multiple systems are supported via the standard - ``dpdata`` APIs. + +def _read_metadata(txn: lmdb.Transaction) -> dict[str | bytes, Any]: + raw = txn.get(METADATA_KEY) + if raw is None: + raise LMDBMetadataError("LMDB database does not contain __metadata__.") + try: + metadata = msgpack.unpackb(raw, raw=False) + except ( + msgpack.exceptions.ExtraData, + msgpack.exceptions.FormatError, + msgpack.exceptions.StackError, + ValueError, + ) as exc: + raise LMDBMetadataError(f"Cannot decode __metadata__: {exc}") from exc + if not isinstance(metadata, dict): + raise LMDBMetadataError("__metadata__ must contain a mapping.") + return metadata + + +def _packb(value: Any) -> bytes: + """Pack a value and enforce the byte-oriented LMDB contract.""" + packed = msgpack.packb(value, use_bin_type=True) + if not isinstance(packed, bytes): + raise TypeError("msgpack.packb did not return bytes.") + return packed + + +def _require_integer( + value: Any, label: str, *, minimum: int +) -> int: + """Return a strictly validated integer metadata value.""" + if isinstance(value, (bool, np.bool_)) or not isinstance( + value, numbers.Integral + ): + raise LMDBMetadataError(f"{label} must be an integer.") + result = int(value) + if result < minimum: + raise LMDBMetadataError(f"{label} must be at least {minimum}.") + return result + + +def _all_data_types() -> dict[str, DataType]: + """Return registered and protocol-level data types by dpdata name.""" + from dpdata.system import LabeledSystem, System + + dtypes = {dt.name: dt for dt in _PROTOCOL_DTYPES} + for dt in (*System.DTYPES, *LabeledSystem.DTYPES): + dtypes[dt.name] = dt + return dtypes + + +def _disk_name(dtype: DataType) -> str: + """Return the LMDB key associated with a dpdata data type.""" + return _CORE_DISK_NAMES.get(dtype.name, dtype.deepmd_name) + + +def _validate_field_namespace(dp_name: str, disk_name: str) -> None: + """Prevent custom fields from shadowing DeePMD structural data.""" + if dp_name in _CORE_DISK_NAMES: + expected = _CORE_DISK_NAMES[dp_name] + if disk_name != expected: + raise LMDBError( + f"Core field '{dp_name}' must use LMDB key '{expected}', not " + f"'{disk_name}'." + ) + return + protocol_owner = _PROTOCOL_DISK_NAMES.get(disk_name) + if protocol_owner is not None and protocol_owner != dp_name: + raise LMDBError( + f"Custom field '{dp_name}' cannot use protocol LMDB key " + f"'{disk_name}', which belongs to '{protocol_owner}'." + ) + if disk_name in _FORBIDDEN_CUSTOM_DISK_NAMES or disk_name.startswith("find_"): + raise LMDBError( + f"Custom field '{dp_name}' cannot use reserved LMDB key " + f"'{disk_name}'." + ) + + +def _field_spec( + dtype: DataType, *, validate_namespace: bool = True +) -> _FieldSpec: + """Build a field specification from a registered data type.""" + if dtype.shape is None: + raise LMDBError( + f"Data type '{dtype.name}' has no declared shape and cannot be " + "represented unambiguously in per-frame LMDB records." + ) + disk_name = _disk_name(dtype) + if validate_namespace: + _validate_field_namespace(dtype.name, disk_name) + return _FieldSpec( + dp_name=dtype.name, + disk_name=disk_name, + shape=tuple(dtype.shape), + deepmd_name=dtype.deepmd_name, + ) + + +def _field_registries( + *, prefer_protocol: bool = False +) -> tuple[ + dict[str, _FieldSpec], dict[str, _FieldSpec] +]: + """Return field specifications indexed by dpdata and disk names.""" + data_types = _all_data_types() + if prefer_protocol: + for dtype in _PROTOCOL_DTYPES: + data_types[dtype.name] = dtype + by_dp: dict[str, _FieldSpec] = {} + by_disk: dict[str, _FieldSpec] = {} + ambiguous_disk_names: set[str] = set() + for dtype in data_types.values(): + if dtype.name in _RESERVED_KEYS or dtype.shape is None: + continue + spec = _field_spec(dtype, validate_namespace=False) + by_dp[spec.dp_name] = spec + if spec.disk_name in ambiguous_disk_names: + continue + previous = by_disk.get(spec.disk_name) + if previous is not None and previous.dp_name != spec.dp_name: + del by_disk[spec.disk_name] + ambiguous_disk_names.add(spec.disk_name) + continue + by_disk[spec.disk_name] = spec + return by_dp, by_disk + + +def _validate_names(names: list[str] | tuple[str, ...], label: str) -> tuple[str, ...]: + """Validate and normalize an element table.""" + if isinstance(names, (str, bytes)): + raise LMDBError(f"{label} must be a sequence of element names.") + normalized = tuple(str(name) for name in names) + if not normalized: + raise LMDBError(f"{label} must contain at least one element.") + if len(set(normalized)) != len(normalized): + raise LMDBError(f"{label} contains duplicate element names: {normalized}.") + return normalized + + +def _validate_integer_types( + values: np.ndarray, + *, + name: str, + upper_bound: int, + allow_virtual: bool, +) -> np.ndarray: + """Validate an atom-type array without silently truncating values.""" + array = np.asarray(values) + if ( + np.issubdtype(array.dtype, np.bool_) + or not np.issubdtype(array.dtype, np.integer) + ): + raise LMDBError(f"{name} must use an integer dtype, got {array.dtype}.") + lower_bound = -1 if allow_virtual else 0 + if np.issubdtype(array.dtype, np.unsignedinteger): + out_of_range = np.any(array >= upper_bound) + else: + out_of_range = np.any(array < lower_bound) or np.any(array >= upper_bound) + if out_of_range: + constraint = ( + f"-1 or indices in [0, {upper_bound})" + if allow_virtual + else f"indices in [0, {upper_bound})" + ) + raise LMDBError(f"{name} must contain only {constraint}.") + return array.astype(np.int64, copy=False) + + +def _prepare_atom_source( + data: dict[str, Any], nframes: int, natoms: int +) -> _AtomSource: + """Validate standard or mixed atom types for one source system.""" + if "real_atom_types" in data or "real_atom_names" in data: + if "real_atom_types" not in data or "real_atom_names" not in data: + raise LMDBError( + "Mixed-type data requires both real_atom_types and real_atom_names." + ) + names = _validate_names(data["real_atom_names"], "real_atom_names") + real_types = _validate_integer_types( + np.asarray(data["real_atom_types"]), + name="real_atom_types", + upper_bound=len(names), + allow_virtual=True, + ) + if real_types.shape != (nframes, natoms): + raise LMDBError( + "real_atom_types must have shape " + f"({nframes}, {natoms}), got {real_types.shape}." + ) + masks = real_types >= 0 + if np.any(np.sum(masks, axis=1) == 0): + raise LMDBError("Every frame must contain at least one real atom.") + return _AtomSource(names, real_types, masks) + + names = _validate_names(data["atom_names"], "atom_names") + atom_types = _validate_integer_types( + np.asarray(data["atom_types"]), + name="atom_types", + upper_bound=len(names), + allow_virtual=False, + ) + if atom_types.shape != (natoms,): + raise LMDBError( + f"atom_types must have shape ({natoms},), got {atom_types.shape}." + ) + expected_counts = np.bincount(atom_types, minlength=len(names)).tolist() + if list(data["atom_numbs"]) != expected_counts: + raise LMDBError( + "atom_numbs is inconsistent with atom_types: " + f"expected {expected_counts}, got {data['atom_numbs']}." + ) + return _AtomSource(names, atom_types, None) + + +def _source_data(system: Any) -> dict[str, Any]: + """Return a raw data dictionary from a System-like object.""" + if isinstance(system, dict): + return system + if hasattr(system, "data") and isinstance(system.data, dict): + return system.data + raise TypeError("Each input must be a dpdata System or a data dictionary.") + + +def _source_active_names(data: dict[str, Any]) -> tuple[str, ...]: + """Return active elements in first-appearance order for type-map discovery.""" + coords = np.asarray(data["coords"]) + if coords.ndim != 3 or coords.shape[2] != 3: + raise LMDBError( + f"coords must have shape (nframes, natoms, 3), got {coords.shape}." + ) + source = _prepare_atom_source(data, coords.shape[0], coords.shape[1]) + types = source.types[source.types >= 0] + active_indices = set(int(value) for value in np.unique(types)) + return tuple( + name for index, name in enumerate(source.names) if index in active_indices + ) + + +def _remove_path(path: Path) -> None: + """Remove a file or directory if it exists.""" + if path.is_dir() and not path.is_symlink(): + shutil.rmtree(path) + elif path.exists() or path.is_symlink(): + path.unlink() + + +def _normalized_path(path: str | Path) -> str: + """Return one cache key for relative, absolute, and case aliases.""" + return os.path.normcase(os.path.realpath(os.fspath(path))) + + +def _reset_read_cache_after_fork() -> None: + """Discard inherited LMDB environments and synchronization state.""" + global _READ_ENV_CACHE, _READ_ENV_LOCK, _PUBLISHING_PATHS, _READ_ENV_PID + + for env, _ in _READ_ENV_CACHE.values(): + try: + env.close() + except lmdb.Error: + pass + _READ_ENV_CACHE = {} + _PUBLISHING_PATHS = set() + _READ_ENV_LOCK = threading.Lock() + _READ_ENV_PID = os.getpid() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reset_read_cache_after_fork) + + +def _open_read_env(path: str) -> tuple[str, lmdb.Environment]: + """Open or share a read-only LMDB environment within dpdata.""" + global _READ_ENV_PID + + if _READ_ENV_PID != os.getpid(): + _reset_read_cache_after_fork() + resolved = _normalized_path(path) + with _READ_ENV_LOCK: + if resolved in _PUBLISHING_PATHS: + raise LMDBError( + f"LMDB '{path}' is being replaced and cannot be opened." + ) + cached = _READ_ENV_CACHE.get(resolved) + if cached is not None: + env, count = cached + _READ_ENV_CACHE[resolved] = (env, count + 1) + return resolved, env + try: + env = lmdb.open( + path, + readonly=True, + lock=False, + subdir=True, + readahead=False, + meminit=False, + ) + except lmdb.Error as exc: + raise LMDBError( + f"Cannot open LMDB '{path}'. If another library has already opened " + "this path in the current process, close that reader before using " + "dpdata." + ) from exc + _READ_ENV_CACHE[resolved] = (env, 1) + return resolved, env + + +def _close_read_env(resolved: str) -> None: + """Release a cached read-only LMDB environment.""" + with _READ_ENV_LOCK: + env, count = _READ_ENV_CACHE[resolved] + if count == 1: + del _READ_ENV_CACHE[resolved] + env.close() + else: + _READ_ENV_CACHE[resolved] = (env, count - 1) + + +def _begin_publish(path: Path) -> str: + """Reserve a destination for publication after checking active readers.""" + resolved = _normalized_path(path) + with _READ_ENV_LOCK: + if resolved in _READ_ENV_CACHE: + raise LMDBError( + f"Cannot replace LMDB '{path}' while a dpdata reader is active." + ) + if resolved in _PUBLISHING_PATHS: + raise LMDBError(f"LMDB '{path}' is already being replaced.") + _PUBLISHING_PATHS.add(resolved) + return resolved + + +def _end_publish(resolved: str) -> None: + """Release a publication reservation.""" + with _READ_ENV_LOCK: + _PUBLISHING_PATHS.discard(resolved) + + +def _open_publish_guard(path: Path) -> lmdb.Environment: + """Hold the process-wide LMDB slot until atomic replacement completes.""" + try: + return lmdb.open( + str(path), + readonly=True, + lock=False, + subdir=True, + readahead=False, + meminit=False, + ) + except lmdb.Error as exc: + raise LMDBError( + f"Cannot replace LMDB '{path}' while another LMDB environment is " + "open in this process. Close DeePMD-kit and other readers first." + ) from exc + + +class _LMDBWriter: + """Write a complete LMDB to a temporary sibling and publish it safely.""" + + def __init__( + self, + directory: str, + *, + map_size: int, + frame_idx_fmt: str, + type_map: list[str] | None, + write_batch_size: int, + overwrite: bool, + ) -> None: + if write_batch_size <= 0: + raise ValueError("write_batch_size must be positive.") + normalized_type_map = ( + list(_validate_names(type_map, "type_map")) + if type_map is not None + else None + ) + self.directory = Path(directory).resolve() + if self.directory.exists() and not overwrite: + raise FileExistsError( + f"LMDB destination '{directory}' already exists; pass " + "overwrite=True to replace it." + ) + self.directory.parent.mkdir(parents=True, exist_ok=True) + self.temp_directory = Path( + tempfile.mkdtemp( + prefix=f".{self.directory.name}.tmp-", + dir=self.directory.parent, + ) + ) + self.overwrite = overwrite + self.frame_idx_fmt = frame_idx_fmt + self.type_map = normalized_type_map + self.write_batch_size = write_batch_size + self.frame_index = 0 + self.system_index = 0 + self.frame_nlocs: list[int] = [] + self.frame_system_ids: list[int] = [] + self.data_shapes: dict[str, list[int | str]] = {} + self.data_names: dict[str, str] = {} + self._closed = False + self._published = False + try: + self.env = lmdb.open( + str(self.temp_directory), map_size=map_size, subdir=True + ) + except BaseException: + _remove_path(self.temp_directory) + raise + + def _build_specs( + self, + data: dict[str, Any], + *, + nframes: int, + natoms: int, + has_virtual_atoms: bool, + ) -> list[_FieldSpec]: + """Build and validate field specifications for one system.""" + by_dp, _ = _field_registries() + all_dtypes = _all_data_types() + specs: list[_FieldSpec] = [] + used_disk_names: dict[str, str] = {} + for name, value in data.items(): + if name in _RESERVED_KEYS or not isinstance(value, np.ndarray): + continue + array = np.asarray(value) + registered_dtype = all_dtypes.get(name) + if registered_dtype is not None and registered_dtype.shape is None: + raise LMDBError( + f"Data type '{name}' has no declared shape and cannot be " + "represented unambiguously in per-frame LMDB records." + ) + registered = name in by_dp + if registered: + spec = by_dp[name] + else: + if array.ndim > 0 and array.shape[0] == nframes: + shape: tuple[int | Axis, ...] = ( + Axis.NFRAMES, + *array.shape[1:], + ) + else: + shape = tuple(array.shape) + spec = _FieldSpec(name, name, shape, name) + payload_shape = list(shape) + if Axis.NFRAMES in payload_shape: + payload_shape.remove(Axis.NFRAMES) + if has_virtual_atoms and natoms in payload_shape: + raise LMDBError( + f"Unregistered field '{name}' may contain an atom axis, " + "but its shape is ambiguous for mixed data with virtual " + "atoms. Register a DataType with Axis.NATOMS before writing." + ) + _validate_field_namespace(spec.dp_name, spec.disk_name) + previous = used_disk_names.get(spec.disk_name) + if previous is not None and previous != spec.dp_name: + raise LMDBError( + f"Fields '{previous}' and '{spec.dp_name}' both map to LMDB " + f"key '{spec.disk_name}'." + ) + used_disk_names[spec.disk_name] = spec.dp_name + self._validate_source_shape(spec, array, nframes, natoms) + specs.append(spec) + if spec.disk_name not in _CORE_FRAME_KEYS: + existing_shape = self.data_shapes.get(spec.disk_name) + tokens = _shape_tokens(spec.shape) + if existing_shape is not None and existing_shape != tokens: + raise LMDBError( + f"Field '{spec.dp_name}' has incompatible shapes across " + f"systems: {existing_shape} and {tokens}." + ) + self.data_shapes[spec.disk_name] = tokens + if spec.disk_name != spec.dp_name: + existing_name = self.data_names.get(spec.disk_name) + if existing_name is not None and existing_name != spec.dp_name: + raise LMDBError( + f"LMDB key '{spec.disk_name}' maps to conflicting dpdata " + f"names '{existing_name}' and '{spec.dp_name}'." + ) + self.data_names[spec.disk_name] = spec.dp_name + if "coords" not in used_disk_names: + raise LMDBError("coords is required for every system.") + return specs + + @staticmethod + def _validate_source_shape( + spec: _FieldSpec, array: np.ndarray, nframes: int, natoms: int + ) -> None: + """Validate an in-memory field against its symbolic shape.""" + if array.ndim != len(spec.shape): + raise LMDBError( + f"Field '{spec.dp_name}' has rank {array.ndim}, expected " + f"{len(spec.shape)} from {spec.shape}." + ) + for axis, (actual, expected) in enumerate(zip(array.shape, spec.shape)): + if expected is Axis.NFRAMES and actual != nframes: + raise LMDBError( + f"Field '{spec.dp_name}' axis {axis} must have {nframes} " + f"frames, got {actual}." + ) + if expected is Axis.NATOMS and actual != natoms: + raise LMDBError( + f"Field '{spec.dp_name}' axis {axis} must have {natoms} " + f"atoms, got {actual}." + ) + if isinstance(expected, int) and expected != -1 and actual != expected: + raise LMDBError( + f"Field '{spec.dp_name}' axis {axis} must have length " + f"{expected}, got {actual}." + ) + + @staticmethod + def _frame_payload( + array: np.ndarray, + spec: _FieldSpec, + frame_index: int, + atom_mask: np.ndarray | None, + ) -> np.ndarray: + """Extract one frame and remove virtual atoms on every atom axis.""" + frame_axis = spec.frame_axis + payload = ( + np.take(array, frame_index, axis=frame_axis) + if frame_axis is not None + else array + ) + if atom_mask is not None: + for axis in spec.payload_atom_axes: + payload = np.compress(atom_mask, payload, axis=axis) + return payload + + def _put_records(self, records: list[tuple[bytes, bytes]]) -> None: + """Commit one batch, growing the LMDB map and retrying when necessary.""" + while True: + try: + with self.env.begin(write=True) as txn: + for key, value in records: + txn.put(key, value) + return + except lmdb.MapFullError: + current_size = int(self.env.info()["map_size"]) + self.env.set_mapsize( + max(current_size * 2, current_size + (1 << 20)) + ) + + def write_system(self, data: dict[str, Any]) -> None: + """Append one validated source system.""" + coords = np.asarray(data["coords"]) + if coords.ndim != 3 or coords.shape[2] != 3: + raise LMDBError( + f"coords must have shape (nframes, natoms, 3), got {coords.shape}." + ) + nframes, natoms = coords.shape[:2] + if nframes == 0: + raise LMDBError("A source system must contain at least one frame.") + atom_source = _prepare_atom_source(data, nframes, natoms) + + if self.type_map is None: + self.type_map = list(atom_source.names) + type_map = _validate_names(self.type_map, "type_map") + name_to_global = {name: index for index, name in enumerate(type_map)} + active_indices = np.unique(atom_source.types[atom_source.types >= 0]) + active_names = {atom_source.names[int(index)] for index in active_indices} + missing_active = sorted(active_names - set(type_map)) + if missing_active: + raise LMDBError( + f"Elements {missing_active} are not present in type_map {type_map}." + ) + local_to_global = np.array( + [name_to_global.get(name, -1) for name in atom_source.names], + dtype=np.int64, + ) + has_virtual_atoms = ( + atom_source.masks is not None and not np.all(atom_source.masks) + ) + specs = self._build_specs( + data, + nframes=nframes, + natoms=natoms, + has_virtual_atoms=has_virtual_atoms, + ) + nopbc = bool(data.get("nopbc", False)) + fields = [(spec, np.asarray(data[spec.dp_name])) for spec in specs] + + for batch_start in range(0, nframes, self.write_batch_size): + batch_end = min(batch_start + self.write_batch_size, nframes) + records: list[tuple[bytes, bytes]] = [] + batch_nlocs: list[int] = [] + for offset, frame_i in enumerate(range(batch_start, batch_end)): + if atom_source.masks is None: + local_types = atom_source.types + atom_mask = None + else: + atom_mask = atom_source.masks[frame_i] + local_types = atom_source.types[frame_i, atom_mask] + global_types = local_to_global[local_types] + if np.any(global_types < 0): + missing_indices = np.unique(local_types[global_types < 0]) + missing_names = [ + atom_source.names[int(index)] for index in missing_indices + ] + raise LMDBError( + f"Elements {missing_names} are not present in type_map " + f"{list(type_map)}." + ) + global_types = global_types.astype(np.int32) + frame: dict[str, Any] = { + "atom_types": _encode_array(global_types) + } + for spec, array in fields: + if spec.dp_name == "cells" and nopbc: + continue + payload = self._frame_payload( + array, spec, frame_i, atom_mask + ) + frame[spec.disk_name] = _encode_array(payload) + counts = np.bincount( + global_types, minlength=len(type_map) + )[: len(type_map)] + frame["atom_numbs"] = [int(count) for count in counts] + key = format( + self.frame_index + offset, self.frame_idx_fmt + ).encode("ascii") + records.append((key, _packb(frame))) + batch_nlocs.append(int(global_types.size)) + self._put_records(records) + self.frame_nlocs.extend(batch_nlocs) + self.frame_system_ids.extend( + [self.system_index] * len(records) + ) + self.frame_index += len(records) + self.system_index += 1 + + def _metadata(self) -> dict[str, Any]: + """Return validated metadata for the completed database.""" + if self.system_index == 0 or self.frame_index == 0: + raise LMDBError("Cannot write an empty LMDB dataset.") + if self.type_map is None: + raise LMDBError("Cannot write LMDB metadata without a type_map.") + metadata: dict[str, Any] = { + "nframes": self.frame_index, + "frame_idx_fmt": self.frame_idx_fmt, + "type_map": self.type_map, + "frame_nlocs": self.frame_nlocs, + "frame_system_ids": self.frame_system_ids, + } + if self.data_shapes: + metadata["dp_data_shapes"] = self.data_shapes + if self.data_names: + metadata["dp_data_names"] = self.data_names + return metadata + + def finish(self) -> None: + """Commit metadata, close the temporary database, and publish it.""" + self._put_records([(METADATA_KEY, _packb(self._metadata()))]) + self.env.sync() + self.env.close() + self._closed = True + self._validate_staged_database() + self._publish() + self._published = True + + def _validate_staged_database(self) -> None: + """Reopen the staged database and verify its complete key sequence.""" + env = lmdb.open( + str(self.temp_directory), + readonly=True, + lock=False, + subdir=True, + readahead=False, + meminit=False, + ) + try: + with env.begin() as txn: + entries = int(txn.stat()["entries"]) + if entries != self.frame_index + 1: + raise LMDBError( + f"Staged LMDB has {entries} entries; expected " + f"{self.frame_index + 1}." + ) + metadata = _read_metadata(txn) + if _meta_get(metadata, "nframes") != self.frame_index: + raise LMDBError("Staged LMDB nframes does not match written frames.") + if len(_meta_get(metadata, "frame_nlocs") or []) != self.frame_index: + raise LMDBError( + "Staged LMDB frame_nlocs length is inconsistent." + ) + if ( + len(_meta_get(metadata, "frame_system_ids") or []) + != self.frame_index + ): + raise LMDBError( + "Staged LMDB frame_system_ids length is inconsistent." + ) + for index in range(self.frame_index): + key = format(index, self.frame_idx_fmt).encode("ascii") + if txn.get(key) is None: + raise LMDBError( + f"Staged LMDB is missing frame key {key!r}." + ) + finally: + env.close() + + def _publish(self) -> None: + """Publish by one atomic rename of a directory or its data file.""" + reservation = _begin_publish(self.directory) + try: + if not self.directory.exists(): + os.replace(self.temp_directory, self.directory) + return + if not self.overwrite: + raise FileExistsError( + f"LMDB destination '{self.directory}' already exists." + ) + target_data = self.directory / "data.mdb" + staged_data = self.temp_directory / "data.mdb" + if not self.directory.is_dir() or not target_data.is_file(): + raise LMDBError( + f"Existing destination '{self.directory}' is not an LMDB " + "directory containing data.mdb." + ) + guard = _open_publish_guard(self.directory) + try: + os.replace(staged_data, target_data) + finally: + guard.close() + _remove_path(self.temp_directory) + finally: + _end_publish(reservation) + + def abort(self) -> None: + """Close and discard an incomplete temporary database.""" + if not self._closed: + self.env.close() + self._closed = True + if not self._published: + _remove_path(self.temp_directory) + + +class LMDBFormat(Format): + """DeePMD-kit compatible LMDB format. + + A single flat LMDB stores all frames from one or many systems. The + same on-disk format is produced regardless of whether the source is a + standard or a mixed-type system, so the output is always readable by + DeePMD-kit's ``LmdbDataReader``. + + The ``mixed_type`` keyword controls only how frames are mapped back to + ``dpdata`` objects on read (see :meth:`from_multi_systems`). Examples -------- - **Saving a single LabeledSystem** + Write a single labeled system:: - >>> import dpdata - >>> system = dpdata.LabeledSystem("path/to/input.vasp", fmt="vasp/outcar") - >>> system.to("lmdb", "my_single_system.lmdb") + >>> import dpdata + >>> ls = dpdata.LabeledSystem("OUTCAR", fmt="vasp/outcar") + >>> ls.to("lmdb", "data.lmdb") - **Loading a single LabeledSystem** + Write many systems into one LMDB, forcing a global type map:: - >>> loaded_system = dpdata.LabeledSystem("my_single_system.lmdb", fmt="lmdb") + >>> ms = dpdata.MultiSystems(s1, s2, s3) + >>> ms.to("lmdb", "data.lmdb", type_map=["H", "C", "N", "O"]) - **Saving multiple systems to a single LMDB database** + Read back as standard (per-composition) systems:: - >>> import dpdata - >>> system_1 = dpdata.LabeledSystem("path/to/system1/OUTCAR", fmt="vasp/outcar") - >>> system_2 = dpdata.LabeledSystem("path/to/system2/OUTCAR", fmt="vasp/outcar") - >>> multi_systems_obj = dpdata.MultiSystems(system_1, system_2) - >>> multi_systems_obj.to("lmdb", "my_multi_system_db.lmdb") + >>> ms = dpdata.MultiSystems.from_file("data.lmdb", fmt="lmdb") - **Loading multiple systems from a single LMDB database** + Read back keeping the full global type map on every system:: - >>> import dpdata - >>> loaded_multi_systems = dpdata.MultiSystems.from_file("my_multi_system_db.lmdb", fmt="lmdb") + >>> ms = dpdata.MultiSystems.from_file( + ... "data.lmdb", fmt="lmdb", mixed_type=True + ... ) + + Note that loading through :class:`dpdata.MultiSystems` normalises the + ``atom_names`` order (the element set is kept, but reordered); a direct + single-system load preserves the stored order. """ + # ------------------------------------------------------------------ # + # Writing + # ------------------------------------------------------------------ # def to_multi_systems( - self, formulas, directory, map_size=1000000000, frame_idx_fmt="012d", **kwargs + self, + formulas, + directory, + map_size: int = DEFAULT_MAP_SIZE, + frame_idx_fmt: str = DEFAULT_FRAME_FMT, + type_map: list[str] | None = None, + write_batch_size: int = DEFAULT_WRITE_BATCH_SIZE, + overwrite: bool = False, + **kwargs, ): - """Implement MultiSystems.to for LMDB format. + """Write multiple dpdata systems to one LMDB. Parameters ---------- formulas : list[str] - list of formulas + One handle per system (the value is not used on disk). directory : str - directory of system + Output LMDB directory. map_size : int, optional - Maximum size of the LMDB database in bytes. Default is 1GB. + Maximum LMDB size in bytes. Default is 1 TiB (sparse). frame_idx_fmt : str, optional - The format string used to encode the frame index as a key. Default is "012d". + Format used for the per-frame integer key. Default ``"012d"``. + type_map : list[str], optional + Global element table. If ``None``, the element list of the + first system written is used (for a ``MultiSystems`` this is + the union of all systems' elements). + write_batch_size : int, optional + Number of frames committed per LMDB write transaction. + overwrite : bool, optional + Whether to replace an existing destination after the new database + has been written and validated. The default is ``False``. **kwargs : dict other parameters Yields ------ tuple - (self, formula) to be used by to_system + ``(self, formula)`` handle consumed by :meth:`to_system`. """ - self._frame_idx_fmt = frame_idx_fmt - self._global_frame_idx = 0 - self._system_info = [] - os.makedirs(directory, exist_ok=True) - with lmdb.open(directory, map_size=map_size) as env: - with env.begin(write=True) as txn: - self._txn = txn - for ff in formulas: - yield (self, ff) - # Finalize metadata - metadata = { - "nframes": self._global_frame_idx, - "system_info": self._system_info, - "frame_idx_fmt": self._frame_idx_fmt, - } - txn.put(b"__metadata__", msgpack.packb(metadata, use_bin_type=True)) - self._txn = None - - def _dump_to_txn(self, data, txn, formula, dtypes): - from dpdata.data_type import Axis - - nframes = data["coords"].shape[0] - - # Identify symbolic shapes and frame-dependent keys - data_shapes = {} - frame_dependent_keys = [] - for dt in dtypes: - if dt.name in data: - if dt.shape is not None: - data_shapes[dt.name] = [ - s.value if isinstance(s, Axis) else s for s in dt.shape - ] - if Axis.NFRAMES in dt.shape: - frame_dependent_keys.append(dt.name) - else: - data_shapes[dt.name] = None - - # Record system info - # natoms needs to be extracted from data - if "atom_numbs" in data: - natoms_list = data["atom_numbs"] - else: - # Fallback for systems without atom_numbs (should not happen in valid dpdata systems) - natoms_list = [] - - self._system_info.append( - { - "formula": formula, - "natoms": natoms_list, - "nframes": nframes, - "start_idx": self._global_frame_idx, - "data_shapes": data_shapes, - "frame_dependent_keys": frame_dependent_keys, - } + writer = _LMDBWriter( + directory, + map_size=map_size, + frame_idx_fmt=frame_idx_fmt, + type_map=type_map, + write_batch_size=write_batch_size, + overwrite=overwrite, ) + self._writer = writer + completed = False + try: + for ff in formulas: + yield (self, ff) + writer.finish() + completed = True + finally: + if not completed: + writer.abort() + self._writer = None - for i in range(nframes): - frame_data = {} - for key, val in data.items(): - if key in frame_dependent_keys: - frame_data[key] = val[i] - else: - frame_data[key] = val - - key = f"{self._global_frame_idx:{self._frame_idx_fmt}}".encode("ascii") - value = msgpack.packb(frame_data, use_bin_type=True) - txn.put(key, value) - self._global_frame_idx += 1 + def dump_systems( + self, + systems, + directory, + map_size: int = DEFAULT_MAP_SIZE, + frame_idx_fmt: str = DEFAULT_FRAME_FMT, + type_map: list[str] | None = None, + write_batch_size: int = DEFAULT_WRITE_BATCH_SIZE, + overwrite: bool = False, + ): + """Write an ordered sequence of systems, one ``frame_system_id`` each. - def to_labeled_system(self, data, file_name, **kwargs): - """Save a single LabeledSystem to an LMDB database.""" - from dpdata.system import LabeledSystem + Unlike :meth:`to_multi_systems` (the path used by + ``MultiSystems.to('lmdb', ...)``), the systems are **not** merged by + formula: every element of ``systems`` becomes exactly one source + system in the database, numbered ``0, 1, 2, ...`` in iteration order. + This preserves the system partition recorded in ``frame_system_ids``, + which DeePMD-kit uses for ``prob_sys_size`` based sampling. - if isinstance(file_name, tuple) and file_name[0] is self: - txn, formula = self._txn, file_name[1] - self._dump_to_txn(data, txn, formula, LabeledSystem.DTYPES) + Parameters + ---------- + systems : Iterable[System] + An ordered iterable of :class:`dpdata.System` / + :class:`dpdata.LabeledSystem` objects (or raw data dicts). A + :class:`dpdata.MultiSystems` must not be used here, because its + frames are already merged by formula. + directory : str + Output LMDB directory. + map_size : int, optional + Maximum LMDB size in bytes. Default is 1 TiB (sparse). + frame_idx_fmt : str, optional + Format used for the per-frame integer key. Default ``"012d"``. + type_map : list[str], optional + Global element table. If ``None``, the union of the elements of + all systems is used, in first-appearance order. When the systems + are produced lazily (a generator) a ``type_map`` should be given + so that the whole sequence need not be held in memory. + write_batch_size : int, optional + Number of frames committed per LMDB write transaction. + overwrite : bool, optional + Whether to replace an existing destination after the new database + has been written and validated. The default is ``False``. + """ + if hasattr(systems, "systems"): + raise TypeError( + "dump_systems expects the original ordered systems, not a " + "MultiSystems object whose systems are already merged by formula." + ) + if type_map is None: + systems = list(systems) + if not systems: + raise LMDBError("Cannot write an empty LMDB dataset.") + resolved: list[str] = [] + seen: set[str] = set() + for ss in systems: + data = _source_data(ss) + for name in _source_active_names(data): + if name not in seen: + seen.add(name) + resolved.append(name) else: - # Single system call: use to_multi_systems logic - # Infer formula from data if possible, or use default - formula = kwargs.get("formula", "unknown") - gen = self.to_multi_systems([formula], file_name, **kwargs) - handle = next(gen) - self.to_labeled_system(data, handle, **kwargs) + resolved = [str(n) for n in type_map] + iterator = iter(systems) try: - next(gen) + first = next(iterator) except StopIteration: - pass + raise LMDBError("Cannot write an empty LMDB dataset.") from None + systems = itertools.chain((first,), iterator) + + writer = _LMDBWriter( + directory, + map_size=map_size, + frame_idx_fmt=frame_idx_fmt, + type_map=resolved, + write_batch_size=write_batch_size, + overwrite=overwrite, + ) + try: + for ss in systems: + writer.write_system(_source_data(ss)) + writer.finish() + except BaseException: + writer.abort() + raise def to_system(self, data, file_name, **kwargs): - """Save a single System to an LMDB database.""" - from dpdata.system import System + """Save a single (unlabeled) System to an LMDB database.""" + self._to_any(data, file_name, **kwargs) + + def to_labeled_system(self, data, file_name, **kwargs): + """Save a single LabeledSystem to an LMDB database.""" + self._to_any(data, file_name, **kwargs) + def _to_any(self, data, file_name, **kwargs): if isinstance(file_name, tuple) and file_name[0] is self: - txn, formula = self._txn, file_name[1] - self._dump_to_txn(data, txn, formula, System.DTYPES) - else: - # Single system call - formula = kwargs.get("formula", "unknown") - gen = self.to_multi_systems([formula], file_name, **kwargs) - handle = next(gen) - self.to_system(data, handle, **kwargs) - try: - next(gen) - except StopIteration: - pass + if self._writer is None: + raise LMDBError("LMDB writer is not active.") + self._writer.write_system(data) + return + gen = self.to_multi_systems(["system"], file_name, **kwargs) + try: + next(gen) + writer = self._writer + if writer is None: + raise LMDBError("LMDB writer is not active.") + writer.write_system(data) + next(gen) + except StopIteration: + return + except BaseException: + gen.close() + raise - def from_multi_systems(self, file_name, map_size=1000000000, **kwargs): - """Load multiple systems from a single LMDB database. + # ------------------------------------------------------------------ # + # Reading + # ------------------------------------------------------------------ # + def from_multi_systems( + self, + directory, + mixed_type: bool = False, + type_map: list[str] | None = None, + max_frames: int | None = DEFAULT_MAX_FRAMES, + **kwargs, + ): + """Load systems from a flat LMDB. + + Frames are grouped by atom-count composition. Atom order is + canonicalized by a stable sort on the global atom type, and every + registered atomic field follows the same permutation. Each + composition becomes one ``dpdata`` system. Parameters ---------- - file_name : str - The path to the LMDB database directory. - map_size : int, optional - Maximum size of the LMDB database in bytes. + directory : str + Path to the LMDB directory. + mixed_type : bool, optional + If ``False`` (default) each system's ``atom_names`` is the + compact set of elements it actually contains. If ``True`` every + system keeps the full global ``type_map`` as ``atom_names`` + (with zero counts for absent elements). + type_map : list[str], optional + Requested element table for the returned systems. When the file + stores a ``type_map``, the stored global indices are remapped to + this table **by element name** (consistent with DeePMD-kit); every + element in the file must be present in ``type_map``. When the file + has no ``type_map``, the indices are named positionally from this + argument. Defaults to the ``type_map`` stored in the file. + max_frames : int or None, optional + Maximum number of frames loaded into memory. The default is + 100,000. Set to ``None`` only when sufficient memory is available. **kwargs : dict other parameters Yields ------ dict - data dictionary for each system + system data dictionary for each composition group. """ - from dpdata.data_type import Axis, DataType + frames, meta = self._read_all_frames(directory, max_frames=max_frames) + file_type_map = _meta_get(meta, "type_map") + if file_type_map: + names = list( + _validate_names( + [ + n.decode() if isinstance(n, bytes) else n + for n in file_type_map + ], + "file type_map", + ) + ) + self._validate_frame_types(frames, names, meta) + if type_map is not None: + requested = list(_validate_names(type_map, "type_map")) + missing = [name for name in names if name not in requested] + if missing: + raise LMDBError( + f"Elements {missing} from the file type_map are missing " + f"from the requested type_map {requested}." + ) + remap = np.array([requested.index(n) for n in names], dtype=int) + for fr in frames: + fr["atom_types"] = remap[np.asarray(fr["atom_types"]).astype(int)] + fr["atom_numbs"] = np.bincount( + fr["atom_types"], minlength=len(requested) + ).tolist() + names = requested + elif type_map is not None: + names = list(_validate_names(type_map, "type_map")) + self._validate_frame_types(frames, names, meta) + else: + names = self._infer_type_names(frames) + self._validate_frame_types(frames, names, meta) + self._normalize_reference_fields(frames) + specs = self._resolve_read_specs(frames, meta) + frames = self._rename_frame_fields(frames, specs) + datasets = self._group_frames(frames, names, mixed_type, specs) + dp_specs = {spec.dp_name: spec for spec in specs.values()} from dpdata.system import LabeledSystem, System - with lmdb.open(file_name, readonly=True) as env: - with env.begin() as txn: - metadata_packed = txn.get(b"__metadata__") - if metadata_packed is None: - raise LMDBMetadataError("LMDB database does not contain metadata.") - metadata = msgpack.unpackb(metadata_packed, raw=False) - frame_idx_fmt = metadata.get("frame_idx_fmt", "012d") - - for sys_info in metadata["system_info"]: - system_frames = [] - start_idx = sys_info["start_idx"] - nframes = sys_info["nframes"] - data_shapes = sys_info.get("data_shapes", {}) - frame_dependent_keys = sys_info.get("frame_dependent_keys", []) - - for i in range(start_idx, start_idx + nframes): - key = f"{i:{frame_idx_fmt}}".encode("ascii") - value = txn.get(key) - if value is None: - raise LMDBFrameError(f"Frame data not found for key: {key}") - frame_data = msgpack.unpackb(value, raw=False) - system_frames.append(frame_data) - - # Aggregate data for one system - first_frame = system_frames[0] - is_labeled = "energies" in first_frame - cls = LabeledSystem if is_labeled else System - - # Auto-register unknown data types - existing_dt_names = [dt.name for dt in cls.DTYPES] - new_dts = [] - axis_map = {a.value: a for a in Axis} - for key, val in first_frame.items(): - if key not in existing_dt_names and key in data_shapes: - shape_raw = data_shapes[key] - if shape_raw is not None: - shape = tuple([axis_map.get(s, s) for s in shape_raw]) - else: - shape = None - - v_arr = np.array(val) - new_dts.append( - DataType(key, type(v_arr), shape=shape, required=False) - ) - - if new_dts: - cls.register_data_type(*new_dts) - - agg_data = {} - for key, val in first_frame.items(): - if key in frame_dependent_keys: - agg_data[key] = np.array([d[key] for d in system_frames]) - else: - agg_data[key] = val - - yield agg_data + original_system_dtypes = System.DTYPES + original_labeled_dtypes = LabeledSystem.DTYPES + self._register_specs(datasets, dp_specs) + completed = False + try: + yield from datasets + completed = True + finally: + if not completed: + System.DTYPES = original_system_dtypes + LabeledSystem.DTYPES = original_labeled_dtypes - def from_labeled_system(self, file_name, **kwargs): - """Load data for a single LabeledSystem from an LMDB database.""" - if isinstance(file_name, dict): - return file_name - # from_multi_systems returns a generator of dicts - gen = self.from_multi_systems(file_name, **kwargs) - return next(gen) + @staticmethod + def _normalize_reference_fields(frames: list[dict[str, Any]]) -> None: + """Normalize flattened atomic parameters emitted by reference converters.""" + for index, frame in enumerate(frames): + if "aparam" not in frame: + continue + value = np.asarray(frame["aparam"]) + natoms = len(frame["atom_types"]) + if value.ndim == 1: + if value.size == 0 or value.size % natoms: + raise LMDBFrameError( + f"Frame {index} aparam length {value.size} is not a " + f"positive multiple of natoms={natoms}." + ) + frame["aparam"] = value.reshape(natoms, value.size // natoms) + elif value.ndim != 2 or value.shape[0] != natoms: + raise LMDBFrameError( + f"Frame {index} aparam must have shape (natoms, ndof) or " + f"a flattened length divisible by natoms={natoms}, got " + f"{value.shape}." + ) def from_system(self, file_name, **kwargs): """Load data for a single System from an LMDB database.""" if isinstance(file_name, dict): return file_name - # from_multi_systems returns a generator of dicts - gen = self.from_multi_systems(file_name, **kwargs) - return next(gen) + return self._first_system(file_name, require_labeled=False, **kwargs) + + def from_labeled_system(self, file_name, **kwargs): + """Load data for a single LabeledSystem from an LMDB database.""" + if isinstance(file_name, dict): + return file_name + return self._first_system(file_name, require_labeled=True, **kwargs) + + def _first_system(self, file_name, *, require_labeled: bool, **kwargs): + """Return the first composition group, warning if more than one exists. + + A single ``System`` can only hold one composition; an LMDB usually + stores several. Callers that need every system should use + ``MultiSystems.from_file``. + """ + from dpdata.system import LabeledSystem, System + + original_system_dtypes = System.DTYPES + original_labeled_dtypes = LabeledSystem.DTYPES + systems = list(self.from_multi_systems(file_name, **kwargs)) + first = systems[0] + if require_labeled and "energies" not in first: + System.DTYPES = original_system_dtypes + LabeledSystem.DTYPES = original_labeled_dtypes + raise LMDBFrameError( + "LMDB does not contain energies required by LabeledSystem. " + "Load it as System or pass labeled=False to MultiSystems." + ) + if len(systems) > 1: + warnings.warn( + f"LMDB '{file_name}' contains more than one composition; only the " + "first is loaded into a single System. Use " + "dpdata.MultiSystems.from_file(..., fmt='lmdb') to load all of them.", + stacklevel=2, + ) + return first + + def _read_all_frames( + self, directory: str, *, max_frames: int | None + ) -> tuple[list[dict[str, Any]], dict[str | bytes, Any]]: + resolved, env = _open_read_env(directory) + try: + with env.begin() as txn: + meta = _read_metadata(txn) + frame_fmt = _meta_get(meta, "frame_idx_fmt") or DEFAULT_FRAME_FMT + if not isinstance(frame_fmt, str): + raise LMDBMetadataError("frame_idx_fmt must be a string.") + try: + format(0, frame_fmt) + except (TypeError, ValueError) as exc: + raise LMDBMetadataError( + f"Invalid frame_idx_fmt {frame_fmt!r}: {exc}" + ) from exc + nframes = _meta_get(meta, "nframes") + if nframes is None: + raise LMDBMetadataError("Metadata does not contain 'nframes'.") + nframes = _require_integer(nframes, "nframes", minimum=1) + if max_frames is not None: + if max_frames <= 0: + raise ValueError("max_frames must be positive or None.") + if nframes > max_frames: + raise LMDBError( + f"LMDB '{directory}' contains {nframes:,} frames, " + f"exceeding max_frames={max_frames:,}. Use DeePMD-kit's " + "streaming reader, or set max_frames=None when sufficient " + "memory is available." + ) + frame_nlocs = _meta_get(meta, "frame_nlocs") + if frame_nlocs is not None: + if not isinstance(frame_nlocs, (list, tuple)): + raise LMDBMetadataError( + "frame_nlocs must be a sequence of integers." + ) + if len(frame_nlocs) != nframes: + raise LMDBMetadataError( + "frame_nlocs length does not match nframes." + ) + frame_nlocs = [ + _require_integer( + value, + f"frame_nlocs[{index}]", + minimum=1, + ) + for index, value in enumerate(frame_nlocs) + ] + meta["frame_nlocs"] = frame_nlocs + frame_system_ids = _meta_get(meta, "frame_system_ids") + if frame_system_ids is not None: + if not isinstance(frame_system_ids, (list, tuple)): + raise LMDBMetadataError( + "frame_system_ids must be a sequence of integers." + ) + if len(frame_system_ids) != nframes: + raise LMDBMetadataError( + "frame_system_ids length does not match nframes." + ) + frame_system_ids = [ + _require_integer( + value, + f"frame_system_ids[{index}]", + minimum=0, + ) + for index, value in enumerate(frame_system_ids) + ] + meta["frame_system_ids"] = frame_system_ids + frames: list[dict[str, Any]] = [] + for i in range(int(nframes)): + key = format(i, frame_fmt).encode("ascii") + raw = txn.get(key) + if raw is None: + raise LMDBFrameError(f"Frame data not found for key: {key!r}") + frames.append(_decode_frame(bytes(raw))) + finally: + _close_read_env(resolved) + return frames, meta + + @staticmethod + def _infer_type_names(frames: list[dict[str, Any]]) -> list[str]: + """Create positional names for a file without type_map metadata.""" + maximum = -1 + for frame in frames: + raw = np.asarray(frame["atom_types"]) + if not np.issubdtype(raw.dtype, np.integer) or raw.ndim != 1: + raise LMDBFrameError("atom_types must be a one-dimensional integer array.") + if raw.size and int(raw.min()) < 0: + raise LMDBFrameError("atom_types must not contain negative indices.") + if raw.size: + maximum = max(maximum, int(raw.max())) + if maximum < 0: + raise LMDBFrameError("Every frame must contain at least one atom.") + return [f"Type_{index}" for index in range(maximum + 1)] + + @staticmethod + def _validate_frame_types( + frames: list[dict[str, Any]], + names: list[str], + metadata: dict[str | bytes, Any], + ) -> None: + """Validate per-frame atom types and count metadata.""" + frame_nlocs = _meta_get(metadata, "frame_nlocs") + for index, frame in enumerate(frames): + try: + atom_types = _validate_integer_types( + np.asarray(frame["atom_types"]), + name=f"frame {index} atom_types", + upper_bound=len(names), + allow_virtual=False, + ) + except (KeyError, LMDBError) as exc: + raise LMDBFrameError(str(exc)) from exc + if atom_types.ndim != 1 or atom_types.size == 0: + raise LMDBFrameError( + f"Frame {index} atom_types must be a non-empty 1-D array." + ) + if frame_nlocs is not None and int(frame_nlocs[index]) != atom_types.size: + raise LMDBMetadataError( + f"frame_nlocs[{index}] does not match atom_types length." + ) + expected = np.bincount(atom_types, minlength=len(names)).tolist() + stored = frame.get("atom_numbs") + if stored is not None and list(stored) != expected: + raise LMDBFrameError( + f"Frame {index} atom_numbs is inconsistent with atom_types." + ) + frame["atom_types"] = atom_types + frame["atom_numbs"] = expected + + @staticmethod + def _metadata_mapping( + metadata: dict[str | bytes, Any], key: str + ) -> dict[str, Any]: + """Return a metadata mapping with decoded string keys.""" + raw = _meta_get(metadata, key) or {} + if not isinstance(raw, dict): + raise LMDBMetadataError(f"{key} must be a mapping.") + return { + k.decode() if isinstance(k, bytes) else str(k): v + for k, v in raw.items() + } + + def _resolve_read_specs( + self, + frames: list[dict[str, Any]], + metadata: dict[str | bytes, Any], + ) -> dict[str, _FieldSpec]: + """Resolve every disk field to a dpdata name and symbolic shape.""" + by_dp, by_disk = _field_registries(prefer_protocol=True) + shape_hints = self._metadata_mapping(metadata, "dp_data_shapes") + name_hints = self._metadata_mapping(metadata, "dp_data_names") + disk_names = { + key + for frame in frames + for key in frame + if key not in {"atom_types", "atom_numbs"} + } + specs: dict[str, _FieldSpec] = {} + used_dp_names: dict[str, str] = {} + for disk_name in sorted(disk_names): + hinted_name = name_hints.get(disk_name) + if isinstance(hinted_name, bytes): + hinted_name = hinted_name.decode() + known = by_disk.get(disk_name) + dp_name = str(hinted_name) if hinted_name is not None else ( + known.dp_name if known is not None else disk_name + ) + if known is not None and hinted_name is not None and dp_name != known.dp_name: + raise LMDBMetadataError( + f"dp_data_names maps '{disk_name}' to '{dp_name}', but the " + f"registered protocol maps it to '{known.dp_name}'." + ) + expected_spec = by_dp.get(dp_name) + if ( + expected_spec is not None + and disk_name != expected_spec.disk_name + ): + raise LMDBMetadataError( + f"LMDB key '{disk_name}' cannot map to registered field " + f"'{dp_name}', whose protocol key is " + f"'{expected_spec.disk_name}'." + ) + try: + _validate_field_namespace(dp_name, disk_name) + except LMDBError as exc: + raise LMDBFrameError(str(exc)) from exc + previous_disk = used_dp_names.get(dp_name) + if previous_disk is not None and previous_disk != disk_name: + raise LMDBMetadataError( + f"LMDB keys '{previous_disk}' and '{disk_name}' both map to " + f"dpdata field '{dp_name}'." + ) + used_dp_names[dp_name] = disk_name + + if disk_name in shape_hints: + shape = _tokens_to_shape(shape_hints[disk_name]) + if ( + known is not None + and not self._shape_hint_compatible(known.shape, shape) + ): + raise LMDBMetadataError( + f"dp_data_shapes for known field '{disk_name}' changes " + f"its protocol shape from {known.shape} to {shape}." + ) + elif ( + known is not None + and known.dp_name in _PROTOCOL_FIELD_NAMES + ): + shape = known.shape + else: + shape = self._infer_field_shape(disk_name, frames) + if shape is None: + raise LMDBMetadataError( + f"No shape is available for LMDB field '{disk_name}'." + ) + deepmd_name = ( + known.deepmd_name + if known is not None + else by_dp[dp_name].deepmd_name + if dp_name in by_dp + else disk_name + ) + specs[disk_name] = _FieldSpec( + dp_name=dp_name, + disk_name=disk_name, + shape=tuple(shape), + deepmd_name=deepmd_name, + ) + return specs + + @staticmethod + def _shape_hint_compatible( + protocol: tuple[int | Axis, ...], + hinted: tuple[int | Axis, ...], + ) -> bool: + """Allow metadata to refine wildcard dimensions but not protocol axes.""" + if len(protocol) != len(hinted): + return False + for expected, actual in zip(protocol, hinted): + if expected == -1: + if isinstance(actual, int): + continue + return False + if expected != actual: + return False + return True + + @staticmethod + def _infer_field_shape( + disk_name: str, frames: list[dict[str, Any]] + ) -> tuple[int | Axis, ...]: + """Infer unknown shapes across all atom counts without local coincidences.""" + observed = [ + (np.asarray(frame[disk_name]).shape, len(frame["atom_types"])) + for frame in frames + if disk_name in frame + ] + ranks = {len(shape) for shape, _ in observed} + if len(ranks) != 1: + raise LMDBFrameError( + f"Field '{disk_name}' has inconsistent ranks across frames." + ) + rank = ranks.pop() + inferred: list[int | Axis] = [Axis.NFRAMES] + for axis in range(rank): + dimensions = [shape[axis] for shape, _ in observed] + follows_natoms = all( + dim == natoms + for dim, (_, natoms) in zip(dimensions, observed) + ) + if follows_natoms: + if len({natoms for _, natoms in observed}) == 1: + raise LMDBFrameError( + f"Field '{disk_name}' axis {axis} has the same length as " + "the atom count in every observed frame; it is ambiguous " + "without dp_data_shapes or a registered DataType." + ) + inferred.append(Axis.NATOMS) + elif len(set(dimensions)) == 1: + inferred.append(dimensions[0]) + else: + raise LMDBFrameError( + f"Field '{disk_name}' axis {axis} varies independently of " + "the atom count and cannot be represented by one DataType." + ) + return tuple(inferred) + + @staticmethod + def _rename_frame_fields( + frames: list[dict[str, Any]], + specs: dict[str, _FieldSpec], + ) -> list[dict[str, Any]]: + """Map disk field names to dpdata names and validate payload shapes.""" + renamed_frames: list[dict[str, Any]] = [] + for frame_index, frame in enumerate(frames): + renamed: dict[str, Any] = { + "atom_types": frame["atom_types"], + "atom_numbs": frame["atom_numbs"], + } + natoms = len(frame["atom_types"]) + for disk_name, spec in specs.items(): + if disk_name not in frame: + continue + if spec.dp_name in renamed: + raise LMDBFrameError( + f"Frame {frame_index} contains duplicate field " + f"'{spec.dp_name}'." + ) + value = np.asarray(frame[disk_name]) + LMDBFormat._validate_payload_shape( + spec, value, natoms, frame_index + ) + renamed[spec.dp_name] = value + renamed_frames.append(renamed) + return renamed_frames + + @staticmethod + def _validate_payload_shape( + spec: _FieldSpec, + value: np.ndarray, + natoms: int, + frame_index: int, + ) -> None: + """Validate a decoded per-frame payload against its field specification.""" + expected = list(spec.shape) + if spec.frame_axis is not None: + del expected[spec.frame_axis] + if value.ndim != len(expected): + raise LMDBFrameError( + f"Frame {frame_index} field '{spec.disk_name}' has shape " + f"{value.shape}, incompatible with {spec.shape}." + ) + for axis, (actual, dimension) in enumerate(zip(value.shape, expected)): + if dimension is Axis.NATOMS and actual != natoms: + raise LMDBFrameError( + f"Frame {frame_index} field '{spec.disk_name}' axis {axis} " + f"must have {natoms} atoms, got {actual}." + ) + if isinstance(dimension, int) and dimension != -1 and actual != dimension: + raise LMDBFrameError( + f"Frame {frame_index} field '{spec.disk_name}' axis {axis} " + f"must have length {dimension}, got {actual}." + ) + + def _group_frames( + self, + frames: list[dict[str, Any]], + names: list[str], + mixed_type: bool, + specs_by_disk: dict[str, _FieldSpec], + ) -> list[dict[str, Any]]: + """Canonicalize atom order and group frames by composition.""" + specs = {spec.dp_name: spec for spec in specs_by_disk.values()} + canonical_frames: list[dict[str, Any]] = [] + groups: OrderedDict[tuple[int, ...], list[int]] = OrderedDict() + for index, frame in enumerate(frames): + canonical = self._canonicalize_frame(frame, specs) + canonical_frames.append(canonical) + composition = tuple( + np.bincount( + canonical["atom_types"], minlength=len(names) + ).tolist() + ) + groups.setdefault(composition, []).append(index) + + results: list[dict[str, Any]] = [] + for composition, indices in groups.items(): + self._validate_group_schema(canonical_frames, indices) + global_atom_types = canonical_frames[indices[0]]["atom_types"] + if mixed_type: + atom_names = list(names) + atom_types = global_atom_types.copy() + else: + present = [index for index, count in enumerate(composition) if count] + atom_names = [names[index] for index in present] + remap = np.full(len(names), -1, dtype=int) + remap[present] = np.arange(len(present)) + atom_types = remap[global_atom_types] + data: dict[str, Any] = { + "atom_names": atom_names, + "atom_numbs": np.bincount( + atom_types, minlength=len(atom_names) + ).tolist(), + "atom_types": atom_types, + "orig": np.zeros(3), + } + self._aggregate_group(data, canonical_frames, indices, specs) + results.append(data) + return results + + @staticmethod + def _canonicalize_frame( + frame: dict[str, Any], specs: dict[str, _FieldSpec] + ) -> dict[str, Any]: + """Sort atoms stably by global type and reorder every atomic field.""" + atom_types = np.asarray(frame["atom_types"], dtype=np.int64) + permutation = np.argsort(atom_types, kind="stable") + canonical: dict[str, Any] = { + "atom_types": atom_types[permutation], + "atom_numbs": frame["atom_numbs"], + } + for name, value in frame.items(): + if name in {"atom_types", "atom_numbs"}: + continue + spec = specs[name] + reordered = np.asarray(value) + for axis in spec.payload_atom_axes: + reordered = np.take(reordered, permutation, axis=axis) + canonical[name] = reordered + return canonical + + @staticmethod + def _validate_group_schema( + frames: list[dict[str, Any]], indices: list[int] + ) -> None: + """Reject label or PBC inconsistencies that a System cannot represent.""" + expected = set(frames[indices[0]]) - {"atom_types", "atom_numbs"} + for index in indices[1:]: + actual = set(frames[index]) - {"atom_types", "atom_numbs"} + if actual != expected: + missing = sorted(expected - actual) + extra = sorted(actual - expected) + raise LMDBFrameError( + "Frames with the same composition must contain identical " + f"fields; frame {index} is missing {missing} and adds {extra}." + ) + + @staticmethod + def _aggregate_group( + data: dict[str, Any], + frames: list[dict[str, Any]], + indices: list[int], + specs: dict[str, _FieldSpec], + ) -> None: + """Aggregate canonical per-frame payloads into one System data dictionary.""" + fields = set(frames[indices[0]]) - {"atom_types", "atom_numbs"} + for name in fields: + values = [np.asarray(frames[index][name]) for index in indices] + data[name] = LMDBFormat._aggregate_values(values, specs[name]) + if "cells" not in fields: + data["nopbc"] = True + data["cells"] = np.zeros((len(indices), 3, 3)) + + @staticmethod + def _aggregate_values( + values: list[np.ndarray], spec: _FieldSpec + ) -> np.ndarray: + """Aggregate values while preserving byte order and static semantics.""" + first = values[0] + for value in values[1:]: + if value.dtype != first.dtype or value.shape != first.shape: + raise LMDBFrameError( + f"Field '{spec.disk_name}' has inconsistent dtype or shape " + "within one composition." + ) + frame_axis = spec.frame_axis + if frame_axis is None: + if any(not np.array_equal(first, value, equal_nan=True) for value in values[1:]): + raise LMDBFrameError( + f"Static field '{spec.disk_name}' differs between frames." + ) + return first.copy() + output_shape = list(first.shape) + output_shape.insert(frame_axis, len(values)) + output = np.empty(output_shape, dtype=first.dtype) + for index, value in enumerate(values): + selection: list[slice | int] = [slice(None)] * output.ndim + selection[frame_axis] = index + output[tuple(selection)] = value + return output + + @staticmethod + def _register_specs( + datasets: list[dict[str, Any]], specs: dict[str, _FieldSpec] + ) -> None: + """Validate all schemas, then commit process-global registrations.""" + from dpdata.system import LabeledSystem, System + + existing = {dt.name: dt for dt in (*System.DTYPES, *LabeledSystem.DTYPES)} + protocol_dtypes = {dtype.name: dtype for dtype in _PROTOCOL_DTYPES} + used_names = { + name + for data in datasets + for name in data + if name in specs + } + pending: list[DataType] = [] + for name in sorted(used_names): + if name not in specs: + continue + spec = specs[name] + current = existing.get(name) + if current is not None: + merged_shape = LMDBFormat._merge_shapes( + current.shape, spec.shape + ) + protocol_dtype = protocol_dtypes.get(name) + if ( + merged_shape is None + and protocol_dtype is not None + and LMDBFormat._shapes_compatible( + current.shape, protocol_dtype.shape + ) + and LMDBFormat._shapes_compatible( + spec.shape, protocol_dtype.shape + ) + ): + merged_shape = protocol_dtype.shape + if ( + merged_shape is None + or current.deepmd_name != spec.deepmd_name + ): + raise LMDBError( + f"Data type '{name}' conflicts with the process-global " + f"definition: existing shape/name " + f"{current.shape}/{current.deepmd_name}, LMDB shape/name " + f"{spec.shape}/{spec.deepmd_name}." + ) + if merged_shape != current.shape: + dtype = DataType( + name, + np.ndarray, + shape=merged_shape, + required=False, + deepmd_name=spec.deepmd_name, + ) + pending.append(dtype) + existing[name] = dtype + continue + dtype = DataType( + name, + np.ndarray, + shape=spec.shape, + required=False, + deepmd_name=spec.deepmd_name, + ) + pending.append(dtype) + existing[name] = dtype + if pending: + replacements = {dtype.name: dtype for dtype in pending} + for cls in (System, LabeledSystem): + current = {dtype.name: dtype for dtype in cls.DTYPES} + current.update(replacements) + cls.DTYPES = tuple(current.values()) + + @staticmethod + def _shapes_compatible( + first: tuple[int | Axis, ...] | None, + second: tuple[int | Axis, ...] | None, + ) -> bool: + """Return whether two shapes differ only through wildcard dimensions.""" + if first is None or second is None or len(first) != len(second): + return first == second + return all( + left == right + or left == -1 + or right == -1 + for left, right in zip(first, second) + ) + + @staticmethod + def _merge_shapes( + first: tuple[int | Axis, ...] | None, + second: tuple[int | Axis, ...] | None, + ) -> tuple[int | Axis, ...] | None: + """Return the least restrictive compatible shape.""" + if not LMDBFormat._shapes_compatible(first, second): + return None + if first is None or second is None: + return first + return tuple( + -1 if left == -1 or right == -1 else left + for left, right in zip(first, second) + ) + + +def dump_systems( + systems, + file_name, + type_map: list[str] | None = None, + map_size: int = DEFAULT_MAP_SIZE, + frame_idx_fmt: str = DEFAULT_FRAME_FMT, + write_batch_size: int = DEFAULT_WRITE_BATCH_SIZE, + overwrite: bool = False, +): + """Write an ordered sequence of systems to one LMDB, preserving identity. + + Each element of ``systems`` is stored as a distinct source system, + numbered ``0, 1, 2, ...`` in iteration order, and recorded in the + ``frame_system_ids`` metadata. In contrast to + ``MultiSystems.to('lmdb', ...)``, systems are not merged by formula, so + the system partition used by DeePMD-kit's ``prob_sys_size`` is kept. + + Parameters + ---------- + systems : Iterable[System] + An ordered iterable of :class:`dpdata.System` / + :class:`dpdata.LabeledSystem` objects or raw system data dictionaries. + Do not pass a :class:`dpdata.MultiSystems`, whose frames are already + merged by formula. + file_name : str + Output LMDB directory. + type_map : list[str], optional + Global element table. If ``None``, the union of the elements of all + systems is used. Provide it explicitly to stream a generator without + materialising the whole sequence. + map_size : int, optional + Maximum LMDB size in bytes. Default is 1 TiB (sparse). + frame_idx_fmt : str, optional + Format used for the per-frame integer key. Default ``"012d"``. + write_batch_size : int, optional + Number of frames committed per LMDB write transaction. + overwrite : bool, optional + Whether to replace an existing destination after the new database has + been written and validated. The default is ``False``. + + Examples + -------- + >>> import dpdata + >>> from dpdata.formats.lmdb import dump_systems + >>> systems = [ + ... dpdata.LabeledSystem(d, fmt="deepmd/npy") for d in directories + ... ] + >>> dump_systems(systems, "data.lmdb", type_map=["H", "C", "N", "O"]) + """ + LMDBFormat().dump_systems( + systems, + file_name, + map_size=map_size, + frame_idx_fmt=frame_idx_fmt, + type_map=type_map, + write_batch_size=write_batch_size, + overwrite=overwrite, + ) diff --git a/pyproject.toml b/pyproject.toml index b94fce2e9..754dd9cdd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ dependencies = [ 'h5py', 'wcmatch', 'lmdb', - 'msgpack-numpy', + 'msgpack', ] requires-python = ">=3.10" readme = "README.md" diff --git a/tests/test_lmdb.py b/tests/test_lmdb.py index bc0fdeecc..ff56c5b05 100644 --- a/tests/test_lmdb.py +++ b/tests/test_lmdb.py @@ -2,11 +2,14 @@ import os import shutil +import threading +import typing import unittest +from pathlib import Path +from unittest import mock import lmdb import msgpack -import msgpack_numpy as m import numpy as np from comp_sys import ( CompLabeledMultiSys, @@ -17,7 +20,32 @@ ) from context import dpdata -from dpdata.formats.lmdb.format import LMDBFrameError, LMDBMetadataError +from dpdata.data_type import Axis, DataError, DataType +from dpdata.formats.lmdb.format import ( + LMDBError, + LMDBFrameError, + LMDBMetadataError, + _close_read_env, + _open_read_env, +) + + +def _reference_encode_array(array): + """Encode an array independently of the production implementation.""" + value = np.asarray(array) + return { + "type": str(value.dtype), + "shape": list(value.shape), + "data": value.tobytes(order="C"), + } + + +def _reference_packb(value): + """Pack a protocol value and assert the LMDB byte contract.""" + packed = msgpack.packb(value, use_bin_type=True) + if not isinstance(packed, bytes): + raise TypeError("msgpack.packb did not return bytes") + return packed class TestLMDBLabeledSystem(unittest.TestCase, CompLabeledSys, IsPBC): @@ -66,11 +94,9 @@ def setUp(self): ) self.ms_1 = dpdata.MultiSystems(system_1, system_2, system_3) - - # Standard API + for system in self.ms_1: + system.sort_atom_types() self.ms_1.to("lmdb", self.lmdb_path) - - # Standard API self.ms_2 = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") self.places = 6 @@ -83,60 +109,1373 @@ def tearDown(self): shutil.rmtree(self.lmdb_path) -class TestLMDBErrorHandling(unittest.TestCase): +class TestLMDBOnDiskFormat(unittest.TestCase): + """The on-disk layout must match the DeePMD-kit / reference converters.""" + + def setUp(self): + self.lmdb_path = "tmp_format.lmdb" + self.system = dpdata.LabeledSystem("poscars/OUTCAR.h2o.md", fmt="vasp/outcar") + self.system.to("lmdb", self.lmdb_path) + + def tearDown(self): + if os.path.exists(self.lmdb_path): + shutil.rmtree(self.lmdb_path) + + def test_metadata_schema(self): + with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: + with env.begin() as txn: + meta = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) + self.assertEqual(meta["nframes"], self.system.get_nframes()) + self.assertEqual(meta["frame_idx_fmt"], "012d") + self.assertEqual(meta["type_map"], self.system.data["atom_names"]) + self.assertEqual(len(meta["frame_nlocs"]), self.system.get_nframes()) + self.assertTrue(all(n == self.system.get_natoms() for n in meta["frame_nlocs"])) + self.assertEqual(len(meta["frame_system_ids"]), self.system.get_nframes()) + self.assertTrue(all(s == 0 for s in meta["frame_system_ids"])) + # legacy field must be gone + self.assertNotIn("system_info", meta) + + def test_frame_encoding(self): + with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: + with env.begin() as txn: + raw = txn.get(b"000000000000") + meta = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) + frame = msgpack.unpackb(raw, raw=False) + # string keys, manual encoding + for key in ("atom_types", "coords", "cells", "energies", "forces"): + self.assertIn(key, frame) + self.assertIsInstance(frame[key], dict) + self.assertIn("type", frame[key]) + self.assertIn("shape", frame[key]) + self.assertIn("data", frame[key]) + # atom_types are int32 global indices + self.assertEqual(frame["atom_types"]["type"], "int32") + # energies stored as a 0-d array + self.assertEqual(list(frame["energies"]["shape"]), []) + # atom_numbs is a plain list over the full type_map + self.assertIsInstance(frame["atom_numbs"], list) + self.assertEqual(len(frame["atom_numbs"]), len(meta["type_map"])) + + def test_atom_types_are_global(self): + with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: + with env.begin() as txn: + raw = txn.get(b"000000000000") + meta = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) + frame = msgpack.unpackb(raw, raw=False) + atype = np.frombuffer(frame["atom_types"]["data"], dtype=np.int32).reshape( + frame["atom_types"]["shape"] + ) + type_map = meta["type_map"] + decoded_names = [type_map[i] for i in atype] + expected = [ + self.system.data["atom_names"][t] for t in self.system.data["atom_types"] + ] + self.assertEqual(decoded_names, expected) + + +class TestLMDBMixedTypeRead(unittest.TestCase): + """mixed_type=True keeps the full global type_map on every system.""" + def setUp(self): - self.lmdb_path_missing_metadata = "tmp_missing_metadata.lmdb" - self.lmdb_path_missing_frame = "tmp_missing_frame.lmdb" - - # Ensure cleanup in case of previous test failures - if os.path.exists(self.lmdb_path_missing_metadata): - shutil.rmtree(self.lmdb_path_missing_metadata) - if os.path.exists(self.lmdb_path_missing_frame): - shutil.rmtree(self.lmdb_path_missing_frame) - - # For test_load_missing_frame_data, create a valid LMDB environment - # and write metadata, but no actual frames. - env = lmdb.open(self.lmdb_path_missing_frame, map_size=1000000000) + self.lmdb_path = "tmp_mixed.lmdb" + s1 = dpdata.LabeledSystem("gaussian/methane.gaussianlog", fmt="gaussian/log") + s2 = dpdata.LabeledSystem( + "gaussian/methane_sub.gaussianlog", fmt="gaussian/log" + ) + self.ms = dpdata.MultiSystems(s1, s2) + self.type_map = ["H", "C", "N", "O"] + self.ms.to("lmdb", self.lmdb_path, type_map=self.type_map) + + def tearDown(self): + if os.path.exists(self.lmdb_path): + shutil.rmtree(self.lmdb_path) + + def test_mixed_preserves_full_type_map(self): + # via MultiSystems the element order is normalized (sorted), but the + # full type_map set must be preserved on every system. + ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", mixed_type=True) + for ss in ms: + names = [str(n) for n in ss.data["atom_names"]] + self.assertEqual(sorted(names), sorted(self.type_map)) + self.assertEqual(len(ss.data["atom_numbs"]), len(self.type_map)) + + def test_mixed_single_system_preserves_order(self): + # a direct single-system load of a single-composition file is not + # reordered, so the global type_map order is preserved exactly. + path = "tmp_mixed_single.lmdb" + try: + s = dpdata.LabeledSystem("gaussian/methane.gaussianlog", fmt="gaussian/log") + s.to("lmdb", path, type_map=self.type_map) + ls = dpdata.LabeledSystem(path, fmt="lmdb", mixed_type=True) + self.assertEqual([str(n) for n in ls.data["atom_names"]], self.type_map) + finally: + if os.path.exists(path): + shutil.rmtree(path) + + def test_standard_compresses_type_map(self): + ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + # methane only contains C and H, not N/O + for ss in ms: + self.assertNotIn("N", ss.data["atom_names"]) + self.assertNotIn("O", ss.data["atom_names"]) + + +class TestLMDBTypeMapOverride(unittest.TestCase): + def setUp(self): + self.lmdb_path = "tmp_type_map.lmdb" + self.system = dpdata.LabeledSystem("poscars/OUTCAR.h2o.md", fmt="vasp/outcar") + + def tearDown(self): + if os.path.exists(self.lmdb_path): + shutil.rmtree(self.lmdb_path) + + def test_explicit_type_map(self): + type_map = ["H", "He", "Li", "Be", "B", "C", "N", "O"] + self.system.to("lmdb", self.lmdb_path, type_map=type_map) + with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: + with env.begin() as txn: + meta = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) + raw = txn.get(b"000000000000") + self.assertEqual(meta["type_map"], type_map) + frame = msgpack.unpackb(raw, raw=False) + atype = np.frombuffer(frame["atom_types"]["data"], dtype=np.int32).reshape( + frame["atom_types"]["shape"] + ) + decoded = [type_map[i] for i in atype] + expected = [ + self.system.data["atom_names"][t] for t in self.system.data["atom_types"] + ] + self.assertEqual(decoded, expected) + + def test_missing_element_raises(self): + from dpdata.formats.lmdb.format import LMDBError + + # water needs O and H; this type_map omits O. + with self.assertRaises(LMDBError): + self.system.to("lmdb", self.lmdb_path, type_map=["H", "He"]) + + +class TestLMDBReferenceFormatInterop(unittest.TestCase): + """dpdata must read an LMDB written exactly like the reference converters.""" + + def setUp(self): + self.lmdb_path = "tmp_reference.lmdb" + self.type_map = ["H", "C", "N", "O"] + # two frames: H2O (3 atoms) and CH4 (5 atoms) + self.frames = [ + { + "atom_types": np.array([3, 0, 0], dtype=np.int32), # O H H + "coords": np.random.rand(3, 3).astype(np.float32), + "cells": (np.eye(3) * 10).astype(np.float32), + "energies": np.array(-10.0, dtype=np.float32), + "forces": np.random.rand(3, 3).astype(np.float32), + }, + { + "atom_types": np.array([1, 0, 0, 0, 0], dtype=np.int32), # C H H H H + "coords": np.random.rand(5, 3).astype(np.float32), + "cells": (np.eye(3) * 12).astype(np.float32), + "energies": np.array(-20.0, dtype=np.float32), + "forces": np.random.rand(5, 3).astype(np.float32), + }, + ] + env = lmdb.open(self.lmdb_path, map_size=1 << 30) + frame_nlocs = [] with env.begin(write=True) as txn: - metadata = { - "nframes": 1, - "system_info": [ - { - "formula": "H2O", - "natoms": [1, 2], - "nframes": 1, - "start_idx": 0, - } - ], + for i, fr in enumerate(self.frames): + out: dict[str, object] = {} + for k, v in fr.items(): + out[k] = _reference_encode_array(v) + ntypes = len(self.type_map) + counts = np.bincount(fr["atom_types"], minlength=ntypes)[:ntypes] + out["atom_numbs"] = [int(c) for c in counts] + txn.put( + format(i, "012d").encode(), _reference_packb(out) + ) + frame_nlocs.append(len(fr["atom_types"])) + meta: dict[str, object] = { + "nframes": len(self.frames), + "frame_idx_fmt": "012d", + "type_map": self.type_map, + "frame_nlocs": frame_nlocs, + "frame_system_ids": list(range(len(self.frames))), } - m.patch() # Ensure numpy patching for metadata - txn.put(b"__metadata__", msgpack.packb(metadata, use_bin_type=True)) + txn.put(b"__metadata__", _reference_packb(meta)) env.close() def tearDown(self): - if os.path.exists(self.lmdb_path_missing_metadata): - shutil.rmtree(self.lmdb_path_missing_metadata) - if os.path.exists(self.lmdb_path_missing_frame): - shutil.rmtree(self.lmdb_path_missing_frame) + if os.path.exists(self.lmdb_path): + shutil.rmtree(self.lmdb_path) - def test_load_missing_metadata(self): - # Create a valid, empty LMDB environment, then test for missing metadata - lmdb.open( - self.lmdb_path_missing_metadata, map_size=1000000000 - ).close() # Creates empty LMDB environment + def test_read_reference(self): + ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + self.assertEqual(ms.get_nframes(), 2) + self.assertEqual(len(ms), 2) + # composition-based grouping: water (3 atoms) + methane (5 atoms) + self.assertEqual(sorted(s.get_natoms() for s in ms), [3, 5]) + def test_read_reference_values(self): + ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + water = next(s for s in ms if s.get_natoms() == 3) + # Frames are canonicalized by global type while coordinates and labels + # follow the same stable permutation. + names = [water.data["atom_names"][t] for t in water.data["atom_types"]] + self.assertEqual(names, ["H", "H", "O"]) + permutation = np.argsort(self.frames[0]["atom_types"], kind="stable") + np.testing.assert_allclose( + water.data["coords"][0], + self.frames[0]["coords"][permutation], + rtol=1e-6, + atol=1e-6, + ) + np.testing.assert_allclose( + water.data["energies"][0], self.frames[0]["energies"], rtol=1e-6 + ) + + +def _make_labeled_system( + atom_types: list[int], atom_names: list[str], nframes: int +) -> dpdata.LabeledSystem: + atype = np.array(atom_types, dtype=int) + natoms = len(atype) + numbs = [ + int(np.count_nonzero(atype == i)) for i in range(len(atom_names)) + ] + data = { + "atom_numbs": numbs, + "atom_names": list(atom_names), + "atom_types": atype, + "orig": np.zeros(3), + "cells": np.tile(np.eye(3) * 10.0, (nframes, 1, 1)), + "coords": np.random.rand(nframes, natoms, 3), + "energies": np.arange(nframes, dtype=float), + "forces": np.random.rand(nframes, natoms, 3), + } + return dpdata.LabeledSystem(data=data) + + +def _write_raw_lmdb(path, frames, type_map, metadata_extra=None): + """Write frames using the reference encoding (one record per frame).""" + if os.path.exists(path): + shutil.rmtree(path) + env = lmdb.open(path, map_size=1 << 30) + nlocs = [] + with env.begin(write=True) as txn: + for i, fr in enumerate(frames): + out: dict[str, object] = { + k: _reference_encode_array(np.asarray(v)) for k, v in fr.items() + } + at = np.asarray(fr["atom_types"]) + counts = np.bincount(at, minlength=len(type_map))[: len(type_map)] + out["atom_numbs"] = [int(c) for c in counts] + txn.put(format(i, "012d").encode(), _reference_packb(out)) + nlocs.append(len(at)) + meta: dict[str, object] = { + "nframes": len(frames), + "frame_idx_fmt": "012d", + "type_map": type_map, + "frame_nlocs": nlocs, + "frame_system_ids": list(range(len(frames))), + } + if metadata_extra: + meta.update(metadata_extra) + txn.put(b"__metadata__", _reference_packb(meta)) + env.close() + + +class TestLMDBReadRobustness(unittest.TestCase): + """Regression tests for the read path.""" + + def setUp(self): + self.lmdb_path = "tmp_robust.lmdb" + self.original_system_dtypes = dpdata.System.DTYPES + self.original_labeled_system_dtypes = dpdata.LabeledSystem.DTYPES + + def tearDown(self): + dpdata.System.DTYPES = self.original_system_dtypes + dpdata.LabeledSystem.DTYPES = self.original_labeled_system_dtypes + if os.path.exists(self.lmdb_path): + shutil.rmtree(self.lmdb_path) + + def test_type_map_remap_by_name(self): + # file stores type_map [H, O] (so O has global index 1); reading with a + # reordered type_map must remap by element name, not by index. + frames = [ + { + "atom_types": np.array([1, 0, 0], dtype=np.int32), # O H H + "coords": np.random.rand(3, 3).astype(np.float32), + "cells": (np.eye(3) * 10).astype(np.float32), + "energies": np.array(-1.0, dtype=np.float32), + } + ] + _write_raw_lmdb(self.lmdb_path, frames, ["H", "O"]) + ms = dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="lmdb", type_map=["O", "N", "C", "H"] + ) + s = ms[0] + names = [str(s.data["atom_names"][t]) for t in s.data["atom_types"]] + self.assertEqual(names, ["O", "H", "H"]) + + def test_type_map_missing_element_raises(self): + frames = [ + { + "atom_types": np.array([0, 1, 1], dtype=np.int32), + "coords": np.random.rand(3, 3).astype(np.float32), + "cells": (np.eye(3) * 10).astype(np.float32), + "energies": np.array(-1.0, dtype=np.float32), + } + ] + _write_raw_lmdb(self.lmdb_path, frames, ["O", "H"]) + from dpdata.formats.lmdb.format import LMDBError + + with self.assertRaises(LMDBError): + dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="lmdb", type_map=["C", "N"] + ) + + def test_inconsistent_keys_raise(self): + # same composition, but only the first frame carries virials. + frames = [ + { + "atom_types": np.array([0, 1, 1], dtype=np.int32), + "coords": np.random.rand(3, 3).astype(np.float32), + "cells": (np.eye(3) * 10).astype(np.float32), + "energies": np.array(-1.0, dtype=np.float32), + "forces": np.random.rand(3, 3).astype(np.float32), + "virials": np.random.rand(3, 3).astype(np.float32), + }, + { + "atom_types": np.array([0, 1, 1], dtype=np.int32), + "coords": np.random.rand(3, 3).astype(np.float32), + "cells": (np.eye(3) * 10).astype(np.float32), + "energies": np.array(-2.0, dtype=np.float32), + "forces": np.random.rand(3, 3).astype(np.float32), + }, + ] + _write_raw_lmdb(self.lmdb_path, frames, ["O", "H"]) + with self.assertRaisesRegex( + LMDBFrameError, "must contain identical fields" + ): + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + + def test_single_system_multi_composition_warns(self): + frames = [ + { + "atom_types": np.array([0, 1, 1], dtype=np.int32), + "coords": np.random.rand(3, 3).astype(np.float32), + "cells": (np.eye(3) * 10).astype(np.float32), + "energies": np.array(-1.0, dtype=np.float32), + }, + { + "atom_types": np.array([0, 1, 1, 1, 1], dtype=np.int32), + "coords": np.random.rand(5, 3).astype(np.float32), + "cells": (np.eye(3) * 10).astype(np.float32), + "energies": np.array(-2.0, dtype=np.float32), + }, + ] + _write_raw_lmdb(self.lmdb_path, frames, ["C", "H"]) + with self.assertWarns(UserWarning): + ls = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + # only the first composition is returned + self.assertEqual(ls.get_natoms(), 3) + + def test_same_composition_different_atom_order_is_canonicalized(self): + first_types = np.array([1, 0, 1], dtype=np.int32) + second_types = np.array([0, 1, 1], dtype=np.int32) + first_values = np.array([10.0, 20.0, 30.0]) + second_values = np.array([40.0, 50.0, 60.0]) + frames = [ + { + "atom_types": first_types, + "coords": np.repeat(first_values[:, None], 3, axis=1), + "cells": np.eye(3), + "energies": np.array(-1.0), + "forces": np.repeat(first_values[:, None], 3, axis=1), + }, + { + "atom_types": second_types, + "coords": np.repeat(second_values[:, None], 3, axis=1), + "cells": np.eye(3), + "energies": np.array(-2.0), + "forces": np.repeat(second_values[:, None], 3, axis=1), + }, + ] + _write_raw_lmdb(self.lmdb_path, frames, ["H", "O"]) + ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + self.assertEqual(len(ms), 1) + system = ms[0] + self.assertEqual(system.get_nframes(), 2) + np.testing.assert_array_equal(system["atom_types"], [0, 1, 1]) + np.testing.assert_array_equal( + system["coords"][:, :, 0], + [[20.0, 10.0, 30.0], [40.0, 50.0, 60.0]], + ) + np.testing.assert_array_equal( + system["forces"][:, :, 0], + [[20.0, 10.0, 30.0], [40.0, 50.0, 60.0]], + ) + + def test_same_composition_mixed_pbc_raises(self): + frames = [ + { + "atom_types": np.array([0, 1, 1], dtype=np.int32), + "coords": np.zeros((3, 3)), + "cells": np.eye(3), + "energies": np.array(-1.0), + }, + { + "atom_types": np.array([0, 1, 1], dtype=np.int32), + "coords": np.ones((3, 3)), + "energies": np.array(-2.0), + }, + ] + _write_raw_lmdb(self.lmdb_path, frames, ["O", "H"]) + with self.assertRaisesRegex( + LMDBFrameError, "must contain identical fields" + ): + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + + def test_max_frames_guard(self): + frames = [ + { + "atom_types": np.array([0], dtype=np.int32), + "coords": np.zeros((1, 3)), + }, + { + "atom_types": np.array([0], dtype=np.int32), + "coords": np.ones((1, 3)), + }, + ] + _write_raw_lmdb(self.lmdb_path, frames, ["H"]) + with self.assertRaisesRegex(LMDBError, "exceeding max_frames"): + dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="lmdb", max_frames=1 + ) + loaded = dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="lmdb", max_frames=None, labeled=False + ) + self.assertEqual(loaded.get_nframes(), 2) + + def test_big_endian_dtype_preserved(self): + frames = [ + { + "atom_types": np.array([0], dtype=np.int32), + "coords": np.array([[1.0, 2.0, 3.0]], dtype=">f8"), + "cells": np.eye(3, dtype=">f8"), + "energies": np.array(-1.0, dtype=">f8"), + }, + { + "atom_types": np.array([0], dtype=np.int32), + "coords": np.array([[4.0, 5.0, 6.0]], dtype=">f8"), + "cells": np.eye(3, dtype=">f8"), + "energies": np.array(-2.0, dtype=">f8"), + }, + ] + _write_raw_lmdb(self.lmdb_path, frames, ["H"]) + system = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb")[0] + self.assertEqual(system["coords"].dtype.str, ">f8") + self.assertEqual(system["cells"].dtype.str, ">f8") + self.assertEqual(system["energies"].dtype.str, ">f8") + + def test_reference_fparam_shape_does_not_collide_with_natoms(self): + frames = [ + { + "atom_types": np.array([0, 0], dtype=np.int32), + "coords": np.zeros((2, 3)), + "fparam": np.array([1.0, 2.0]), + }, + { + "atom_types": np.array([0, 0], dtype=np.int32), + "coords": np.ones((2, 3)), + "fparam": np.array([3.0, 4.0]), + }, + ] + _write_raw_lmdb(self.lmdb_path, frames, ["H"]) + dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="lmdb", labeled=False + ) + dtype = next( + dt for dt in dpdata.System.DTYPES if dt.name == "fparam" + ) + shape = dtype.shape + self.assertIsNotNone(shape) + assert shape is not None + self.assertEqual(shape[0], Axis.NFRAMES) + self.assertNotIn(Axis.NATOMS, shape) + + def test_flattened_reference_aparam_is_normalized_and_reordered(self): + fixed_aparam = DataType( + "aparam", + np.ndarray, + (Axis.NFRAMES, Axis.NATOMS, 3), + required=False, + ) + dpdata.System.register_data_type(fixed_aparam) + dpdata.LabeledSystem.register_data_type(fixed_aparam) + for raw, expected in ( + ( + np.array([10.0, 20.0, 30.0]), + np.array([[[20.0], [10.0], [30.0]]]), + ), + ( + np.array([10.0, 11.0, 20.0, 21.0, 30.0, 31.0]), + np.array( + [[[20.0, 21.0], [10.0, 11.0], [30.0, 31.0]]] + ), + ), + ): + with self.subTest(size=raw.size): + frames = [ + { + "atom_types": np.array( + [1, 0, 1], dtype=np.int32 + ), + "coords": np.zeros((3, 3)), + "aparam": raw, + } + ] + _write_raw_lmdb(self.lmdb_path, frames, ["H", "O"]) + system = dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="lmdb", labeled=False + )[0] + np.testing.assert_array_equal(system["aparam"], expected) + registered = next( + dt for dt in dpdata.System.DTYPES if dt.name == "aparam" + ) + self.assertEqual( + registered.shape, + (Axis.NFRAMES, Axis.NATOMS, -1), + ) + + def test_reference_spin_key_maps_to_dpdata_name(self): + frames = [ + { + "atom_types": np.array([0, 0], dtype=np.int32), + "coords": np.zeros((2, 3)), + "spin": np.ones((2, 3)), + } + ] + _write_raw_lmdb(self.lmdb_path, frames, ["H"]) + system = dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="lmdb", labeled=False + )[0] + self.assertIn("spins", system.data) + self.assertNotIn("spin", system.data) + np.testing.assert_array_equal(system["spins"], np.ones((1, 2, 3))) + + def test_known_field_shape_hint_cannot_remove_atom_axis(self): + frames = [ + { + "atom_types": np.array([0, 0, 0], dtype=np.int32), + "coords": np.zeros((3, 3)), + "spin": np.ones((3, 3)), + } + ] + _write_raw_lmdb( + self.lmdb_path, + frames, + ["H"], + metadata_extra={ + "dp_data_shapes": { + "spin": ["nframes", 3, 3], + } + }, + ) + with self.assertRaisesRegex( + LMDBMetadataError, "changes its protocol shape" + ): + dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="lmdb", labeled=False + ) + + def test_data_name_hint_must_match_registered_protocol(self): + frames = [ + { + "atom_types": np.array([0], dtype=np.int32), + "coords": np.zeros((1, 3)), + "foo": np.ones((1, 3)), + } + ] + _write_raw_lmdb( + self.lmdb_path, + frames, + ["H"], + metadata_extra={"dp_data_names": {"foo": "spins"}}, + ) with self.assertRaisesRegex( - LMDBMetadataError, "LMDB database does not contain metadata." + LMDBMetadataError, "protocol key is 'spin'" ): - # Standard API - dpdata.MultiSystems.from_file(self.lmdb_path_missing_metadata, fmt="lmdb") + dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="lmdb", labeled=False + ) - def test_load_missing_frame_data(self): + def test_custom_deepmd_core_key_rejected_on_read(self): + frames = [ + { + "atom_types": np.array([0], dtype=np.int32), + "coords": np.zeros((1, 3)), + "coord": np.full((1, 3), 101.0), + } + ] + _write_raw_lmdb(self.lmdb_path, frames, ["H"]) with self.assertRaisesRegex( - LMDBFrameError, "Frame data not found for key: b'000000000000'" + LMDBFrameError, "reserved LMDB key 'coord'" ): - # Standard API - dpdata.MultiSystems.from_file(self.lmdb_path_missing_frame, fmt="lmdb") + dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="lmdb", labeled=False + ) + + def test_unknown_atomic_axis_inferred_across_atom_counts(self): + frames = [ + { + "atom_types": np.array([0, 0], dtype=np.int32), + "coords": np.zeros((2, 3)), + "mystery": np.zeros((2, 1)), + }, + { + "atom_types": np.array([0, 0, 0], dtype=np.int32), + "coords": np.ones((3, 3)), + "mystery": np.ones((3, 1)), + }, + ] + _write_raw_lmdb(self.lmdb_path, frames, ["H"]) + systems = dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="lmdb", labeled=False + ) + self.assertEqual(systems.get_nframes(), 2) + dtype = next( + dt for dt in dpdata.System.DTYPES if dt.name == "mystery" + ) + self.assertEqual( + dtype.shape, + ( + Axis.NFRAMES, + Axis.NATOMS, + 1, + ), + ) + + def test_unknown_same_nloc_atom_axis_is_rejected_as_ambiguous(self): + frames = [ + { + "atom_types": np.array([1, 0, 1], dtype=np.int32), + "coords": np.zeros((3, 3)), + "mystery_atomic": np.array([10.0, 20.0, 30.0]), + } + ] + _write_raw_lmdb(self.lmdb_path, frames, ["H", "O"]) + with self.assertRaisesRegex( + LMDBFrameError, "ambiguous without dp_data_shapes" + ): + dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="lmdb", labeled=False + ) + + def test_process_global_schema_conflict_raises(self): + second_path = "tmp_robust_second.lmdb" + try: + first_frames = [ + { + "atom_types": np.array([0], dtype=np.int32), + "coords": np.zeros((1, 3)), + "custom": np.array([1.0, 2.0]), + } + ] + second_frames = [ + { + "atom_types": np.array([0], dtype=np.int32), + "coords": np.zeros((1, 3)), + "custom": np.array([1.0, 2.0, 3.0]), + } + ] + _write_raw_lmdb(self.lmdb_path, first_frames, ["H"]) + _write_raw_lmdb(second_path, second_frames, ["H"]) + dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="lmdb", labeled=False + ) + with self.assertRaisesRegex( + LMDBError, "process-global definition" + ): + dpdata.MultiSystems.from_file( + second_path, fmt="lmdb", labeled=False + ) + finally: + if os.path.exists(second_path): + shutil.rmtree(second_path) + + def test_failed_read_does_not_partially_register_data_types(self): + frames = [ + { + "atom_types": np.array([0], dtype=np.int32), + "coords": np.zeros((1, 3)), + "new_scalar": np.array(1.0), + }, + { + "atom_types": np.array([0, 0], dtype=np.int32), + "coords": np.zeros((2, 3)), + "forces": np.zeros((2, 3)), + }, + { + "atom_types": np.array([0, 0], dtype=np.int32), + "coords": np.ones((2, 3)), + }, + ] + _write_raw_lmdb(self.lmdb_path, frames, ["H"]) + with self.assertRaisesRegex( + LMDBFrameError, "must contain identical fields" + ): + dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="lmdb", labeled=False + ) + self.assertNotIn( + "new_scalar", [dt.name for dt in dpdata.System.DTYPES] + ) + + def test_caller_construction_failure_rolls_back_registration(self): + frames = [ + { + "atom_types": np.array([0], dtype=np.int32), + "coords": np.zeros((1, 3)), + "unlabeled_custom": np.array(1.0), + } + ] + _write_raw_lmdb(self.lmdb_path, frames, ["H"]) + with self.assertRaises(DataError): + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + self.assertNotIn( + "unlabeled_custom", + [dt.name for dt in dpdata.System.DTYPES], + ) + self.assertNotIn( + "unlabeled_custom", + [dt.name for dt in dpdata.LabeledSystem.DTYPES], + ) + + def test_direct_labeled_read_rejects_unlabeled_before_registration(self): + frames = [ + { + "atom_types": np.array([0], dtype=np.int32), + "coords": np.zeros((1, 3)), + "unlabeled_custom": np.array(1.0), + } + ] + _write_raw_lmdb(self.lmdb_path, frames, ["H"]) + with self.assertRaisesRegex( + LMDBFrameError, "required by LabeledSystem" + ): + dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + self.assertNotIn( + "unlabeled_custom", + [dt.name for dt in dpdata.System.DTYPES], + ) + + +class TestLMDBDumpSystems(unittest.TestCase): + """dump_systems keeps every source system as a distinct frame_system_id.""" + + def setUp(self): + self.lmdb_path = "tmp_dump_systems.lmdb" + self._remove_temporary_directories() + # A and C share a formula (water); a plain MultiSystems would merge them. + self.A = _make_labeled_system([0, 1, 1], ["O", "H"], 2) + self.B = _make_labeled_system([0, 1, 1, 1, 1], ["C", "H"], 1) + self.C = _make_labeled_system([0, 1, 1], ["O", "H"], 3) + + def tearDown(self): + if os.path.exists(self.lmdb_path): + shutil.rmtree(self.lmdb_path) + self._remove_temporary_directories() + + def _remove_temporary_directories(self): + for path in Path(".").glob(f".{self.lmdb_path}.tmp-*"): + shutil.rmtree(path) + for path in Path(".").glob(f".{self.lmdb_path}.backup-*"): + shutil.rmtree(path) + + def test_system_ids_preserved(self): + from dpdata.formats.lmdb import dump_systems + + dump_systems([self.A, self.B, self.C], self.lmdb_path) + with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: + with env.begin() as txn: + meta = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) + self.assertEqual(meta["nframes"], 6) + # three distinct source systems, in order, none merged + self.assertEqual(meta["frame_system_ids"], [0, 0, 1, 2, 2, 2]) + self.assertEqual(meta["frame_nlocs"], [3, 3, 5, 3, 3, 3]) + # union type_map in first-appearance order + self.assertEqual(meta["type_map"], ["O", "H", "C"]) + + def test_contrast_with_multisystems_merge(self): + # the default MultiSystems path merges A and C by formula. + ms = dpdata.MultiSystems(self.A, self.B, self.C) + ms.to("lmdb", self.lmdb_path) + with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: + with env.begin() as txn: + meta = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) + # only two systems remain after formula merging + self.assertEqual(sorted(set(meta["frame_system_ids"])), [0, 1]) + + def test_roundtrip_total_frames(self): + from dpdata.formats.lmdb import dump_systems + + dump_systems([self.A, self.B, self.C], self.lmdb_path) + ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + self.assertEqual(ms.get_nframes(), 6) + + def test_generator_with_type_map(self): + from dpdata.formats.lmdb import dump_systems + + def gen(): + yield self.A + yield self.B + yield self.C + + dump_systems(gen(), self.lmdb_path, type_map=["H", "C", "N", "O"]) + with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: + with env.begin() as txn: + meta = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) + self.assertEqual(meta["type_map"], ["H", "C", "N", "O"]) + self.assertEqual(meta["frame_system_ids"], [0, 0, 1, 2, 2, 2]) + + def test_empty_input_rejected(self): + from dpdata.formats.lmdb import dump_systems + + with self.assertRaisesRegex(LMDBError, "empty"): + dump_systems([], self.lmdb_path) + self.assertFalse(os.path.exists(self.lmdb_path)) + + def test_empty_multisystems_rejected(self): + with self.assertRaisesRegex(LMDBError, "empty"): + dpdata.MultiSystems().to("lmdb", self.lmdb_path) + self.assertFalse(os.path.exists(self.lmdb_path)) + + def test_multisystems_input_rejected(self): + from dpdata.formats.lmdb import dump_systems + + with self.assertRaisesRegex(TypeError, "original ordered systems"): + dump_systems(dpdata.MultiSystems(self.A, self.B), self.lmdb_path) + + def test_negative_standard_atom_type_rejected(self): + from dpdata.formats.lmdb import dump_systems + + data = self.A.data.copy() + data["atom_types"] = np.array([-1, 1, 1]) + with self.assertRaisesRegex(LMDBError, "indices in"): + dump_systems([data], self.lmdb_path) + self.assertFalse(os.path.exists(self.lmdb_path)) + + def test_floating_atom_type_rejected(self): + from dpdata.formats.lmdb import dump_systems + + data = self.A.data.copy() + data["atom_types"] = np.array([0.0, 1.0, 1.0]) + with self.assertRaisesRegex(LMDBError, "integer dtype"): + dump_systems([data], self.lmdb_path) + + def test_uint64_overflow_not_treated_as_virtual_atom(self): + from dpdata.formats.lmdb import dump_systems + + data = { + "atom_numbs": [1], + "atom_names": ["MIXED_TOKEN"], + "atom_types": np.array([0]), + "real_atom_names": ["H"], + "real_atom_types": np.array( + [[np.iinfo(np.uint64).max]], dtype=np.uint64 + ), + "orig": np.zeros(3), + "cells": np.eye(3)[None], + "coords": np.zeros((1, 1, 3)), + } + with self.assertRaisesRegex(LMDBError, "indices in"): + dump_systems([data], self.lmdb_path) + + def test_inconsistent_atom_numbs_rejected(self): + from dpdata.formats.lmdb import dump_systems + + data = self.A.data.copy() + data["atom_numbs"] = [2, 1] + with self.assertRaisesRegex(LMDBError, "inconsistent"): + dump_systems([data], self.lmdb_path) + + def test_nonportable_array_dtypes_rejected(self): + from dpdata.formats.lmdb import dump_systems + + for dtype in ( + object, + np.dtype([("x", np.float64)]), + ): + with self.subTest(dtype=dtype): + data = self.A.data.copy() + data["coords"] = np.zeros( + self.A["coords"].shape, dtype=dtype + ) + with self.assertRaisesRegex( + LMDBError, "not a portable raw-byte dtype" + ): + dump_systems([data], self.lmdb_path) + self.assertFalse(os.path.exists(self.lmdb_path)) + + def test_mixed_raw_virtual_atoms_are_removed(self): + from dpdata.formats.lmdb import dump_systems + + data = { + "atom_numbs": [3], + "atom_names": ["MIXED_TOKEN"], + "atom_types": np.zeros(3, dtype=int), + "real_atom_names": ["H", "O"], + "real_atom_types": np.array([[1, 0, -1]], dtype=int), + "orig": np.zeros(3), + "cells": np.eye(3)[None], + "coords": np.arange(9, dtype=float).reshape(1, 3, 3), + "energies": np.array([-1.0]), + "forces": np.arange(9, dtype=float).reshape(1, 3, 3), + } + dump_systems([data], self.lmdb_path) + with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: + with env.begin() as txn: + meta = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) + frame = msgpack.unpackb( + txn.get(b"000000000000"), raw=False + ) + self.assertEqual(meta["type_map"], ["H", "O"]) + self.assertEqual(meta["frame_nlocs"], [2]) + self.assertEqual(frame["atom_types"]["shape"], [2]) + self.assertEqual(frame["coords"]["shape"], [2, 3]) + self.assertEqual(frame["forces"]["shape"], [2, 3]) + + def test_ambiguous_mixed_custom_field_rejected(self): + from dpdata.formats.lmdb import dump_systems + + data = { + "atom_numbs": [3], + "atom_names": ["MIXED_TOKEN"], + "atom_types": np.zeros(3, dtype=int), + "real_atom_names": ["H", "O"], + "real_atom_types": np.array([[1, 0, -1]], dtype=int), + "orig": np.zeros(3), + "cells": np.eye(3)[None], + "coords": np.zeros((1, 3, 3)), + "mystery": np.zeros((1, 3, 2)), + } + with self.assertRaisesRegex( + LMDBError, "shape is ambiguous" + ): + dump_systems([data], self.lmdb_path) + + def test_existing_destination_requires_overwrite(self): + from dpdata.formats.lmdb import dump_systems + + dump_systems([self.A], self.lmdb_path) + with self.assertRaises(FileExistsError): + dump_systems([self.B], self.lmdb_path) + + def test_failed_overwrite_preserves_old_database(self): + from dpdata.formats.lmdb import dump_systems + + old = _make_labeled_system([0], ["H"], 1) + old.data["energies"][:] = 10.0 + dump_systems([old], self.lmdb_path) + + replacement = _make_labeled_system([0], ["H"], 1) + replacement.data["energies"][:] = 20.0 + invalid = _make_labeled_system([0], ["H"], 1).data.copy() + invalid["atom_types"] = np.array([-1]) + with self.assertRaises(LMDBError): + dump_systems( + [replacement, invalid], + self.lmdb_path, + type_map=["H"], + overwrite=True, + write_batch_size=1, + ) + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + np.testing.assert_array_equal(loaded["energies"], [10.0]) + temporary = list( + Path(self.lmdb_path).parent.glob( + f".{Path(self.lmdb_path).name}.tmp-*" + ) + ) + self.assertEqual(temporary, []) + + def test_successful_overwrite_replaces_database(self): + from dpdata.formats.lmdb import dump_systems + + dump_systems([self.A], self.lmdb_path) + dump_systems( + [self.B], + self.lmdb_path, + overwrite=True, + write_batch_size=1, + ) + loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + self.assertEqual(loaded.get_nframes(), 1) + self.assertEqual(loaded[0].get_natoms(), 5) + with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: + with env.begin() as txn: + self.assertEqual(txn.stat()["entries"], 2) + self.assertIsNone(txn.get(b"000000000001")) + + def test_duplicate_type_map_rejected(self): + from dpdata.formats.lmdb import dump_systems + + with self.assertRaisesRegex(LMDBError, "duplicate"): + dump_systems( + [self.A], + self.lmdb_path, + type_map=["H", "H", "O"], + ) + + def test_invalid_write_batch_size_rejected_without_artifact(self): + from dpdata.formats.lmdb import dump_systems + + with self.assertRaisesRegex(ValueError, "write_batch_size"): + dump_systems( + [self.A], + self.lmdb_path, + write_batch_size=0, + ) + self.assertFalse(os.path.exists(self.lmdb_path)) + + def test_map_full_grows_and_retries_batch(self): + from dpdata.formats.lmdb import dump_systems + + large = _make_labeled_system([0] * 128, ["H"], 8) + dump_systems( + [large], + self.lmdb_path, + map_size=4096, + write_batch_size=3, + ) + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + self.assertEqual(loaded.get_nframes(), 8) + with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: + self.assertGreater(env.info()["map_size"], 4096) + + def test_active_dpdata_reader_blocks_overwrite_for_path_alias(self): + from dpdata.formats.lmdb import dump_systems + + dump_systems([self.A], self.lmdb_path) + resolved, _ = _open_read_env(str(Path(self.lmdb_path).resolve())) + try: + with self.assertRaisesRegex( + LMDBError, "reader is active" + ): + dump_systems( + [self.B], + self.lmdb_path, + overwrite=True, + ) + finally: + _close_read_env(resolved) + loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + self.assertEqual(loaded.get_nframes(), self.A.get_nframes()) + + def test_external_lmdb_reader_blocks_overwrite(self): + from dpdata.formats.lmdb import dump_systems + + dump_systems([self.A], self.lmdb_path) + external_env = lmdb.open( + self.lmdb_path, + readonly=True, + lock=False, + readahead=False, + meminit=False, + ) + try: + with self.assertRaisesRegex( + LMDBError, "Close DeePMD-kit and other readers" + ): + dump_systems( + [self.B], + self.lmdb_path, + overwrite=True, + ) + finally: + external_env.close() + loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + self.assertEqual(loaded.get_nframes(), self.A.get_nframes()) + + def test_publish_guard_closes_external_reader_race(self): + from dpdata.formats.lmdb import dump_systems + + dump_systems([self.A], self.lmdb_path) + replace_entered = threading.Event() + reader_finished = threading.Event() + reader_was_blocked: list[bool] = [] + real_replace = os.replace + + def concurrent_reader(): + replace_entered.wait(timeout=5) + try: + env = lmdb.open( + self.lmdb_path, + readonly=True, + lock=False, + readahead=False, + meminit=False, + ) + except lmdb.Error: + reader_was_blocked.append(True) + else: + reader_was_blocked.append(False) + env.close() + finally: + reader_finished.set() + + def synchronized_replace(source, destination): + replace_entered.set() + if not reader_finished.wait(timeout=5): + raise TimeoutError("concurrent reader did not finish") + return real_replace(source, destination) + + thread = threading.Thread(target=concurrent_reader) + thread.start() + try: + with mock.patch( + "dpdata.formats.lmdb.format.os.replace", + side_effect=synchronized_replace, + ): + dump_systems( + [self.B], + self.lmdb_path, + overwrite=True, + ) + finally: + thread.join(timeout=5) + self.assertEqual(reader_was_blocked, [True]) + loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + self.assertEqual(loaded[0].get_natoms(), self.B.get_natoms()) + + def test_staged_validation_failure_preserves_old_database(self): + from dpdata.formats.lmdb import dump_systems + from dpdata.formats.lmdb.format import _LMDBWriter + + dump_systems([self.A], self.lmdb_path) + with ( + mock.patch.object( + _LMDBWriter, + "_validate_staged_database", + side_effect=LMDBError("injected validation failure"), + ), + self.assertRaisesRegex(LMDBError, "injected"), + ): + dump_systems( + [self.B], + self.lmdb_path, + overwrite=True, + ) + loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + self.assertEqual(loaded.get_nframes(), self.A.get_nframes()) + + @unittest.skipUnless(hasattr(os, "fork"), "requires os.fork") + def test_read_cache_resets_after_fork(self): + from dpdata.formats.lmdb import dump_systems + + dump_systems([self.A], self.lmdb_path) + resolved, _ = _open_read_env(self.lmdb_path) + pid = os.fork() + if pid == 0: + try: + child_resolved, _ = _open_read_env(self.lmdb_path) + _close_read_env(child_resolved) + except BaseException: + os._exit(1) + os._exit(0) + _, status = os.waitpid(pid, 0) + _close_read_env(resolved) + self.assertEqual(os.waitstatus_to_exitcode(status), 0) + + +class TestLMDBErrorHandling(unittest.TestCase): + def setUp(self): + self.lmdb_missing_meta = "tmp_missing_meta.lmdb" + self.lmdb_missing_frame = "tmp_missing_frame.lmdb" + for p in (self.lmdb_missing_meta, self.lmdb_missing_frame): + if os.path.exists(p): + shutil.rmtree(p) + + env = lmdb.open(self.lmdb_missing_frame, map_size=1 << 30) + with env.begin(write=True) as txn: + meta = { + "nframes": 1, + "frame_idx_fmt": "012d", + "type_map": ["H", "O"], + "frame_nlocs": [3], + "frame_system_ids": [0], + } + txn.put(b"__metadata__", _reference_packb(meta)) + env.close() + + def tearDown(self): + for p in (self.lmdb_missing_meta, self.lmdb_missing_frame): + if os.path.exists(p): + shutil.rmtree(p) + + def test_missing_metadata(self): + lmdb.open(self.lmdb_missing_meta, map_size=1 << 30).close() + with self.assertRaises(LMDBMetadataError): + dpdata.MultiSystems.from_file(self.lmdb_missing_meta, fmt="lmdb") + + def test_missing_frame(self): + with self.assertRaises(LMDBFrameError): + dpdata.MultiSystems.from_file(self.lmdb_missing_frame, fmt="lmdb") + + def test_non_mapping_metadata_rejected(self): + path = "tmp_bad_metadata.lmdb" + try: + env = lmdb.open(path, map_size=1 << 30) + with env.begin(write=True) as txn: + txn.put(b"__metadata__", _reference_packb(["not", "a", "mapping"])) + env.close() + with self.assertRaisesRegex( + LMDBMetadataError, "must contain a mapping" + ): + dpdata.MultiSystems.from_file(path, fmt="lmdb") + finally: + if os.path.exists(path): + shutil.rmtree(path) + + def test_invalid_msgpack_metadata_wrapped(self): + path = "tmp_invalid_msgpack_metadata.lmdb" + try: + env = lmdb.open(path, map_size=1 << 30) + with env.begin(write=True) as txn: + txn.put(b"__metadata__", b"\xc1") + env.close() + with self.assertRaisesRegex( + LMDBMetadataError, "Cannot decode __metadata__" + ): + dpdata.MultiSystems.from_file(path, fmt="lmdb") + finally: + if os.path.exists(path): + shutil.rmtree(path) + + def test_fractional_nframes_rejected(self): + path = "tmp_fractional_nframes.lmdb" + try: + env = lmdb.open(path, map_size=1 << 30) + with env.begin(write=True) as txn: + txn.put( + b"__metadata__", + _reference_packb( + { + "nframes": 1.9, + "frame_idx_fmt": "012d", + "type_map": ["H"], + } + ), + ) + env.close() + with self.assertRaisesRegex( + LMDBMetadataError, "nframes must be an integer" + ): + dpdata.MultiSystems.from_file(path, fmt="lmdb") + finally: + if os.path.exists(path): + shutil.rmtree(path) + + def test_non_integer_metadata_vectors_rejected(self): + for key, value in ( + ("frame_nlocs", [1.5]), + ("frame_system_ids", [False]), + ): + with self.subTest(key=key): + path = f"tmp_invalid_{key}.lmdb" + try: + frames = [ + { + "atom_types": np.array([0], dtype=np.int32), + "coords": np.zeros((1, 3)), + } + ] + _write_raw_lmdb( + path, + frames, + ["H"], + metadata_extra={key: value}, + ) + with self.assertRaisesRegex( + LMDBMetadataError, "must be an integer" + ): + dpdata.MultiSystems.from_file(path, fmt="lmdb") + finally: + if os.path.exists(path): + shutil.rmtree(path) + + def test_malformed_array_payload_rejected(self): + path = "tmp_bad_array.lmdb" + try: + env = lmdb.open(path, map_size=1 << 30) + with env.begin(write=True) as txn: + frame = { + "atom_types": { + "type": "int32", + "shape": [2], + "data": b"\x00", + }, + "atom_numbs": [2], + } + txn.put(b"000000000000", _reference_packb(frame)) + txn.put( + b"__metadata__", + _reference_packb( + { + "nframes": 1, + "frame_idx_fmt": "012d", + "type_map": ["H"], + "frame_nlocs": [2], + } + ), + ) + env.close() + with self.assertRaisesRegex( + LMDBFrameError, "Cannot decode array" + ): + dpdata.MultiSystems.from_file(path, fmt="lmdb") + finally: + if os.path.exists(path): + shutil.rmtree(path) + + def test_invalid_frame_system_ids_length(self): + path = "tmp_invalid_system_ids.lmdb" + try: + system = _make_labeled_system([0], ["H"], 1) + system.to("lmdb", path) + env = lmdb.open(path, map_size=1 << 30) + with env.begin(write=True) as txn: + metadata = msgpack.unpackb( + txn.get(b"__metadata__"), raw=False + ) + metadata["frame_system_ids"] = [] + txn.put( + b"__metadata__", + _reference_packb(metadata), + ) + env.close() + with self.assertRaisesRegex( + LMDBMetadataError, "frame_system_ids length" + ): + dpdata.MultiSystems.from_file(path, fmt="lmdb") + finally: + if os.path.exists(path): + shutil.rmtree(path) + + def test_existing_external_reader_has_actionable_error(self): + path = "tmp_external_reader.lmdb" + try: + _make_labeled_system([0], ["H"], 1).to("lmdb", path) + env = lmdb.open(path, readonly=True, lock=False) + try: + with self.assertRaisesRegex( + LMDBError, "another library has already opened" + ): + dpdata.MultiSystems.from_file(path, fmt="lmdb") + finally: + env.close() + finally: + if os.path.exists(path): + shutil.rmtree(path) class TestLMDBConfig(unittest.TestCase): @@ -149,25 +1488,23 @@ def tearDown(self): shutil.rmtree(self.lmdb_path) def test_custom_frame_idx_fmt(self): - fmt = "06d" - # Standard API with custom kwarg ms = dpdata.MultiSystems(self.system) - ms.to("lmdb", self.lmdb_path, frame_idx_fmt=fmt) - - # 1. Verify key format in database - with lmdb.open(self.lmdb_path, readonly=True) as env: + ms.to("lmdb", self.lmdb_path, frame_idx_fmt="06d") + with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: with env.begin() as txn: - # Frame 0 should be "000000" self.assertIsNotNone(txn.get(b"000000")) - # Frame 0 should NOT be "000000000000" self.assertIsNone(txn.get(b"000000000000")) + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + self.assertEqual(len(loaded), len(self.system)) + np.testing.assert_allclose(loaded.data["coords"], self.system.data["coords"]) - # 2. Verify loading works automatically via standard API - system_loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") - self.assertEqual(len(system_loaded), len(self.system)) - np.testing.assert_allclose( - system_loaded.data["coords"], self.system.data["coords"] - ) + def test_format_annotations_resolve_at_runtime(self): + from dpdata.format import Format + + to_hints = typing.get_type_hints(Format.to_multi_systems) + from_hints = typing.get_type_hints(Format.from_multi_systems) + self.assertIn("return", to_hints) + self.assertIn("return", from_hints) if __name__ == "__main__": diff --git a/tests/test_lmdb_custom_dtype.py b/tests/test_lmdb_custom_dtype.py index 4786a30c7..f2bf1a170 100644 --- a/tests/test_lmdb_custom_dtype.py +++ b/tests/test_lmdb_custom_dtype.py @@ -4,30 +4,30 @@ import shutil import unittest +import lmdb +import msgpack import numpy as np import dpdata from dpdata.data_type import Axis, DataType +from dpdata.formats.lmdb.format import LMDBError, LMDBFrameError -class TestLMDBCustomDType(unittest.TestCase): +class TestLMDBFrameData(unittest.TestCase): + """Frame-dependent custom data must round-trip through LMDB.""" + def setUp(self): self.original_system_dtypes = dpdata.System.DTYPES self.original_labeled_system_dtypes = dpdata.LabeledSystem.DTYPES - # Register custom data types as optional self.dt_frame = DataType( "frame_data", np.ndarray, shape=(Axis.NFRAMES, 2), required=False ) - self.dt_static = DataType("static_data", np.ndarray, shape=(2,), required=False) - - dpdata.System.register_data_type(self.dt_frame, self.dt_static) - dpdata.LabeledSystem.register_data_type(self.dt_frame, self.dt_static) + dpdata.System.register_data_type(self.dt_frame) + dpdata.LabeledSystem.register_data_type(self.dt_frame) self.lmdb_path = "tmp_custom_dtype.lmdb" - # Create a system with custom data - # Assuming running from tests/ directory try: self.system = dpdata.LabeledSystem( "poscars/OUTCAR.h2o.md", fmt="vasp/outcar" @@ -39,7 +39,6 @@ def setUp(self): nframes = self.system.get_nframes() self.system.data["frame_data"] = np.random.rand(nframes, 2) - self.system.data["static_data"] = np.array([1.0, 2.0]) self.system.check_data() def tearDown(self): @@ -48,53 +47,25 @@ def tearDown(self): if os.path.exists(self.lmdb_path): shutil.rmtree(self.lmdb_path) - def test_custom_dtype_preservation(self): + def test_frame_data_preservation(self): self.system.to("lmdb", self.lmdb_path) - system_loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") - - np.testing.assert_allclose( - system_loaded.data["frame_data"], self.system.data["frame_data"] - ) - np.testing.assert_allclose( - system_loaded.data["static_data"], self.system.data["static_data"] - ) - - def test_multi_systems_custom_dtype(self): - ms = dpdata.MultiSystems(self.system) - ms.to("lmdb", self.lmdb_path) - ms_loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") - - system_loaded = list(ms_loaded.systems.values())[0] - np.testing.assert_allclose( - system_loaded.data["frame_data"], self.system.data["frame_data"] - ) + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") np.testing.assert_allclose( - system_loaded.data["static_data"], self.system.data["static_data"] + loaded.data["frame_data"], self.system.data["frame_data"] ) - def test_custom_dtype_auto_registration(self): - # Save with custom data types registered + def test_frame_data_auto_registration(self): self.system.to("lmdb", self.lmdb_path) - # Simulate a clean session by unregistering the custom types + # simulate a clean session dpdata.System.DTYPES = self.original_system_dtypes dpdata.LabeledSystem.DTYPES = self.original_labeled_system_dtypes - - # Verify they are currently missing self.assertNotIn("frame_data", [dt.name for dt in dpdata.LabeledSystem.DTYPES]) - # Load from LMDB - should trigger auto-registration - system_loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") - - # Verify data is loaded and types are registered + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") self.assertIn("frame_data", [dt.name for dt in dpdata.LabeledSystem.DTYPES]) - self.assertIn("static_data", [dt.name for dt in dpdata.LabeledSystem.DTYPES]) - np.testing.assert_allclose( - system_loaded.data["frame_data"], self.system.data["frame_data"] - ) - np.testing.assert_allclose( - system_loaded.data["static_data"], self.system.data["static_data"] + loaded.data["frame_data"], self.system.data["frame_data"] ) @@ -104,12 +75,7 @@ def setUp(self): self.original_labeled_system_dtypes = dpdata.LabeledSystem.DTYPES new_datatypes = [ - DataType( - "fparam", - np.ndarray, - shape=(Axis.NFRAMES, 2), - required=False, - ), + DataType("fparam", np.ndarray, shape=(Axis.NFRAMES, 2), required=False), DataType( "aparam", np.ndarray, @@ -117,7 +83,6 @@ def setUp(self): required=False, ), ] - for datatype in new_datatypes: dpdata.System.register_data_type(datatype) dpdata.LabeledSystem.register_data_type(datatype) @@ -147,57 +112,37 @@ def tearDown(self): def test_fparam_aparam_preservation(self): self.system.to("lmdb", self.lmdb_path) - system_loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") - - np.testing.assert_allclose( - system_loaded.data["fparam"], self.system.data["fparam"] - ) - np.testing.assert_allclose( - system_loaded.data["aparam"], self.system.data["aparam"] - ) + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + np.testing.assert_allclose(loaded.data["fparam"], self.system.data["fparam"]) + np.testing.assert_allclose(loaded.data["aparam"], self.system.data["aparam"]) def test_fparam_aparam_auto_registration(self): - # Save with fparam/aparam registered self.system.to("lmdb", self.lmdb_path) - # Simulate a clean session by restoring original DTYPES dpdata.System.DTYPES = self.original_system_dtypes dpdata.LabeledSystem.DTYPES = self.original_labeled_system_dtypes - # Load from LMDB - system_loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") - - # Verify auto-registration and data correctness + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") self.assertIn("fparam", [dt.name for dt in dpdata.LabeledSystem.DTYPES]) self.assertIn("aparam", [dt.name for dt in dpdata.LabeledSystem.DTYPES]) - - np.testing.assert_allclose( - system_loaded.data["fparam"], self.system.data["fparam"] - ) - np.testing.assert_allclose( - system_loaded.data["aparam"], self.system.data["aparam"] - ) + np.testing.assert_allclose(loaded.data["fparam"], self.system.data["fparam"]) + np.testing.assert_allclose(loaded.data["aparam"], self.system.data["aparam"]) def test_symbolic_axis_natoms_preservation(self): - # 1. Save system with aparam (which uses Axis.NATOMS) + # natoms == 3 collides with the trailing coordinate dim; the stored + # symbolic shape must still recover (NFRAMES, NATOMS, 3), not + # (NFRAMES, NATOMS, NATOMS). self.system.to("lmdb", self.lmdb_path) - # 2. Simulate new session dpdata.System.DTYPES = self.original_system_dtypes dpdata.LabeledSystem.DTYPES = self.original_labeled_system_dtypes - # 3. Load triggers auto-registration dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") - - # 4. Find the newly registered DataType for 'aparam' aparam_dt = next( dt for dt in dpdata.LabeledSystem.DTYPES if dt.name == "aparam" ) + self.assertEqual(aparam_dt.shape, (Axis.NFRAMES, Axis.NATOMS, 3)) - # 5. Assert that it contains Axis.NATOMS, not a fixed integer - self.assertIn(Axis.NATOMS, aparam_dt.shape) - - # 6. Functional verification data_diff = { "atom_numbs": [5], "atom_names": ["H"], @@ -215,5 +160,264 @@ def test_symbolic_axis_natoms_preservation(self): self.fail(f"DataError raised despite symbolic NATOMS: {e}") +class TestLMDBFieldProtocol(unittest.TestCase): + def setUp(self): + self.original_system_dtypes = dpdata.System.DTYPES + self.original_labeled_system_dtypes = dpdata.LabeledSystem.DTYPES + self.lmdb_path = "tmp_field_protocol.lmdb" + try: + self.system = dpdata.LabeledSystem( + "poscars/OUTCAR.h2o.md", fmt="vasp/outcar" + ) + except FileNotFoundError: + self.system = dpdata.LabeledSystem( + "tests/poscars/OUTCAR.h2o.md", fmt="vasp/outcar" + ) + + def tearDown(self): + dpdata.System.DTYPES = self.original_system_dtypes + dpdata.LabeledSystem.DTYPES = self.original_labeled_system_dtypes + if os.path.exists(self.lmdb_path): + shutil.rmtree(self.lmdb_path) + + def _register(self, *dtypes): + dpdata.System.register_data_type(*dtypes) + dpdata.LabeledSystem.register_data_type(*dtypes) + + def test_deepmd_name_used_on_disk_and_restored(self): + spin = DataType( + "spins", + np.ndarray, + (Axis.NFRAMES, Axis.NATOMS, 3), + required=False, + deepmd_name="spin", + ) + self._register(spin) + self.system.data["spins"] = np.arange( + self.system.get_nframes() * self.system.get_natoms() * 3, + dtype=float, + ).reshape(self.system.get_nframes(), self.system.get_natoms(), 3) + expected = self.system.data["spins"].copy() + self.system.to("lmdb", self.lmdb_path) + + with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: + with env.begin() as txn: + metadata = msgpack.unpackb( + txn.get(b"__metadata__"), raw=False + ) + frame = msgpack.unpackb( + txn.get(b"000000000000"), raw=False + ) + self.assertIn("spin", frame) + self.assertNotIn("spins", frame) + self.assertEqual(metadata["dp_data_names"], {"spin": "spins"}) + + dpdata.System.DTYPES = self.original_system_dtypes + dpdata.LabeledSystem.DTYPES = self.original_labeled_system_dtypes + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + np.testing.assert_array_equal(loaded["spins"], expected) + registered = next( + dt for dt in dpdata.LabeledSystem.DTYPES if dt.name == "spins" + ) + self.assertEqual(registered.deepmd_name, "spin") + + def test_duplicate_deepmd_name_rejected(self): + first = DataType( + "first_field", + np.ndarray, + (Axis.NFRAMES, 1), + required=False, + deepmd_name="shared", + ) + second = DataType( + "second_field", + np.ndarray, + (Axis.NFRAMES, 1), + required=False, + deepmd_name="shared", + ) + self._register(first, second) + nframes = self.system.get_nframes() + self.system.data["first_field"] = np.zeros((nframes, 1)) + self.system.data["second_field"] = np.ones((nframes, 1)) + with self.assertRaisesRegex( + LMDBError, + "both map to", + ): + self.system.to("lmdb", self.lmdb_path) + + def test_deepmd_core_name_collision_rejected(self): + malicious = DataType( + "malicious_coords", + np.ndarray, + (Axis.NFRAMES, Axis.NATOMS, 3), + required=False, + deepmd_name="coord", + ) + self._register(malicious) + self.system.data["malicious_coords"] = np.full( + ( + self.system.get_nframes(), + self.system.get_natoms(), + 3, + ), + 101.0, + ) + with self.assertRaisesRegex(LMDBError, "reserved LMDB key"): + self.system.to("lmdb", self.lmdb_path) + + def test_additional_protocol_alias_rejected(self): + alias = DataType( + "magmom_alias", + np.ndarray, + (Axis.NFRAMES, Axis.NATOMS, 3), + required=False, + deepmd_name="spin", + ) + self._register(alias) + self.system.data["magmom_alias"] = np.zeros( + ( + self.system.get_nframes(), + self.system.get_natoms(), + 3, + ) + ) + with self.assertRaisesRegex( + LMDBError, "belongs to 'spins'" + ): + self.system.to("lmdb", self.lmdb_path) + + def test_atom_types_key_collision_rejected(self): + malicious = DataType( + "malicious_types", + np.ndarray, + (Axis.NFRAMES, Axis.NATOMS), + required=False, + deepmd_name="atom_types", + ) + self._register(malicious) + self.system.data["malicious_types"] = np.zeros( + (self.system.get_nframes(), self.system.get_natoms()) + ) + with self.assertRaisesRegex(LMDBError, "reserved LMDB key"): + self.system.to("lmdb", self.lmdb_path) + + def test_static_field_roundtrip(self): + static = DataType( + "static_data", np.ndarray, shape=(2,), required=False + ) + self._register(static) + self.system.data["static_data"] = np.array([1.0, 2.0]) + self.system.to("lmdb", self.lmdb_path) + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + np.testing.assert_array_equal(loaded["static_data"], [1.0, 2.0]) + + def test_inconsistent_static_field_raises(self): + static = DataType( + "static_data", np.ndarray, shape=(2,), required=False + ) + self._register(static) + self.system.data["static_data"] = np.array([1.0, 2.0]) + self.system.to("lmdb", self.lmdb_path) + + env = lmdb.open(self.lmdb_path, map_size=1 << 30) + with env.begin(write=True) as txn: + key = b"000000000001" + frame = msgpack.unpackb(txn.get(key), raw=False) + replacement = np.array([3.0, 4.0]) + frame["static_data"] = { + "type": str(replacement.dtype), + "shape": list(replacement.shape), + "data": replacement.tobytes(), + } + txn.put(key, msgpack.packb(frame, use_bin_type=True)) + env.close() + + with self.assertRaisesRegex( + LMDBFrameError, + "Static field", + ): + dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + + def test_nonleading_frame_axis_roundtrip(self): + transposed = DataType( + "transposed_frames", + np.ndarray, + shape=(2, Axis.NFRAMES), + required=False, + ) + self._register(transposed) + values = np.arange( + 2 * self.system.get_nframes(), dtype=float + ).reshape(2, self.system.get_nframes()) + self.system.data["transposed_frames"] = values + self.system.to("lmdb", self.lmdb_path) + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + np.testing.assert_array_equal(loaded["transposed_frames"], values) + + def test_unused_shape_none_data_type_does_not_block_write(self): + undefined = DataType( + "unused_undefined_shape", + np.ndarray, + shape=None, + required=False, + ) + self._register(undefined) + self.system.to("lmdb", self.lmdb_path) + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + self.assertEqual(loaded.get_nframes(), self.system.get_nframes()) + + def test_used_shape_none_data_type_rejected(self): + undefined = DataType( + "used_undefined_shape", + np.ndarray, + shape=None, + required=False, + ) + self._register(undefined) + self.system.data["used_undefined_shape"] = np.zeros( + (self.system.get_nframes(), 2) + ) + with self.assertRaisesRegex(LMDBError, "no declared shape"): + self.system.to("lmdb", self.lmdb_path) + + def test_multiple_atom_axes_remove_virtual_atoms(self): + hessian = DataType( + "hessian", + np.ndarray, + ( + Axis.NFRAMES, + Axis.NATOMS, + 3, + Axis.NATOMS, + 3, + ), + required=False, + deepmd_name="hessian", + ) + self._register(hessian) + data = { + "atom_numbs": [3], + "atom_names": ["MIXED_TOKEN"], + "atom_types": np.zeros(3, dtype=int), + "real_atom_names": ["H", "O"], + "real_atom_types": np.array([[1, 0, -1]], dtype=int), + "orig": np.zeros(3), + "cells": np.eye(3)[None], + "coords": np.zeros((1, 3, 3)), + "energies": np.array([-1.0]), + "hessian": np.zeros((1, 3, 3, 3, 3)), + } + from dpdata.formats.lmdb import dump_systems + + dump_systems([data], self.lmdb_path) + with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: + with env.begin() as txn: + frame = msgpack.unpackb( + txn.get(b"000000000000"), raw=False + ) + self.assertEqual(frame["hessian"]["shape"], [2, 3, 2, 3]) + + if __name__ == "__main__": unittest.main() From 5c34ca87b00eb4d8454387d10cef550871e997bb Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 12 Jul 2026 11:28:35 +0000 Subject: [PATCH 2/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/systems/lmdb.md | 19 ++-- dpdata/formats/lmdb/format.py | 174 +++++++++++--------------------- tests/test_lmdb.py | 144 +++++++------------------- tests/test_lmdb_custom_dtype.py | 30 ++---- 4 files changed, 114 insertions(+), 253 deletions(-) diff --git a/docs/systems/lmdb.md b/docs/systems/lmdb.md index 3a9399970..1259e5528 100644 --- a/docs/systems/lmdb.md +++ b/docs/systems/lmdb.md @@ -10,15 +10,15 @@ A database contains one metadata record together with one record per frame. The metadata record is stored under the key `__metadata__` and holds the following entries. -| key | description | -| ------------------ | ---------------------------------------------------- | -| `nframes` | total number of frames | -| `frame_idx_fmt` | format string of the integer frame key (default `012d`) | -| `type_map` | the global element table | -| `frame_nlocs` | the number of atoms of each frame | -| `frame_system_ids` | the index of the source system of each frame | -| `dp_data_shapes` | optional symbolic shapes of additional dpdata fields | -| `dp_data_names` | optional mapping from DeePMD-kit field names to dpdata field names | +| key | description | +| ------------------ | ------------------------------------------------------------------ | +| `nframes` | total number of frames | +| `frame_idx_fmt` | format string of the integer frame key (default `012d`) | +| `type_map` | the global element table | +| `frame_nlocs` | the number of atoms of each frame | +| `frame_system_ids` | the index of the source system of each frame | +| `dp_data_shapes` | optional symbolic shapes of additional dpdata fields | +| `dp_data_names` | optional mapping from DeePMD-kit field names to dpdata field names | Each frame is stored under its zero-padded global index (`000000000000`, `000000000001`, ...). The per-frame arrays `coords`, `cells`, `energies`, `forces`, `virials`, and `atom_types` are stored together with their data type and shape; `atom_types` records global indices into `type_map`, and `atom_numbs` is a list of per-element counts over `type_map`. The numerical dtype of each source array is retained, except that `atom_types` is stored as `int32`. @@ -69,6 +69,7 @@ def load(): for d in directories: yield dpdata.LabeledSystem(d, fmt="deepmd/npy") + dump_systems(load(), "data.lmdb", type_map=["H", "C", "N", "O"]) ``` diff --git a/dpdata/formats/lmdb/format.py b/dpdata/formats/lmdb/format.py index fe109b6cf..991fab820 100644 --- a/dpdata/formats/lmdb/format.py +++ b/dpdata/formats/lmdb/format.py @@ -126,11 +126,7 @@ class _FieldSpec: @property def frame_axis(self) -> int | None: """Return the frame axis in the in-memory array.""" - return ( - self.shape.index(Axis.NFRAMES) - if Axis.NFRAMES in self.shape - else None - ) + return self.shape.index(Axis.NFRAMES) if Axis.NFRAMES in self.shape else None @property def atom_axes(self) -> tuple[int, ...]: @@ -227,9 +223,7 @@ class _AtomSource: *(dtype.name for dtype in _PROTOCOL_DTYPES), } ) -_PROTOCOL_DISK_NAMES = { - dtype.deepmd_name: dtype.name for dtype in _PROTOCOL_DTYPES -} +_PROTOCOL_DISK_NAMES = {dtype.deepmd_name: dtype.name for dtype in _PROTOCOL_DTYPES} _READ_ENV_CACHE: dict[str, tuple[lmdb.Environment, int]] = {} _READ_ENV_LOCK = threading.Lock() @@ -280,9 +274,7 @@ def _encode_array(arr: np.ndarray) -> dict[str, Any]: """ a = np.asarray(arr) if a.dtype.hasobject or a.dtype.fields is not None: - raise LMDBError( - f"LMDB array dtype {a.dtype} is not a portable raw-byte dtype." - ) + raise LMDBError(f"LMDB array dtype {a.dtype} is not a portable raw-byte dtype.") try: restored_dtype = np.dtype(str(a.dtype)) except TypeError as exc: @@ -355,9 +347,7 @@ def _decode_frame(raw: bytes) -> dict[str, Any]: try: out[name] = _decode_array(val) if _is_encoded_array(val) else val except (TypeError, ValueError) as exc: - raise LMDBFrameError( - f"Cannot decode array field '{name}': {exc}" - ) from exc + raise LMDBFrameError(f"Cannot decode array field '{name}': {exc}") from exc return out @@ -397,13 +387,9 @@ def _packb(value: Any) -> bytes: return packed -def _require_integer( - value: Any, label: str, *, minimum: int -) -> int: +def _require_integer(value: Any, label: str, *, minimum: int) -> int: """Return a strictly validated integer metadata value.""" - if isinstance(value, (bool, np.bool_)) or not isinstance( - value, numbers.Integral - ): + if isinstance(value, (bool, np.bool_)) or not isinstance(value, numbers.Integral): raise LMDBMetadataError(f"{label} must be an integer.") result = int(value) if result < minimum: @@ -444,14 +430,11 @@ def _validate_field_namespace(dp_name: str, disk_name: str) -> None: ) if disk_name in _FORBIDDEN_CUSTOM_DISK_NAMES or disk_name.startswith("find_"): raise LMDBError( - f"Custom field '{dp_name}' cannot use reserved LMDB key " - f"'{disk_name}'." + f"Custom field '{dp_name}' cannot use reserved LMDB key '{disk_name}'." ) -def _field_spec( - dtype: DataType, *, validate_namespace: bool = True -) -> _FieldSpec: +def _field_spec(dtype: DataType, *, validate_namespace: bool = True) -> _FieldSpec: """Build a field specification from a registered data type.""" if dtype.shape is None: raise LMDBError( @@ -471,9 +454,7 @@ def _field_spec( def _field_registries( *, prefer_protocol: bool = False -) -> tuple[ - dict[str, _FieldSpec], dict[str, _FieldSpec] -]: +) -> tuple[dict[str, _FieldSpec], dict[str, _FieldSpec]]: """Return field specifications indexed by dpdata and disk names.""" data_types = _all_data_types() if prefer_protocol: @@ -519,9 +500,8 @@ def _validate_integer_types( ) -> np.ndarray: """Validate an atom-type array without silently truncating values.""" array = np.asarray(values) - if ( - np.issubdtype(array.dtype, np.bool_) - or not np.issubdtype(array.dtype, np.integer) + if np.issubdtype(array.dtype, np.bool_) or not np.issubdtype( + array.dtype, np.integer ): raise LMDBError(f"{name} must use an integer dtype, got {array.dtype}.") lower_bound = -1 if allow_virtual else 0 @@ -650,9 +630,7 @@ def _open_read_env(path: str) -> tuple[str, lmdb.Environment]: resolved = _normalized_path(path) with _READ_ENV_LOCK: if resolved in _PUBLISHING_PATHS: - raise LMDBError( - f"LMDB '{path}' is being replaced and cannot be opened." - ) + raise LMDBError(f"LMDB '{path}' is being replaced and cannot be opened.") cached = _READ_ENV_CACHE.get(resolved) if cached is not None: env, count = cached @@ -910,9 +888,7 @@ def _put_records(self, records: list[tuple[bytes, bytes]]) -> None: return except lmdb.MapFullError: current_size = int(self.env.info()["map_size"]) - self.env.set_mapsize( - max(current_size * 2, current_size + (1 << 20)) - ) + self.env.set_mapsize(max(current_size * 2, current_size + (1 << 20))) def write_system(self, data: dict[str, Any]) -> None: """Append one validated source system.""" @@ -941,8 +917,8 @@ def write_system(self, data: dict[str, Any]) -> None: [name_to_global.get(name, -1) for name in atom_source.names], dtype=np.int64, ) - has_virtual_atoms = ( - atom_source.masks is not None and not np.all(atom_source.masks) + has_virtual_atoms = atom_source.masks is not None and not np.all( + atom_source.masks ) specs = self._build_specs( data, @@ -975,30 +951,24 @@ def write_system(self, data: dict[str, Any]) -> None: f"{list(type_map)}." ) global_types = global_types.astype(np.int32) - frame: dict[str, Any] = { - "atom_types": _encode_array(global_types) - } + frame: dict[str, Any] = {"atom_types": _encode_array(global_types)} for spec, array in fields: if spec.dp_name == "cells" and nopbc: continue - payload = self._frame_payload( - array, spec, frame_i, atom_mask - ) + payload = self._frame_payload(array, spec, frame_i, atom_mask) frame[spec.disk_name] = _encode_array(payload) - counts = np.bincount( - global_types, minlength=len(type_map) - )[: len(type_map)] + counts = np.bincount(global_types, minlength=len(type_map))[ + : len(type_map) + ] frame["atom_numbs"] = [int(count) for count in counts] - key = format( - self.frame_index + offset, self.frame_idx_fmt - ).encode("ascii") + key = format(self.frame_index + offset, self.frame_idx_fmt).encode( + "ascii" + ) records.append((key, _packb(frame))) batch_nlocs.append(int(global_types.size)) self._put_records(records) self.frame_nlocs.extend(batch_nlocs) - self.frame_system_ids.extend( - [self.system_index] * len(records) - ) + self.frame_system_ids.extend([self.system_index] * len(records)) self.frame_index += len(records) self.system_index += 1 @@ -1051,11 +1021,11 @@ def _validate_staged_database(self) -> None: ) metadata = _read_metadata(txn) if _meta_get(metadata, "nframes") != self.frame_index: - raise LMDBError("Staged LMDB nframes does not match written frames.") - if len(_meta_get(metadata, "frame_nlocs") or []) != self.frame_index: raise LMDBError( - "Staged LMDB frame_nlocs length is inconsistent." + "Staged LMDB nframes does not match written frames." ) + if len(_meta_get(metadata, "frame_nlocs") or []) != self.frame_index: + raise LMDBError("Staged LMDB frame_nlocs length is inconsistent.") if ( len(_meta_get(metadata, "frame_system_ids") or []) != self.frame_index @@ -1066,9 +1036,7 @@ def _validate_staged_database(self) -> None: for index in range(self.frame_index): key = format(index, self.frame_idx_fmt).encode("ascii") if txn.get(key) is None: - raise LMDBError( - f"Staged LMDB is missing frame key {key!r}." - ) + raise LMDBError(f"Staged LMDB is missing frame key {key!r}.") finally: env.close() @@ -1373,10 +1341,7 @@ def from_multi_systems( if file_type_map: names = list( _validate_names( - [ - n.decode() if isinstance(n, bytes) else n - for n in file_type_map - ], + [n.decode() if isinstance(n, bytes) else n for n in file_type_map], "file type_map", ) ) @@ -1570,7 +1535,9 @@ def _infer_type_names(frames: list[dict[str, Any]]) -> list[str]: for frame in frames: raw = np.asarray(frame["atom_types"]) if not np.issubdtype(raw.dtype, np.integer) or raw.ndim != 1: - raise LMDBFrameError("atom_types must be a one-dimensional integer array.") + raise LMDBFrameError( + "atom_types must be a one-dimensional integer array." + ) if raw.size and int(raw.min()) < 0: raise LMDBFrameError("atom_types must not contain negative indices.") if raw.size: @@ -1615,16 +1582,13 @@ def _validate_frame_types( frame["atom_numbs"] = expected @staticmethod - def _metadata_mapping( - metadata: dict[str | bytes, Any], key: str - ) -> dict[str, Any]: + def _metadata_mapping(metadata: dict[str | bytes, Any], key: str) -> dict[str, Any]: """Return a metadata mapping with decoded string keys.""" raw = _meta_get(metadata, key) or {} if not isinstance(raw, dict): raise LMDBMetadataError(f"{key} must be a mapping.") return { - k.decode() if isinstance(k, bytes) else str(k): v - for k, v in raw.items() + k.decode() if isinstance(k, bytes) else str(k): v for k, v in raw.items() } def _resolve_read_specs( @@ -1649,19 +1613,22 @@ def _resolve_read_specs( if isinstance(hinted_name, bytes): hinted_name = hinted_name.decode() known = by_disk.get(disk_name) - dp_name = str(hinted_name) if hinted_name is not None else ( - known.dp_name if known is not None else disk_name + dp_name = ( + str(hinted_name) + if hinted_name is not None + else (known.dp_name if known is not None else disk_name) ) - if known is not None and hinted_name is not None and dp_name != known.dp_name: + if ( + known is not None + and hinted_name is not None + and dp_name != known.dp_name + ): raise LMDBMetadataError( f"dp_data_names maps '{disk_name}' to '{dp_name}', but the " f"registered protocol maps it to '{known.dp_name}'." ) expected_spec = by_dp.get(dp_name) - if ( - expected_spec is not None - and disk_name != expected_spec.disk_name - ): + if expected_spec is not None and disk_name != expected_spec.disk_name: raise LMDBMetadataError( f"LMDB key '{disk_name}' cannot map to registered field " f"'{dp_name}', whose protocol key is " @@ -1681,18 +1648,14 @@ def _resolve_read_specs( if disk_name in shape_hints: shape = _tokens_to_shape(shape_hints[disk_name]) - if ( - known is not None - and not self._shape_hint_compatible(known.shape, shape) + if known is not None and not self._shape_hint_compatible( + known.shape, shape ): raise LMDBMetadataError( f"dp_data_shapes for known field '{disk_name}' changes " f"its protocol shape from {known.shape} to {shape}." ) - elif ( - known is not None - and known.dp_name in _PROTOCOL_FIELD_NAMES - ): + elif known is not None and known.dp_name in _PROTOCOL_FIELD_NAMES: shape = known.shape else: shape = self._infer_field_shape(disk_name, frames) @@ -1752,8 +1715,7 @@ def _infer_field_shape( for axis in range(rank): dimensions = [shape[axis] for shape, _ in observed] follows_natoms = all( - dim == natoms - for dim, (_, natoms) in zip(dimensions, observed) + dim == natoms for dim, (_, natoms) in zip(dimensions, observed) ) if follows_natoms: if len({natoms for _, natoms in observed}) == 1: @@ -1794,9 +1756,7 @@ def _rename_frame_fields( f"'{spec.dp_name}'." ) value = np.asarray(frame[disk_name]) - LMDBFormat._validate_payload_shape( - spec, value, natoms, frame_index - ) + LMDBFormat._validate_payload_shape(spec, value, natoms, frame_index) renamed[spec.dp_name] = value renamed_frames.append(renamed) return renamed_frames @@ -1844,9 +1804,7 @@ def _group_frames( canonical = self._canonicalize_frame(frame, specs) canonical_frames.append(canonical) composition = tuple( - np.bincount( - canonical["atom_types"], minlength=len(names) - ).tolist() + np.bincount(canonical["atom_types"], minlength=len(names)).tolist() ) groups.setdefault(composition, []).append(index) @@ -1929,9 +1887,7 @@ def _aggregate_group( data["cells"] = np.zeros((len(indices), 3, 3)) @staticmethod - def _aggregate_values( - values: list[np.ndarray], spec: _FieldSpec - ) -> np.ndarray: + def _aggregate_values(values: list[np.ndarray], spec: _FieldSpec) -> np.ndarray: """Aggregate values while preserving byte order and static semantics.""" first = values[0] for value in values[1:]: @@ -1942,7 +1898,9 @@ def _aggregate_values( ) frame_axis = spec.frame_axis if frame_axis is None: - if any(not np.array_equal(first, value, equal_nan=True) for value in values[1:]): + if any( + not np.array_equal(first, value, equal_nan=True) for value in values[1:] + ): raise LMDBFrameError( f"Static field '{spec.disk_name}' differs between frames." ) @@ -1965,12 +1923,7 @@ def _register_specs( existing = {dt.name: dt for dt in (*System.DTYPES, *LabeledSystem.DTYPES)} protocol_dtypes = {dtype.name: dtype for dtype in _PROTOCOL_DTYPES} - used_names = { - name - for data in datasets - for name in data - if name in specs - } + used_names = {name for data in datasets for name in data if name in specs} pending: list[DataType] = [] for name in sorted(used_names): if name not in specs: @@ -1978,9 +1931,7 @@ def _register_specs( spec = specs[name] current = existing.get(name) if current is not None: - merged_shape = LMDBFormat._merge_shapes( - current.shape, spec.shape - ) + merged_shape = LMDBFormat._merge_shapes(current.shape, spec.shape) protocol_dtype = protocol_dtypes.get(name) if ( merged_shape is None @@ -1988,15 +1939,10 @@ def _register_specs( and LMDBFormat._shapes_compatible( current.shape, protocol_dtype.shape ) - and LMDBFormat._shapes_compatible( - spec.shape, protocol_dtype.shape - ) + and LMDBFormat._shapes_compatible(spec.shape, protocol_dtype.shape) ): merged_shape = protocol_dtype.shape - if ( - merged_shape is None - or current.deepmd_name != spec.deepmd_name - ): + if merged_shape is None or current.deepmd_name != spec.deepmd_name: raise LMDBError( f"Data type '{name}' conflicts with the process-global " f"definition: existing shape/name " @@ -2039,9 +1985,7 @@ def _shapes_compatible( if first is None or second is None or len(first) != len(second): return first == second return all( - left == right - or left == -1 - or right == -1 + left == right or left == -1 or right == -1 for left, right in zip(first, second) ) diff --git a/tests/test_lmdb.py b/tests/test_lmdb.py index ff56c5b05..7f2525f92 100644 --- a/tests/test_lmdb.py +++ b/tests/test_lmdb.py @@ -288,9 +288,7 @@ def setUp(self): ntypes = len(self.type_map) counts = np.bincount(fr["atom_types"], minlength=ntypes)[:ntypes] out["atom_numbs"] = [int(c) for c in counts] - txn.put( - format(i, "012d").encode(), _reference_packb(out) - ) + txn.put(format(i, "012d").encode(), _reference_packb(out)) frame_nlocs.append(len(fr["atom_types"])) meta: dict[str, object] = { "nframes": len(self.frames), @@ -337,9 +335,7 @@ def _make_labeled_system( ) -> dpdata.LabeledSystem: atype = np.array(atom_types, dtype=int) natoms = len(atype) - numbs = [ - int(np.count_nonzero(atype == i)) for i in range(len(atom_names)) - ] + numbs = [int(np.count_nonzero(atype == i)) for i in range(len(atom_names))] data = { "atom_numbs": numbs, "atom_names": list(atom_names), @@ -452,9 +448,7 @@ def test_inconsistent_keys_raise(self): }, ] _write_raw_lmdb(self.lmdb_path, frames, ["O", "H"]) - with self.assertRaisesRegex( - LMDBFrameError, "must contain identical fields" - ): + with self.assertRaisesRegex(LMDBFrameError, "must contain identical fields"): dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") def test_single_system_multi_composition_warns(self): @@ -529,9 +523,7 @@ def test_same_composition_mixed_pbc_raises(self): }, ] _write_raw_lmdb(self.lmdb_path, frames, ["O", "H"]) - with self.assertRaisesRegex( - LMDBFrameError, "must contain identical fields" - ): + with self.assertRaisesRegex(LMDBFrameError, "must contain identical fields"): dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") def test_max_frames_guard(self): @@ -547,9 +539,7 @@ def test_max_frames_guard(self): ] _write_raw_lmdb(self.lmdb_path, frames, ["H"]) with self.assertRaisesRegex(LMDBError, "exceeding max_frames"): - dpdata.MultiSystems.from_file( - self.lmdb_path, fmt="lmdb", max_frames=1 - ) + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", max_frames=1) loaded = dpdata.MultiSystems.from_file( self.lmdb_path, fmt="lmdb", max_frames=None, labeled=False ) @@ -590,12 +580,8 @@ def test_reference_fparam_shape_does_not_collide_with_natoms(self): }, ] _write_raw_lmdb(self.lmdb_path, frames, ["H"]) - dpdata.MultiSystems.from_file( - self.lmdb_path, fmt="lmdb", labeled=False - ) - dtype = next( - dt for dt in dpdata.System.DTYPES if dt.name == "fparam" - ) + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) + dtype = next(dt for dt in dpdata.System.DTYPES if dt.name == "fparam") shape = dtype.shape self.assertIsNotNone(shape) assert shape is not None @@ -618,17 +604,13 @@ def test_flattened_reference_aparam_is_normalized_and_reordered(self): ), ( np.array([10.0, 11.0, 20.0, 21.0, 30.0, 31.0]), - np.array( - [[[20.0, 21.0], [10.0, 11.0], [30.0, 31.0]]] - ), + np.array([[[20.0, 21.0], [10.0, 11.0], [30.0, 31.0]]]), ), ): with self.subTest(size=raw.size): frames = [ { - "atom_types": np.array( - [1, 0, 1], dtype=np.int32 - ), + "atom_types": np.array([1, 0, 1], dtype=np.int32), "coords": np.zeros((3, 3)), "aparam": raw, } @@ -680,12 +662,8 @@ def test_known_field_shape_hint_cannot_remove_atom_axis(self): } }, ) - with self.assertRaisesRegex( - LMDBMetadataError, "changes its protocol shape" - ): - dpdata.MultiSystems.from_file( - self.lmdb_path, fmt="lmdb", labeled=False - ) + with self.assertRaisesRegex(LMDBMetadataError, "changes its protocol shape"): + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) def test_data_name_hint_must_match_registered_protocol(self): frames = [ @@ -701,12 +679,8 @@ def test_data_name_hint_must_match_registered_protocol(self): ["H"], metadata_extra={"dp_data_names": {"foo": "spins"}}, ) - with self.assertRaisesRegex( - LMDBMetadataError, "protocol key is 'spin'" - ): - dpdata.MultiSystems.from_file( - self.lmdb_path, fmt="lmdb", labeled=False - ) + with self.assertRaisesRegex(LMDBMetadataError, "protocol key is 'spin'"): + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) def test_custom_deepmd_core_key_rejected_on_read(self): frames = [ @@ -717,12 +691,8 @@ def test_custom_deepmd_core_key_rejected_on_read(self): } ] _write_raw_lmdb(self.lmdb_path, frames, ["H"]) - with self.assertRaisesRegex( - LMDBFrameError, "reserved LMDB key 'coord'" - ): - dpdata.MultiSystems.from_file( - self.lmdb_path, fmt="lmdb", labeled=False - ) + with self.assertRaisesRegex(LMDBFrameError, "reserved LMDB key 'coord'"): + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) def test_unknown_atomic_axis_inferred_across_atom_counts(self): frames = [ @@ -742,9 +712,7 @@ def test_unknown_atomic_axis_inferred_across_atom_counts(self): self.lmdb_path, fmt="lmdb", labeled=False ) self.assertEqual(systems.get_nframes(), 2) - dtype = next( - dt for dt in dpdata.System.DTYPES if dt.name == "mystery" - ) + dtype = next(dt for dt in dpdata.System.DTYPES if dt.name == "mystery") self.assertEqual( dtype.shape, ( @@ -763,12 +731,8 @@ def test_unknown_same_nloc_atom_axis_is_rejected_as_ambiguous(self): } ] _write_raw_lmdb(self.lmdb_path, frames, ["H", "O"]) - with self.assertRaisesRegex( - LMDBFrameError, "ambiguous without dp_data_shapes" - ): - dpdata.MultiSystems.from_file( - self.lmdb_path, fmt="lmdb", labeled=False - ) + with self.assertRaisesRegex(LMDBFrameError, "ambiguous without dp_data_shapes"): + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) def test_process_global_schema_conflict_raises(self): second_path = "tmp_robust_second.lmdb" @@ -789,15 +753,9 @@ def test_process_global_schema_conflict_raises(self): ] _write_raw_lmdb(self.lmdb_path, first_frames, ["H"]) _write_raw_lmdb(second_path, second_frames, ["H"]) - dpdata.MultiSystems.from_file( - self.lmdb_path, fmt="lmdb", labeled=False - ) - with self.assertRaisesRegex( - LMDBError, "process-global definition" - ): - dpdata.MultiSystems.from_file( - second_path, fmt="lmdb", labeled=False - ) + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) + with self.assertRaisesRegex(LMDBError, "process-global definition"): + dpdata.MultiSystems.from_file(second_path, fmt="lmdb", labeled=False) finally: if os.path.exists(second_path): shutil.rmtree(second_path) @@ -820,15 +778,9 @@ def test_failed_read_does_not_partially_register_data_types(self): }, ] _write_raw_lmdb(self.lmdb_path, frames, ["H"]) - with self.assertRaisesRegex( - LMDBFrameError, "must contain identical fields" - ): - dpdata.MultiSystems.from_file( - self.lmdb_path, fmt="lmdb", labeled=False - ) - self.assertNotIn( - "new_scalar", [dt.name for dt in dpdata.System.DTYPES] - ) + with self.assertRaisesRegex(LMDBFrameError, "must contain identical fields"): + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) + self.assertNotIn("new_scalar", [dt.name for dt in dpdata.System.DTYPES]) def test_caller_construction_failure_rolls_back_registration(self): frames = [ @@ -859,9 +811,7 @@ def test_direct_labeled_read_rejects_unlabeled_before_registration(self): } ] _write_raw_lmdb(self.lmdb_path, frames, ["H"]) - with self.assertRaisesRegex( - LMDBFrameError, "required by LabeledSystem" - ): + with self.assertRaisesRegex(LMDBFrameError, "required by LabeledSystem"): dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") self.assertNotIn( "unlabeled_custom", @@ -980,9 +930,7 @@ def test_uint64_overflow_not_treated_as_virtual_atom(self): "atom_names": ["MIXED_TOKEN"], "atom_types": np.array([0]), "real_atom_names": ["H"], - "real_atom_types": np.array( - [[np.iinfo(np.uint64).max]], dtype=np.uint64 - ), + "real_atom_types": np.array([[np.iinfo(np.uint64).max]], dtype=np.uint64), "orig": np.zeros(3), "cells": np.eye(3)[None], "coords": np.zeros((1, 1, 3)), @@ -1007,12 +955,8 @@ def test_nonportable_array_dtypes_rejected(self): ): with self.subTest(dtype=dtype): data = self.A.data.copy() - data["coords"] = np.zeros( - self.A["coords"].shape, dtype=dtype - ) - with self.assertRaisesRegex( - LMDBError, "not a portable raw-byte dtype" - ): + data["coords"] = np.zeros(self.A["coords"].shape, dtype=dtype) + with self.assertRaisesRegex(LMDBError, "not a portable raw-byte dtype"): dump_systems([data], self.lmdb_path) self.assertFalse(os.path.exists(self.lmdb_path)) @@ -1035,9 +979,7 @@ def test_mixed_raw_virtual_atoms_are_removed(self): with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: with env.begin() as txn: meta = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) - frame = msgpack.unpackb( - txn.get(b"000000000000"), raw=False - ) + frame = msgpack.unpackb(txn.get(b"000000000000"), raw=False) self.assertEqual(meta["type_map"], ["H", "O"]) self.assertEqual(meta["frame_nlocs"], [2]) self.assertEqual(frame["atom_types"]["shape"], [2]) @@ -1058,9 +1000,7 @@ def test_ambiguous_mixed_custom_field_rejected(self): "coords": np.zeros((1, 3, 3)), "mystery": np.zeros((1, 3, 2)), } - with self.assertRaisesRegex( - LMDBError, "shape is ambiguous" - ): + with self.assertRaisesRegex(LMDBError, "shape is ambiguous"): dump_systems([data], self.lmdb_path) def test_existing_destination_requires_overwrite(self): @@ -1092,9 +1032,7 @@ def test_failed_overwrite_preserves_old_database(self): loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") np.testing.assert_array_equal(loaded["energies"], [10.0]) temporary = list( - Path(self.lmdb_path).parent.glob( - f".{Path(self.lmdb_path).name}.tmp-*" - ) + Path(self.lmdb_path).parent.glob(f".{Path(self.lmdb_path).name}.tmp-*") ) self.assertEqual(temporary, []) @@ -1158,9 +1096,7 @@ def test_active_dpdata_reader_blocks_overwrite_for_path_alias(self): dump_systems([self.A], self.lmdb_path) resolved, _ = _open_read_env(str(Path(self.lmdb_path).resolve())) try: - with self.assertRaisesRegex( - LMDBError, "reader is active" - ): + with self.assertRaisesRegex(LMDBError, "reader is active"): dump_systems( [self.B], self.lmdb_path, @@ -1328,9 +1264,7 @@ def test_non_mapping_metadata_rejected(self): with env.begin(write=True) as txn: txn.put(b"__metadata__", _reference_packb(["not", "a", "mapping"])) env.close() - with self.assertRaisesRegex( - LMDBMetadataError, "must contain a mapping" - ): + with self.assertRaisesRegex(LMDBMetadataError, "must contain a mapping"): dpdata.MultiSystems.from_file(path, fmt="lmdb") finally: if os.path.exists(path): @@ -1429,9 +1363,7 @@ def test_malformed_array_payload_rejected(self): ), ) env.close() - with self.assertRaisesRegex( - LMDBFrameError, "Cannot decode array" - ): + with self.assertRaisesRegex(LMDBFrameError, "Cannot decode array"): dpdata.MultiSystems.from_file(path, fmt="lmdb") finally: if os.path.exists(path): @@ -1444,18 +1376,14 @@ def test_invalid_frame_system_ids_length(self): system.to("lmdb", path) env = lmdb.open(path, map_size=1 << 30) with env.begin(write=True) as txn: - metadata = msgpack.unpackb( - txn.get(b"__metadata__"), raw=False - ) + metadata = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) metadata["frame_system_ids"] = [] txn.put( b"__metadata__", _reference_packb(metadata), ) env.close() - with self.assertRaisesRegex( - LMDBMetadataError, "frame_system_ids length" - ): + with self.assertRaisesRegex(LMDBMetadataError, "frame_system_ids length"): dpdata.MultiSystems.from_file(path, fmt="lmdb") finally: if os.path.exists(path): diff --git a/tests/test_lmdb_custom_dtype.py b/tests/test_lmdb_custom_dtype.py index f2bf1a170..2eb83b756 100644 --- a/tests/test_lmdb_custom_dtype.py +++ b/tests/test_lmdb_custom_dtype.py @@ -202,12 +202,8 @@ def test_deepmd_name_used_on_disk_and_restored(self): with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: with env.begin() as txn: - metadata = msgpack.unpackb( - txn.get(b"__metadata__"), raw=False - ) - frame = msgpack.unpackb( - txn.get(b"000000000000"), raw=False - ) + metadata = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) + frame = msgpack.unpackb(txn.get(b"000000000000"), raw=False) self.assertIn("spin", frame) self.assertNotIn("spins", frame) self.assertEqual(metadata["dp_data_names"], {"spin": "spins"}) @@ -282,9 +278,7 @@ def test_additional_protocol_alias_rejected(self): 3, ) ) - with self.assertRaisesRegex( - LMDBError, "belongs to 'spins'" - ): + with self.assertRaisesRegex(LMDBError, "belongs to 'spins'"): self.system.to("lmdb", self.lmdb_path) def test_atom_types_key_collision_rejected(self): @@ -303,9 +297,7 @@ def test_atom_types_key_collision_rejected(self): self.system.to("lmdb", self.lmdb_path) def test_static_field_roundtrip(self): - static = DataType( - "static_data", np.ndarray, shape=(2,), required=False - ) + static = DataType("static_data", np.ndarray, shape=(2,), required=False) self._register(static) self.system.data["static_data"] = np.array([1.0, 2.0]) self.system.to("lmdb", self.lmdb_path) @@ -313,9 +305,7 @@ def test_static_field_roundtrip(self): np.testing.assert_array_equal(loaded["static_data"], [1.0, 2.0]) def test_inconsistent_static_field_raises(self): - static = DataType( - "static_data", np.ndarray, shape=(2,), required=False - ) + static = DataType("static_data", np.ndarray, shape=(2,), required=False) self._register(static) self.system.data["static_data"] = np.array([1.0, 2.0]) self.system.to("lmdb", self.lmdb_path) @@ -347,9 +337,9 @@ def test_nonleading_frame_axis_roundtrip(self): required=False, ) self._register(transposed) - values = np.arange( - 2 * self.system.get_nframes(), dtype=float - ).reshape(2, self.system.get_nframes()) + values = np.arange(2 * self.system.get_nframes(), dtype=float).reshape( + 2, self.system.get_nframes() + ) self.system.data["transposed_frames"] = values self.system.to("lmdb", self.lmdb_path) loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") @@ -413,9 +403,7 @@ def test_multiple_atom_axes_remove_virtual_atoms(self): dump_systems([data], self.lmdb_path) with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: with env.begin() as txn: - frame = msgpack.unpackb( - txn.get(b"000000000000"), raw=False - ) + frame = msgpack.unpackb(txn.get(b"000000000000"), raw=False) self.assertEqual(frame["hessian"]["shape"], [2, 3, 2, 3]) From 3590e023b0432bc9b5055f9e37f1dee9155bb5a8 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 12 Jul 2026 19:44:40 +0800 Subject: [PATCH 3/6] fix(lmdb): stabilize CI and overwrite behavior Make protocol-metadata assertions independent of global plugin registration order and reject unsafe Windows overwrite operations explicitly. --- docs/systems/lmdb.md | 2 +- dpdata/formats/lmdb/format.py | 15 ++++++++++++--- tests/test_lmdb.py | 27 ++++++++++++++++++++++++++- 3 files changed, 39 insertions(+), 5 deletions(-) diff --git a/docs/systems/lmdb.md b/docs/systems/lmdb.md index 1259e5528..374e703de 100644 --- a/docs/systems/lmdb.md +++ b/docs/systems/lmdb.md @@ -44,7 +44,7 @@ dpdata.MultiSystems(*systems).to("lmdb", "data.lmdb", type_map=list(ELEMENTS)) Frames are committed in batches of 1,000 by default. The `write_batch_size` argument changes the transaction size. If a transaction exceeds `map_size`, the map is enlarged and the same encoded batch is retried before frame counters are advanced. The destination must not exist unless `overwrite=True` is supplied. -Data are first written to a temporary sibling database. After closing it, dpdata reopens the staged database and validates the metadata, entry count, and continuous frame-key sequence. A new destination is published by one directory rename. For `overwrite=True`, the staged `data.mdb` replaces the existing `data.mdb` in one file rename, while the LMDB runtime lock file remains in place. A failed conversion or validation leaves the previous destination unchanged. All dpdata readers of the destination must be closed before replacement; readers owned by other libraries must likewise be closed because their environments cannot be coordinated through dpdata's process-local cache. +Data are first written to a temporary sibling database. After closing it, dpdata reopens the staged database and validates the metadata, entry count, and continuous frame-key sequence. A new destination is published by one directory rename. On POSIX systems, `overwrite=True` replaces the existing `data.mdb` with the staged `data.mdb` in one file rename, while the LMDB runtime lock file remains in place. Windows does not permit safe replacement of an open memory-mapped LMDB file, so `overwrite=True` is not supported there. A failed conversion or validation leaves the previous destination unchanged. All dpdata readers of the destination must be closed before replacement; readers owned by other libraries must likewise be closed because their environments cannot be coordinated through dpdata's process-local cache. ## Preserving the system partition diff --git a/dpdata/formats/lmdb/format.py b/dpdata/formats/lmdb/format.py index 991fab820..295de7f5a 100644 --- a/dpdata/formats/lmdb/format.py +++ b/dpdata/formats/lmdb/format.py @@ -61,6 +61,7 @@ DEFAULT_MAX_FRAMES = 100_000 DEFAULT_WRITE_BATCH_SIZE = 1_000 METADATA_KEY = b"__metadata__" +_IS_WINDOWS = os.name == "nt" _RESERVED_KEYS = frozenset( { @@ -730,6 +731,11 @@ def __init__( f"LMDB destination '{directory}' already exists; pass " "overwrite=True to replace it." ) + if self.directory.exists() and overwrite and _IS_WINDOWS: + raise NotImplementedError( + "overwrite=True is not supported on Windows because LMDB uses " + "memory-mapped files that cannot be replaced safely while open." + ) self.directory.parent.mkdir(parents=True, exist_ok=True) self.temp_directory = Path( tempfile.mkdtemp( @@ -1149,7 +1155,8 @@ def to_multi_systems( Number of frames committed per LMDB write transaction. overwrite : bool, optional Whether to replace an existing destination after the new database - has been written and validated. The default is ``False``. + has been written and validated. This option is supported on POSIX + systems only. The default is ``False``. **kwargs : dict other parameters @@ -1219,7 +1226,8 @@ def dump_systems( Number of frames committed per LMDB write transaction. overwrite : bool, optional Whether to replace an existing destination after the new database - has been written and validated. The default is ``False``. + has been written and validated. This option is supported on POSIX + systems only. The default is ``False``. """ if hasattr(systems, "systems"): raise TypeError( @@ -2043,7 +2051,8 @@ def dump_systems( Number of frames committed per LMDB write transaction. overwrite : bool, optional Whether to replace an existing destination after the new database has - been written and validated. The default is ``False``. + been written and validated. This option is supported on POSIX systems + only. The default is ``False``. Examples -------- diff --git a/tests/test_lmdb.py b/tests/test_lmdb.py index 7f2525f92..5afdcd090 100644 --- a/tests/test_lmdb.py +++ b/tests/test_lmdb.py @@ -679,8 +679,10 @@ def test_data_name_hint_must_match_registered_protocol(self): ["H"], metadata_extra={"dp_data_names": {"foo": "spins"}}, ) - with self.assertRaisesRegex(LMDBMetadataError, "protocol key is 'spin'"): + with self.assertRaises(LMDBMetadataError) as context: dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) + self.assertIn("foo", str(context.exception)) + self.assertIn("spins", str(context.exception)) def test_custom_deepmd_core_key_rejected_on_read(self): frames = [ @@ -1010,6 +1012,24 @@ def test_existing_destination_requires_overwrite(self): with self.assertRaises(FileExistsError): dump_systems([self.B], self.lmdb_path) + def test_windows_overwrite_rejected_explicitly(self): + from dpdata.formats.lmdb import dump_systems + + dump_systems([self.A], self.lmdb_path) + with ( + mock.patch( + "dpdata.formats.lmdb.format._IS_WINDOWS", + True, + ), + self.assertRaisesRegex(NotImplementedError, "not supported on Windows"), + ): + dump_systems( + [self.B], + self.lmdb_path, + overwrite=True, + ) + + @unittest.skipIf(os.name == "nt", "LMDB overwrite is POSIX-only") def test_failed_overwrite_preserves_old_database(self): from dpdata.formats.lmdb import dump_systems @@ -1036,6 +1056,7 @@ def test_failed_overwrite_preserves_old_database(self): ) self.assertEqual(temporary, []) + @unittest.skipIf(os.name == "nt", "LMDB overwrite is POSIX-only") def test_successful_overwrite_replaces_database(self): from dpdata.formats.lmdb import dump_systems @@ -1090,6 +1111,7 @@ def test_map_full_grows_and_retries_batch(self): with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: self.assertGreater(env.info()["map_size"], 4096) + @unittest.skipIf(os.name == "nt", "LMDB overwrite is POSIX-only") def test_active_dpdata_reader_blocks_overwrite_for_path_alias(self): from dpdata.formats.lmdb import dump_systems @@ -1107,6 +1129,7 @@ def test_active_dpdata_reader_blocks_overwrite_for_path_alias(self): loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") self.assertEqual(loaded.get_nframes(), self.A.get_nframes()) + @unittest.skipIf(os.name == "nt", "LMDB overwrite is POSIX-only") def test_external_lmdb_reader_blocks_overwrite(self): from dpdata.formats.lmdb import dump_systems @@ -1132,6 +1155,7 @@ def test_external_lmdb_reader_blocks_overwrite(self): loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") self.assertEqual(loaded.get_nframes(), self.A.get_nframes()) + @unittest.skipIf(os.name == "nt", "LMDB overwrite is POSIX-only") def test_publish_guard_closes_external_reader_race(self): from dpdata.formats.lmdb import dump_systems @@ -1183,6 +1207,7 @@ def synchronized_replace(source, destination): loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") self.assertEqual(loaded[0].get_natoms(), self.B.get_natoms()) + @unittest.skipIf(os.name == "nt", "LMDB overwrite is POSIX-only") def test_staged_validation_failure_preserves_old_database(self): from dpdata.formats.lmdb import dump_systems from dpdata.formats.lmdb.format import _LMDBWriter From 096d5d5dd5efc7c129614c4a81f75b7bc42fdacd Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 14 Jul 2026 14:55:05 +0800 Subject: [PATCH 4/6] fix(lmdb): clarify protocol and runtime requirement Document the core-field wire names and require python-lmdb 2.0 so duplicate environment detection reliably protects active readers. --- docs/systems/lmdb.md | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/systems/lmdb.md b/docs/systems/lmdb.md index 374e703de..67dccc85e 100644 --- a/docs/systems/lmdb.md +++ b/docs/systems/lmdb.md @@ -1,6 +1,6 @@ # LMDB Format -The format `lmdb` stores the frames of one or more systems in a single [LMDB](http://www.lmdb.tech/doc/) database, and can be loaded or dumped through {class}`dpdata.System`, {class}`dpdata.LabeledSystem`, and {class}`dpdata.MultiSystems`. The on-disk layout coincides with the LMDB datasets read by the DeePMD-kit data loader. Core fields and registered additional fields are therefore available to DeePMD-kit under their `deepmd_name`. +The format `lmdb` stores the frames of one or more systems in a single [LMDB](http://www.lmdb.tech/doc/) database, and can be loaded or dumped through {class}`dpdata.System`, {class}`dpdata.LabeledSystem`, and {class}`dpdata.MultiSystems`. The on-disk layout coincides with the LMDB datasets read by the DeePMD-kit data loader. Core fields use the plural on-disk names `coords`, `cells`, `energies`, `forces`, and `virials`, which the DeePMD-kit reader maps to its internal names. Registered additional fields use their `deepmd_name`. In contrast to the directory-based `deepmd/npy` format, every frame is stored as an independent record indexed by a global frame number. Frames within one database may therefore differ in the number of atoms and in chemical composition, which is suited to data sets in which the number of frames per system is small. diff --git a/pyproject.toml b/pyproject.toml index 754dd9cdd..0a846ca1d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,7 @@ dependencies = [ 'scipy', 'h5py', 'wcmatch', - 'lmdb', + 'lmdb>=2.0.0', 'msgpack', ] requires-python = ">=3.10" From 67f80072de074394230af6669c1f2557e764ad5a Mon Sep 17 00:00:00 2001 From: OutisLi Date: Wed, 15 Jul 2026 12:27:07 +0800 Subject: [PATCH 5/6] fix(lmdb): preserve legacy and type-axis data Read and migrate the released schema, remap every type axis consistently, and reject cross-system field ownership or unsupported scalar data. --- docs/systems/lmdb.md | 4 + dpdata/formats/lmdb/format.py | 382 ++++++++++++++++++++++++++++++-- dpdata/system.py | 37 ++++ tests/test_lmdb.py | 119 ++++++++++ tests/test_lmdb_custom_dtype.py | 120 +++++++++- 5 files changed, 644 insertions(+), 18 deletions(-) diff --git a/docs/systems/lmdb.md b/docs/systems/lmdb.md index 67dccc85e..96c60facb 100644 --- a/docs/systems/lmdb.md +++ b/docs/systems/lmdb.md @@ -91,6 +91,10 @@ The general {class}`dpdata.System` constructor applies its `type_map` argument a ms = dpdata.MultiSystems.from_file("data.lmdb", fmt="lmdb", mixed_type=True) ``` +## Legacy databases + +dpdata detects the `system_info` metadata used by the LMDB schema released in dpdata v1.0.2 and decodes its msgpack-numpy records without requiring `msgpack-numpy` at runtime. Legacy databases are supported for reading and migration only; all writes use the DeePMD-kit-compatible flat schema described above. Loading a legacy database and writing the resulting systems to a new LMDB path performs the migration. Legacy object arrays are rejected because their msgpack-numpy representation requires unsafe pickle deserialization. + A database that holds a single composition may also be read into a {class}`dpdata.LabeledSystem`. If the database contains several compositions, only the first can be represented by a single system and a warning is issued. ```python diff --git a/dpdata/formats/lmdb/format.py b/dpdata/formats/lmdb/format.py index 295de7f5a..709ff32fb 100644 --- a/dpdata/formats/lmdb/format.py +++ b/dpdata/formats/lmdb/format.py @@ -134,6 +134,11 @@ def atom_axes(self) -> tuple[int, ...]: """Return atom axes in the in-memory array.""" return tuple(i for i, dim in enumerate(self.shape) if dim is Axis.NATOMS) + @property + def type_axes(self) -> tuple[int, ...]: + """Return type axes in the in-memory array.""" + return tuple(i for i, dim in enumerate(self.shape) if dim is Axis.NTYPES) + @property def payload_atom_axes(self) -> tuple[int, ...]: """Return atom axes after the frame axis has been removed.""" @@ -142,6 +147,14 @@ def payload_atom_axes(self) -> tuple[int, ...]: return self.atom_axes return tuple(axis - (axis > frame_axis) for axis in self.atom_axes) + @property + def payload_type_axes(self) -> tuple[int, ...]: + """Return type axes after the frame axis has been removed.""" + frame_axis = self.frame_axis + if frame_axis is None: + return self.type_axes + return tuple(axis - (axis > frame_axis) for axis in self.type_axes) + @dataclass(frozen=True) class _AtomSource: @@ -352,6 +365,75 @@ def _decode_frame(raw: bytes) -> dict[str, Any]: return out +def _unpack_legacy_dtype(value: Any) -> np.dtype: + """Decode a msgpack-numpy dtype descriptor without that dependency.""" + if isinstance(value, list): + value = [ + ( + item[0].decode() if isinstance(item[0], bytes) else item[0], + _unpack_legacy_dtype(item[1]), + *item[2:], + ) + for item in value + ] + return np.dtype(value) + + +def _decode_legacy_object(value: dict[Any, Any]) -> Any: + """Decode one object produced by the released msgpack-numpy writer.""" + nd_key = b"nd" if b"nd" in value else "nd" if "nd" in value else None + if nd_key is not None: + type_key = b"type" if b"type" in value else "type" + data_key = b"data" if b"data" in value else "data" + dtype = _unpack_legacy_dtype(value[type_key]) + if dtype.hasobject: + raise LMDBFrameError( + "Legacy object arrays are not loaded because they require " + "unsafe pickle deserialization." + ) + if value[nd_key]: + shape_key = b"shape" if b"shape" in value else "shape" + return ( + np.frombuffer(value[data_key], dtype=dtype) + .reshape(tuple(value[shape_key])) + .copy() + ) + return np.frombuffer(value[data_key], dtype=dtype, count=1)[0] + complex_key = ( + b"complex" if b"complex" in value else "complex" if "complex" in value else None + ) + if complex_key is not None: + data_key = b"data" if b"data" in value else "data" + raw = value[data_key] + return complex(raw.decode() if isinstance(raw, bytes) else raw) + return value + + +def _decode_legacy_frame(raw: bytes) -> dict[str, Any]: + """Decode a frame from the LMDB schema released in dpdata v1.0.2.""" + try: + frame = msgpack.unpackb( + raw, + raw=False, + strict_map_key=False, + object_hook=_decode_legacy_object, + ) + except ( + msgpack.exceptions.ExtraData, + msgpack.exceptions.FormatError, + msgpack.exceptions.StackError, + TypeError, + ValueError, + ) as exc: + raise LMDBFrameError(f"Cannot decode legacy LMDB frame: {exc}") from exc + if not isinstance(frame, dict): + raise LMDBFrameError("Each legacy LMDB frame must contain a mapping.") + return { + key.decode() if isinstance(key, bytes) else str(key): value + for key, value in frame.items() + } + + def _meta_get(meta: dict[str | bytes, Any], key: str) -> Any: """Read a metadata value, tolerating string or byte keys.""" if key in meta: @@ -576,18 +658,14 @@ def _source_data(system: Any) -> dict[str, Any]: def _source_active_names(data: dict[str, Any]) -> tuple[str, ...]: - """Return active elements in first-appearance order for type-map discovery.""" + """Return source element names in first-appearance order.""" coords = np.asarray(data["coords"]) if coords.ndim != 3 or coords.shape[2] != 3: raise LMDBError( f"coords must have shape (nframes, natoms, 3), got {coords.shape}." ) source = _prepare_atom_source(data, coords.shape[0], coords.shape[1]) - types = source.types[source.types >= 0] - active_indices = set(int(value) for value in np.unique(types)) - return tuple( - name for index, name in enumerate(source.names) if index in active_indices - ) + return source.names def _remove_path(path: Path) -> None: @@ -753,6 +831,7 @@ def __init__( self.frame_system_ids: list[int] = [] self.data_shapes: dict[str, list[int | str]] = {} self.data_names: dict[str, str] = {} + self.disk_owners: dict[str, str] = {} self._closed = False self._published = False try: @@ -769,6 +848,7 @@ def _build_specs( *, nframes: int, natoms: int, + ntypes: int, has_virtual_atoms: bool, ) -> list[_FieldSpec]: """Build and validate field specifications for one system.""" @@ -777,8 +857,13 @@ def _build_specs( specs: list[_FieldSpec] = [] used_disk_names: dict[str, str] = {} for name, value in data.items(): - if name in _RESERVED_KEYS or not isinstance(value, np.ndarray): + if name in _RESERVED_KEYS: continue + if not isinstance(value, np.ndarray): + raise LMDBError( + f"Field '{name}' has unsupported non-array type " + f"{type(value).__name__}; LMDB fields must be numpy arrays." + ) array = np.asarray(value) registered_dtype = all_dtypes.get(name) if registered_dtype is not None and registered_dtype.shape is None: @@ -815,7 +900,15 @@ def _build_specs( f"key '{spec.disk_name}'." ) used_disk_names[spec.disk_name] = spec.dp_name - self._validate_source_shape(spec, array, nframes, natoms) + existing_owner = self.disk_owners.get(spec.disk_name) + if existing_owner is not None and existing_owner != spec.dp_name: + raise LMDBError( + f"LMDB key '{spec.disk_name}' is owned by " + f"'{existing_owner}' in one system and '{spec.dp_name}' " + "in another." + ) + self.disk_owners[spec.disk_name] = spec.dp_name + self._validate_source_shape(spec, array, nframes, natoms, ntypes) specs.append(spec) if spec.disk_name not in _CORE_FRAME_KEYS: existing_shape = self.data_shapes.get(spec.disk_name) @@ -840,7 +933,11 @@ def _build_specs( @staticmethod def _validate_source_shape( - spec: _FieldSpec, array: np.ndarray, nframes: int, natoms: int + spec: _FieldSpec, + array: np.ndarray, + nframes: int, + natoms: int, + ntypes: int, ) -> None: """Validate an in-memory field against its symbolic shape.""" if array.ndim != len(spec.shape): @@ -859,6 +956,11 @@ def _validate_source_shape( f"Field '{spec.dp_name}' axis {axis} must have {natoms} " f"atoms, got {actual}." ) + if expected is Axis.NTYPES and actual != ntypes: + raise LMDBError( + f"Field '{spec.dp_name}' axis {axis} must have {ntypes} " + f"types, got {actual}." + ) if isinstance(expected, int) and expected != -1 and actual != expected: raise LMDBError( f"Field '{spec.dp_name}' axis {axis} must have length " @@ -871,8 +973,10 @@ def _frame_payload( spec: _FieldSpec, frame_index: int, atom_mask: np.ndarray | None, + type_indices: np.ndarray, + target_ntypes: int, ) -> np.ndarray: - """Extract one frame and remove virtual atoms on every atom axis.""" + """Extract one frame and normalize atom and type axes.""" frame_axis = spec.frame_axis payload = ( np.take(array, frame_index, axis=frame_axis) @@ -882,6 +986,14 @@ def _frame_payload( if atom_mask is not None: for axis in spec.payload_atom_axes: payload = np.compress(atom_mask, payload, axis=axis) + for axis in spec.payload_type_axes: + output_shape = list(payload.shape) + output_shape[axis] = target_ntypes + expanded = np.zeros(output_shape, dtype=payload.dtype) + selection: list[slice | np.ndarray] = [slice(None)] * payload.ndim + selection[axis] = type_indices + expanded[tuple(selection)] = payload + payload = expanded return payload def _put_records(self, records: list[tuple[bytes, bytes]]) -> None: @@ -930,8 +1042,18 @@ def write_system(self, data: dict[str, Any]) -> None: data, nframes=nframes, natoms=natoms, + ntypes=len(atom_source.names), has_virtual_atoms=has_virtual_atoms, ) + if any(spec.type_axes for spec in specs) and np.any(local_to_global < 0): + missing_type_names = [ + atom_source.names[index] + for index in np.flatnonzero(local_to_global < 0) + ] + raise LMDBError( + "Type-dependent fields require every source type in type_map; " + f"missing {missing_type_names}." + ) nopbc = bool(data.get("nopbc", False)) fields = [(spec, np.asarray(data[spec.dp_name])) for spec in specs] @@ -961,7 +1083,14 @@ def write_system(self, data: dict[str, Any]) -> None: for spec, array in fields: if spec.dp_name == "cells" and nopbc: continue - payload = self._frame_payload(array, spec, frame_i, atom_mask) + payload = self._frame_payload( + array, + spec, + frame_i, + atom_mask, + local_to_global, + len(type_map), + ) frame[spec.disk_name] = _encode_array(payload) counts = np.bincount(global_types, minlength=len(type_map))[ : len(type_map) @@ -1345,7 +1474,12 @@ def from_multi_systems( system data dictionary for each composition group. """ frames, meta = self._read_all_frames(directory, max_frames=max_frames) + if _meta_get(meta, "system_info") is not None: + datasets, specs = self._legacy_datasets(frames, meta, type_map) + yield from self._yield_registered_datasets(datasets, specs) + return file_type_map = _meta_get(meta, "type_map") + type_axis_remap: np.ndarray | None = None if file_type_map: names = list( _validate_names( @@ -1353,6 +1487,7 @@ def from_multi_systems( "file type_map", ) ) + stored_ntypes = len(names) self._validate_frame_types(frames, names, meta) if type_map is not None: requested = list(_validate_names(type_map, "type_map")) @@ -1363,6 +1498,7 @@ def from_multi_systems( f"from the requested type_map {requested}." ) remap = np.array([requested.index(n) for n in names], dtype=int) + type_axis_remap = remap for fr in frames: fr["atom_types"] = remap[np.asarray(fr["atom_types"]).astype(int)] fr["atom_numbs"] = np.bincount( @@ -1371,20 +1507,37 @@ def from_multi_systems( names = requested elif type_map is not None: names = list(_validate_names(type_map, "type_map")) + stored_ntypes = len(names) self._validate_frame_types(frames, names, meta) else: names = self._infer_type_names(frames) + stored_ntypes = len(names) self._validate_frame_types(frames, names, meta) self._normalize_reference_fields(frames) specs = self._resolve_read_specs(frames, meta) - frames = self._rename_frame_fields(frames, specs) + frames = self._rename_frame_fields(frames, specs, stored_ntypes) + if type_axis_remap is not None: + self._remap_frame_type_axes( + frames, + specs, + type_axis_remap, + len(names), + ) datasets = self._group_frames(frames, names, mixed_type, specs) dp_specs = {spec.dp_name: spec for spec in specs.values()} + yield from self._yield_registered_datasets(datasets, dp_specs) + + def _yield_registered_datasets( + self, + datasets: list[dict[str, Any]], + specs: dict[str, _FieldSpec], + ): + """Register schemas transactionally while yielding parsed systems.""" from dpdata.system import LabeledSystem, System original_system_dtypes = System.DTYPES original_labeled_dtypes = LabeledSystem.DTYPES - self._register_specs(datasets, dp_specs) + self._register_specs(datasets, specs) completed = False try: yield from datasets @@ -1394,6 +1547,148 @@ def from_multi_systems( System.DTYPES = original_system_dtypes LabeledSystem.DTYPES = original_labeled_dtypes + def _legacy_datasets( + self, + frames: list[dict[str, Any]], + metadata: dict[str | bytes, Any], + type_map: list[str] | None, + ) -> tuple[list[dict[str, Any]], dict[str, _FieldSpec]]: + """Reconstruct systems written by the schema released in dpdata v1.0.2.""" + system_info = _meta_get(metadata, "system_info") + if not isinstance(system_info, list) or not system_info: + raise LMDBMetadataError( + "Legacy LMDB metadata must contain a non-empty system_info list." + ) + known_by_dp, _ = _field_registries() + datasets: list[dict[str, Any]] = [] + specs: dict[str, _FieldSpec] = {} + covered: set[int] = set() + for system_index, raw_info in enumerate(system_info): + if not isinstance(raw_info, dict): + raise LMDBMetadataError( + f"system_info[{system_index}] must be a mapping." + ) + start = _require_integer( + raw_info.get("start_idx"), + f"system_info[{system_index}].start_idx", + minimum=0, + ) + count = _require_integer( + raw_info.get("nframes"), + f"system_info[{system_index}].nframes", + minimum=1, + ) + indices = range(start, start + count) + if start + count > len(frames) or covered.intersection(indices): + raise LMDBMetadataError( + f"system_info[{system_index}] references invalid or " + "overlapping frame indices." + ) + covered.update(indices) + system_frames = frames[start : start + count] + frame_keys = set(system_frames[0]) + if any(set(frame) != frame_keys for frame in system_frames[1:]): + raise LMDBFrameError( + f"Legacy system {system_index} has inconsistent frame fields." + ) + frame_dependent = { + str(name) for name in raw_info.get("frame_dependent_keys", []) + } + data: dict[str, Any] = {} + for name in frame_keys: + values = [frame[name] for frame in system_frames] + if name in frame_dependent: + if all(isinstance(value, np.ndarray) for value in values): + data[name] = np.stack(values, axis=0) + else: + data[name] = np.asarray(values) + else: + data[name] = values[0] + + raw_shapes = raw_info.get("data_shapes", {}) + if not isinstance(raw_shapes, dict): + raise LMDBMetadataError( + f"system_info[{system_index}].data_shapes must be a mapping." + ) + for name, value in data.items(): + if name in _RESERVED_KEYS or not isinstance(value, np.ndarray): + continue + raw_shape = raw_shapes.get(name) + if raw_shape is not None: + shape = _tokens_to_shape(raw_shape) + elif name in known_by_dp: + shape = known_by_dp[name].shape + else: + shape = tuple(value.shape) + if shape is None: + continue + known = known_by_dp.get(name) + spec = _FieldSpec( + dp_name=name, + disk_name=known.disk_name if known is not None else name, + shape=tuple(shape), + deepmd_name=known.deepmd_name if known is not None else name, + ) + existing = specs.get(name) + if existing is not None and existing != spec: + raise LMDBMetadataError( + f"Legacy field '{name}' has incompatible schemas." + ) + specs[name] = spec + if type_map is not None: + self._remap_legacy_type_table(data, specs, type_map) + datasets.append(data) + if covered != set(range(len(frames))): + raise LMDBMetadataError( + "Legacy system_info does not cover every frame exactly once." + ) + return datasets, specs + + @staticmethod + def _remap_legacy_type_table( + data: dict[str, Any], + specs: dict[str, _FieldSpec], + type_map: list[str], + ) -> None: + """Apply a requested type table to one reconstructed legacy system.""" + old_names = _validate_names(data["atom_names"], "legacy atom_names") + requested = _validate_names(type_map, "type_map") + missing = [name for name in old_names if name not in requested] + active = { + name for name, count in zip(old_names, data["atom_numbs"]) if count > 0 + } + missing_active = sorted(active.intersection(missing)) + if missing_active: + raise LMDBError( + f"Active legacy elements {missing_active} are missing from " + f"type_map {list(requested)}." + ) + source_to_target = np.array( + [requested.index(name) if name in requested else -1 for name in old_names], + dtype=np.int64, + ) + atom_types = np.asarray(data["atom_types"], dtype=np.int64) + data["atom_types"] = source_to_target[atom_types] + data["atom_names"] = list(requested) + data["atom_numbs"] = np.bincount( + data["atom_types"], minlength=len(requested) + ).tolist() + for name, spec in specs.items(): + if name not in data: + continue + value = np.asarray(data[name]) + for axis in spec.type_axes: + output_shape = list(value.shape) + output_shape[axis] = len(requested) + remapped = np.zeros(output_shape, dtype=value.dtype) + retained = np.flatnonzero(source_to_target >= 0) + selected = np.take(value, retained, axis=axis) + index: list[slice | np.ndarray] = [slice(None)] * value.ndim + index[axis] = source_to_target[retained] + remapped[tuple(index)] = selected + value = remapped + data[name] = value + @staticmethod def _normalize_reference_fields(frames: list[dict[str, Any]]) -> None: """Normalize flattened atomic parameters emitted by reference converters.""" @@ -1464,6 +1759,7 @@ def _read_all_frames( try: with env.begin() as txn: meta = _read_metadata(txn) + legacy_schema = _meta_get(meta, "system_info") is not None frame_fmt = _meta_get(meta, "frame_idx_fmt") or DEFAULT_FRAME_FMT if not isinstance(frame_fmt, str): raise LMDBMetadataError("frame_idx_fmt must be a string.") @@ -1531,7 +1827,11 @@ def _read_all_frames( raw = txn.get(key) if raw is None: raise LMDBFrameError(f"Frame data not found for key: {key!r}") - frames.append(_decode_frame(bytes(raw))) + frames.append( + _decode_legacy_frame(bytes(raw)) + if legacy_schema + else _decode_frame(bytes(raw)) + ) finally: _close_read_env(resolved) return frames, meta @@ -1746,6 +2046,7 @@ def _infer_field_shape( def _rename_frame_fields( frames: list[dict[str, Any]], specs: dict[str, _FieldSpec], + ntypes: int, ) -> list[dict[str, Any]]: """Map disk field names to dpdata names and validate payload shapes.""" renamed_frames: list[dict[str, Any]] = [] @@ -1764,16 +2065,43 @@ def _rename_frame_fields( f"'{spec.dp_name}'." ) value = np.asarray(frame[disk_name]) - LMDBFormat._validate_payload_shape(spec, value, natoms, frame_index) + LMDBFormat._validate_payload_shape( + spec, value, natoms, ntypes, frame_index + ) renamed[spec.dp_name] = value renamed_frames.append(renamed) return renamed_frames + @staticmethod + def _remap_frame_type_axes( + frames: list[dict[str, Any]], + specs_by_disk: dict[str, _FieldSpec], + source_to_target: np.ndarray, + target_ntypes: int, + ) -> None: + """Permute and extend every per-frame type axis.""" + specs = {spec.dp_name: spec for spec in specs_by_disk.values()} + for frame in frames: + for name, spec in specs.items(): + if name not in frame: + continue + value = np.asarray(frame[name]) + for axis in spec.payload_type_axes: + output_shape = list(value.shape) + output_shape[axis] = target_ntypes + remapped = np.zeros(output_shape, dtype=value.dtype) + selection: list[slice | np.ndarray] = [slice(None)] * value.ndim + selection[axis] = source_to_target + remapped[tuple(selection)] = value + value = remapped + frame[name] = value + @staticmethod def _validate_payload_shape( spec: _FieldSpec, value: np.ndarray, natoms: int, + ntypes: int, frame_index: int, ) -> None: """Validate a decoded per-frame payload against its field specification.""" @@ -1791,6 +2119,11 @@ def _validate_payload_shape( f"Frame {frame_index} field '{spec.disk_name}' axis {axis} " f"must have {natoms} atoms, got {actual}." ) + if dimension is Axis.NTYPES and actual != ntypes: + raise LMDBFrameError( + f"Frame {frame_index} field '{spec.disk_name}' axis {axis} " + f"must have {ntypes} types, got {actual}." + ) if isinstance(dimension, int) and dimension != -1 and actual != dimension: raise LMDBFrameError( f"Frame {frame_index} field '{spec.disk_name}' axis {axis} " @@ -1820,6 +2153,7 @@ def _group_frames( for composition, indices in groups.items(): self._validate_group_schema(canonical_frames, indices) global_atom_types = canonical_frames[indices[0]]["atom_types"] + present: list[int] | None = None if mixed_type: atom_names = list(names) atom_types = global_atom_types.copy() @@ -1838,9 +2172,27 @@ def _group_frames( "orig": np.zeros(3), } self._aggregate_group(data, canonical_frames, indices, specs) + if not mixed_type: + assert present is not None + self._compact_type_axes(data, specs, present) results.append(data) return results + @staticmethod + def _compact_type_axes( + data: dict[str, Any], + specs: dict[str, _FieldSpec], + present: list[int], + ) -> None: + """Restrict every in-memory type axis to the system's active types.""" + for name, spec in specs.items(): + if name not in data: + continue + value = np.asarray(data[name]) + for axis in spec.type_axes: + value = np.take(value, present, axis=axis) + data[name] = value + @staticmethod def _canonicalize_frame( frame: dict[str, Any], specs: dict[str, _FieldSpec] diff --git a/dpdata/system.py b/dpdata/system.py index 18b16aaf8..0134fcf16 100644 --- a/dpdata/system.py +++ b/dpdata/system.py @@ -559,6 +559,43 @@ def sort_atom_names(self, type_map: list[str] | None = None): type_map : list type_map """ + old_names = list(self.data["atom_names"]) + new_names = sorted(old_names) if type_map is None else list(type_map) + active_names = { + name for name, count in zip(old_names, self.data["atom_numbs"]) if count > 0 + } + missing_names = active_names - set(new_names) + if missing_names: + raise ValueError( + f"Active atom types {missing_names} not in provided type_map." + ) + old_to_new = np.array( + [new_names.index(name) if name in new_names else -1 for name in old_names], + dtype=np.int64, + ) + retained_old = np.flatnonzero(old_to_new >= 0) + retained_new = old_to_new[retained_old] + for dtype in self.DTYPES: + if ( + dtype.name in {"atom_names", "atom_numbs", "real_atom_names"} + or dtype.name not in self.data + or dtype.shape is None + or Axis.NTYPES not in dtype.shape + ): + continue + value = np.asarray(self.data[dtype.name]) + for axis, dimension in enumerate(dtype.shape): + if dimension is not Axis.NTYPES: + continue + selected = np.take(value, retained_old, axis=axis) + output_shape = list(selected.shape) + output_shape[axis] = len(new_names) + remapped = np.zeros(output_shape, dtype=value.dtype) + index: list[slice | np.ndarray] = [slice(None)] * value.ndim + index[axis] = retained_new + remapped[tuple(index)] = selected + value = remapped + self.data[dtype.name] = value self.data = sort_atom_names(self.data, type_map=type_map) def check_type_map(self, type_map: list[str] | None): diff --git a/tests/test_lmdb.py b/tests/test_lmdb.py index 5afdcd090..afb3806b2 100644 --- a/tests/test_lmdb.py +++ b/tests/test_lmdb.py @@ -48,6 +48,90 @@ def _reference_packb(value): return packed +def _legacy_encode_value(value): + """Encode numeric values using the released msgpack-numpy wire format.""" + if isinstance(value, np.ndarray): + return { + b"nd": True, + b"type": value.dtype.str, + b"kind": b"", + b"shape": list(value.shape), + b"data": value.tobytes(order="C"), + } + if isinstance(value, (np.bool_, np.number)): + return { + b"nd": False, + b"type": value.dtype.str, + b"data": value.tobytes(), + } + if isinstance(value, dict): + return {key: _legacy_encode_value(item) for key, item in value.items()} + if isinstance(value, list): + return [_legacy_encode_value(item) for item in value] + return value + + +def _write_legacy_lmdb(path, systems): + """Write the LMDB schema released in dpdata v1.0.2.""" + if os.path.exists(path): + shutil.rmtree(path) + env = lmdb.open(path, map_size=1 << 30) + frame_index = 0 + system_info = [] + with env.begin(write=True) as txn: + for system in systems: + data = system.data + frame_fields = { + dtype.name + for dtype in system.DTYPES + if dtype.name in data + and dtype.shape is not None + and Axis.NFRAMES in dtype.shape + } + shapes = { + dtype.name: [ + item.value if isinstance(item, Axis) else item + for item in dtype.shape + ] + if dtype.shape is not None + else None + for dtype in system.DTYPES + if dtype.name in data + } + start = frame_index + for local_index in range(system.get_nframes()): + frame = { + name: value[local_index] if name in frame_fields else value + for name, value in data.items() + } + txn.put( + f"{frame_index:012d}".encode(), + _reference_packb(_legacy_encode_value(frame)), + ) + frame_index += 1 + system_info.append( + { + "formula": system.formula, + "natoms": data["atom_numbs"], + "nframes": system.get_nframes(), + "start_idx": start, + "data_shapes": shapes, + "frame_dependent_keys": sorted(frame_fields), + } + ) + txn.put( + b"__metadata__", + _reference_packb( + { + "nframes": frame_index, + "frame_idx_fmt": "012d", + "system_info": system_info, + } + ), + ) + env.close() + + class TestLMDBLabeledSystem(unittest.TestCase, CompLabeledSys, IsPBC): def setUp(self): self.system_1 = dpdata.LabeledSystem("poscars/OUTCAR.h2o.md", fmt="vasp/outcar") @@ -349,6 +433,41 @@ def _make_labeled_system( return dpdata.LabeledSystem(data=data) +class TestLMDBLegacyCompatibility(unittest.TestCase): + """Read and migrate the LMDB schema released in dpdata v1.0.2.""" + + def setUp(self): + self.legacy_path = "tmp_legacy.lmdb" + self.migrated_path = "tmp_legacy_migrated.lmdb" + self.water = _make_labeled_system([0, 1, 1], ["O", "H"], 2) + self.methane = _make_labeled_system([0, 1, 1, 1, 1], ["C", "H"], 1) + _write_legacy_lmdb(self.legacy_path, [self.water, self.methane]) + + def tearDown(self): + for path in (self.legacy_path, self.migrated_path): + if os.path.exists(path): + shutil.rmtree(path) + + def test_read_released_schema(self): + systems = dpdata.MultiSystems.from_file(self.legacy_path, fmt="lmdb") + self.assertEqual(len(systems), 2) + self.assertEqual(systems.get_nframes(), 3) + water = next(system for system in systems if system.get_natoms() == 3) + np.testing.assert_array_equal(water["energies"], self.water["energies"]) + np.testing.assert_array_equal(water["coords"], self.water["coords"]) + + def test_legacy_database_can_be_migrated_by_roundtrip(self): + systems = dpdata.MultiSystems.from_file(self.legacy_path, fmt="lmdb") + systems.to("lmdb", self.migrated_path) + with lmdb.open(self.migrated_path, readonly=True, lock=False) as env: + with env.begin() as txn: + metadata = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) + self.assertNotIn("system_info", metadata) + self.assertIn("type_map", metadata) + loaded = dpdata.MultiSystems.from_file(self.migrated_path, fmt="lmdb") + self.assertEqual(loaded.get_nframes(), 3) + + def _write_raw_lmdb(path, frames, type_map, metadata_extra=None): """Write frames using the reference encoding (one record per frame).""" if os.path.exists(path): diff --git a/tests/test_lmdb_custom_dtype.py b/tests/test_lmdb_custom_dtype.py index 2eb83b756..480ed6957 100644 --- a/tests/test_lmdb_custom_dtype.py +++ b/tests/test_lmdb_custom_dtype.py @@ -9,7 +9,7 @@ import numpy as np import dpdata -from dpdata.data_type import Axis, DataType +from dpdata.data_type import Axis, DataError, DataType from dpdata.formats.lmdb.format import LMDBError, LMDBFrameError @@ -156,7 +156,7 @@ def test_symbolic_axis_natoms_preservation(self): } try: dpdata.LabeledSystem(data=data_diff) - except dpdata.data_type.DataError as e: + except DataError as e: self.fail(f"DataError raised despite symbolic NATOMS: {e}") @@ -320,7 +320,9 @@ def test_inconsistent_static_field_raises(self): "shape": list(replacement.shape), "data": replacement.tobytes(), } - txn.put(key, msgpack.packb(frame, use_bin_type=True)) + packed = msgpack.packb(frame, use_bin_type=True) + assert isinstance(packed, bytes) + txn.put(key, packed) env.close() with self.assertRaisesRegex( @@ -407,5 +409,117 @@ def test_multiple_atom_axes_remove_virtual_atoms(self): self.assertEqual(frame["hessian"]["shape"], [2, 3, 2, 3]) +class TestLMDBTypeAxis(unittest.TestCase): + def setUp(self): + self.original_system_dtypes = dpdata.System.DTYPES + self.original_labeled_system_dtypes = dpdata.LabeledSystem.DTYPES + self.lmdb_path = "tmp_type_axis.lmdb" + type_values = DataType( + "type_values", + np.ndarray, + (Axis.NFRAMES, Axis.NTYPES), + required=False, + ) + static_type_values = DataType( + "static_type_values", + np.ndarray, + (Axis.NTYPES,), + required=False, + ) + dpdata.System.register_data_type(type_values, static_type_values) + dpdata.LabeledSystem.register_data_type(type_values, static_type_values) + self.system = dpdata.LabeledSystem( + data={ + "atom_numbs": [1, 2], + "atom_names": ["H", "O"], + "atom_types": np.array([0, 1, 1]), + "orig": np.zeros(3), + "cells": np.eye(3)[None], + "coords": np.zeros((1, 3, 3)), + "energies": np.array([-1.0]), + "type_values": np.array([[10.0, 20.0]]), + "static_type_values": np.array([100.0, 200.0]), + } + ) + + def tearDown(self): + dpdata.System.DTYPES = self.original_system_dtypes + dpdata.LabeledSystem.DTYPES = self.original_labeled_system_dtypes + if os.path.exists(self.lmdb_path): + shutil.rmtree(self.lmdb_path) + + def test_type_axes_follow_requested_type_map(self): + self.system.to("lmdb", self.lmdb_path, type_map=["H", "O", "C"]) + loaded = dpdata.LabeledSystem( + self.lmdb_path, + fmt="lmdb", + type_map=["O", "H", "C"], + mixed_type=True, + ) + self.assertEqual(loaded["atom_names"], ["O", "H", "C"]) + np.testing.assert_array_equal(loaded["type_values"], [[20.0, 10.0, 0.0]]) + np.testing.assert_array_equal(loaded["static_type_values"], [200.0, 100.0, 0.0]) + multi_loaded = dpdata.MultiSystems.from_file( + self.lmdb_path, + fmt="lmdb", + type_map=["O", "H", "C"], + mixed_type=True, + )[0] + frame_values = dict( + zip( + multi_loaded["atom_names"], + multi_loaded["type_values"][0], + ) + ) + static_values = dict( + zip( + multi_loaded["atom_names"], + multi_loaded["static_type_values"], + ) + ) + self.assertEqual(frame_values, {"C": 0.0, "H": 10.0, "O": 20.0}) + self.assertEqual(static_values, {"C": 0.0, "H": 100.0, "O": 200.0}) + + def test_type_axes_compact_with_absent_types(self): + self.system.to("lmdb", self.lmdb_path, type_map=["H", "O", "C"]) + loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb")[0] + self.assertEqual(loaded["atom_names"], ["H", "O"]) + np.testing.assert_array_equal(loaded["type_values"], [[10.0, 20.0]]) + np.testing.assert_array_equal(loaded["static_type_values"], [100.0, 200.0]) + + def test_writer_rejects_cross_system_disk_owner_change(self): + foo = DataType( + "foo", + np.ndarray, + (Axis.NFRAMES, 1), + required=False, + ) + bar = DataType( + "bar", + np.ndarray, + (Axis.NFRAMES, 1), + required=False, + deepmd_name="foo", + ) + dpdata.System.register_data_type(foo, bar) + dpdata.LabeledSystem.register_data_type(foo, bar) + first = self.system.copy() + first.data["foo"] = np.array([[1.0]]) + second = self.system.copy() + second.data["bar"] = np.array([[2.0]]) + from dpdata.formats.lmdb import dump_systems + + with self.assertRaisesRegex(LMDBError, "owned by 'foo'.*'bar'"): + dump_systems([first, second], self.lmdb_path) + + def test_registered_non_array_field_is_rejected(self): + scalar = DataType("scalar_value", float, shape=None, required=False) + dpdata.System.register_data_type(scalar) + dpdata.LabeledSystem.register_data_type(scalar) + self.system.data["scalar_value"] = 1.5 + with self.assertRaisesRegex(LMDBError, "unsupported non-array type float"): + self.system.to("lmdb", self.lmdb_path) + + if __name__ == "__main__": unittest.main() From 6bc0f7d3903771fb992be0bf13808c48c104dc96 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Wed, 15 Jul 2026 14:15:25 +0800 Subject: [PATCH 6/6] fix(lmdb): validate malformed stored metadata Wrap malformed legacy descriptors, reject invalid stored type maps, and derive legacy active types from validated atom indices. --- dpdata/formats/lmdb/format.py | 61 ++++++++++++++++++++++++++++------ tests/test_lmdb.py | 62 +++++++++++++++++++++++++++++++++++ 2 files changed, 113 insertions(+), 10 deletions(-) diff --git a/dpdata/formats/lmdb/format.py b/dpdata/formats/lmdb/format.py index 709ff32fb..0712b4244 100644 --- a/dpdata/formats/lmdb/format.py +++ b/dpdata/formats/lmdb/format.py @@ -422,6 +422,8 @@ def _decode_legacy_frame(raw: bytes) -> dict[str, Any]: msgpack.exceptions.ExtraData, msgpack.exceptions.FormatError, msgpack.exceptions.StackError, + IndexError, + KeyError, TypeError, ValueError, ) as exc: @@ -1480,13 +1482,21 @@ def from_multi_systems( return file_type_map = _meta_get(meta, "type_map") type_axis_remap: np.ndarray | None = None - if file_type_map: - names = list( - _validate_names( - [n.decode() if isinstance(n, bytes) else n for n in file_type_map], - "file type_map", + if file_type_map is not None: + if not isinstance(file_type_map, (list, tuple)): + raise LMDBMetadataError("type_map must be a sequence of names.") + try: + names = list( + _validate_names( + [ + n.decode() if isinstance(n, bytes) else n + for n in file_type_map + ], + "file type_map", + ) ) - ) + except LMDBError as exc: + raise LMDBMetadataError(str(exc)) from exc stored_ntypes = len(names) self._validate_frame_types(frames, names, meta) if type_map is not None: @@ -1635,6 +1645,7 @@ def _legacy_datasets( f"Legacy field '{name}' has incompatible schemas." ) specs[name] = spec + self._validate_legacy_type_data(data) if type_map is not None: self._remap_legacy_type_table(data, specs, type_map) datasets.append(data) @@ -1644,6 +1655,33 @@ def _legacy_datasets( ) return datasets, specs + @staticmethod + def _validate_legacy_type_data(data: dict[str, Any]) -> None: + """Validate and normalize the type table of one legacy system.""" + names = _validate_names(data["atom_names"], "legacy atom_names") + try: + atom_types = _validate_integer_types( + np.asarray(data["atom_types"]), + name="legacy atom_types", + upper_bound=len(names), + allow_virtual=False, + ) + except LMDBError as exc: + raise LMDBFrameError(str(exc)) from exc + if atom_types.ndim != 1 or atom_types.size == 0: + raise LMDBFrameError( + "legacy atom_types must be a non-empty one-dimensional array." + ) + counts = np.bincount(atom_types, minlength=len(names)).tolist() + if list(data["atom_numbs"]) != counts: + raise LMDBFrameError( + "Legacy atom_numbs is inconsistent with atom_types: " + f"expected {counts}, got {data['atom_numbs']}." + ) + data["atom_names"] = list(names) + data["atom_types"] = atom_types + data["atom_numbs"] = counts + @staticmethod def _remap_legacy_type_table( data: dict[str, Any], @@ -1654,9 +1692,13 @@ def _remap_legacy_type_table( old_names = _validate_names(data["atom_names"], "legacy atom_names") requested = _validate_names(type_map, "type_map") missing = [name for name in old_names if name not in requested] - active = { - name for name, count in zip(old_names, data["atom_numbs"]) if count > 0 - } + atom_types = _validate_integer_types( + np.asarray(data["atom_types"]), + name="legacy atom_types", + upper_bound=len(old_names), + allow_virtual=False, + ) + active = {old_names[index] for index in np.unique(atom_types)} missing_active = sorted(active.intersection(missing)) if missing_active: raise LMDBError( @@ -1667,7 +1709,6 @@ def _remap_legacy_type_table( [requested.index(name) if name in requested else -1 for name in old_names], dtype=np.int64, ) - atom_types = np.asarray(data["atom_types"], dtype=np.int64) data["atom_types"] = source_to_target[atom_types] data["atom_names"] = list(requested) data["atom_numbs"] = np.bincount( diff --git a/tests/test_lmdb.py b/tests/test_lmdb.py index afb3806b2..bc9a74753 100644 --- a/tests/test_lmdb.py +++ b/tests/test_lmdb.py @@ -467,6 +467,39 @@ def test_legacy_database_can_be_migrated_by_roundtrip(self): loaded = dpdata.MultiSystems.from_file(self.migrated_path, fmt="lmdb") self.assertEqual(loaded.get_nframes(), 3) + def test_malformed_legacy_descriptor_is_wrapped(self): + env = lmdb.open(self.legacy_path, map_size=1 << 30) + with env.begin(write=True) as txn: + malformed = { + "atom_types": { + b"nd": True, + b"shape": [3], + b"data": b"\x00" * 24, + } + } + txn.put(b"000000000000", _reference_packb(malformed)) + env.close() + with self.assertRaisesRegex(LMDBFrameError, "Cannot decode legacy LMDB frame"): + dpdata.MultiSystems.from_file(self.legacy_path, fmt="lmdb") + + def test_inconsistent_legacy_counts_are_rejected(self): + env = lmdb.open(self.legacy_path, map_size=1 << 30) + frame_fields = {"cells", "coords", "energies", "forces"} + with env.begin(write=True) as txn: + for index in range(self.water.get_nframes()): + frame = { + name: value[index] if name in frame_fields else value + for name, value in self.water.data.items() + } + frame["atom_numbs"] = [0, 3] + txn.put( + f"{index:012d}".encode(), + _reference_packb(_legacy_encode_value(frame)), + ) + env.close() + with self.assertRaisesRegex(LMDBFrameError, "atom_numbs is inconsistent"): + dpdata.MultiSystems.from_file(self.legacy_path, fmt="lmdb", type_map=["H"]) + def _write_raw_lmdb(path, frames, type_map, metadata_extra=None): """Write frames using the reference encoding (one record per frame).""" @@ -530,6 +563,35 @@ def test_type_map_remap_by_name(self): names = [str(s.data["atom_names"][t]) for t in s.data["atom_types"]] self.assertEqual(names, ["O", "H", "H"]) + def test_empty_stored_type_map_rejected(self): + frames = [ + { + "atom_types": np.array([0], dtype=np.int32), + "coords": np.zeros((1, 3)), + } + ] + _write_raw_lmdb(self.lmdb_path, frames, []) + with self.assertRaisesRegex( + LMDBMetadataError, "must contain at least one element" + ): + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) + + def test_non_sequence_stored_type_map_rejected(self): + frames = [ + { + "atom_types": np.array([0], dtype=np.int32), + "coords": np.zeros((1, 3)), + } + ] + _write_raw_lmdb( + self.lmdb_path, + frames, + ["H"], + metadata_extra={"type_map": "H"}, + ) + with self.assertRaisesRegex(LMDBMetadataError, "must be a sequence"): + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) + def test_type_map_missing_element_raises(self): frames = [ {