feat(data): add dpdata format conversion

deepmd/dpmodel/train/data.py

@@ -14,6 +14,9 @@
     Any,
 )
 
+from deepmd.utils.data_system import (
+    close_data_systems,
+)
 from deepmd.utils.stat_file import (
     StatFileMode,
     StatFileSpec,
@@ -98,15 +101,23 @@ def make_task_maps(
     config: Mapping[str, Any],
     factory: Callable[[TrainingTaskConfig], tuple[Any, Any | None, Any | None]],
 ) -> tuple[dict[str, Any], dict[str, Any | None], dict[str, Any | None]]:
-    """Build training, validation, and stat maps from normalized task configs."""
+    """Build task maps, closing completed data objects if a later task fails.
+
+    Factories own partial results until they return and must clean those up
+    themselves. Successful maps transfer ownership to the caller.
+    """
     training_data: dict[str, Any] = {}
     validation_data: dict[str, Any | None] = {}
     stat_data: dict[str, Any | None] = {}
-    for task_config in iter_training_task_configs(config):
-        train_item, valid_item, stat_item = factory(task_config)
-        training_data[task_config.key] = train_item
-        validation_data[task_config.key] = valid_item
-        stat_data[task_config.key] = stat_item
+    try:
+        for task_config in iter_training_task_configs(config):
+            train_item, valid_item, stat_item = factory(task_config)
+            training_data[task_config.key] = train_item
+            validation_data[task_config.key] = valid_item
+            stat_data[task_config.key] = stat_item
+    except BaseException:
+        close_data_systems(training_data, validation_data)
+        raise
     return training_data, validation_data, stat_data
 
 

deepmd/dpmodel/utils/lmdb_data.py

@@ -54,6 +54,9 @@
     DataRequirementItem,
     DataRequirementSourcePolicy,
 )
+from deepmd.utils.probability import (
+    prob_sys_size_ext,
+)
 
 log = logging.getLogger(__name__)
 
@@ -393,7 +396,9 @@ def _raw_frame_availability(
                     explicit_true |= bit
         else:
             bit = key_bits.get(name)
-            if bit is not None:
+            if bit is not None and (
+                name != "box" or not np.allclose(_decode_value(value), 0.0)
+            ):
                 present |= bit
     return present & (~explicit_known | explicit_true)
 
@@ -737,6 +742,11 @@ def _frame_source_available(frame: dict[str, Any], key: str) -> bool:
     source_present = _is_encoded_array(value) or isinstance(
         value, (np.ndarray, np.generic, int, float, bool)
     )
+    if source_present and key == "box":
+        # A zero cell is DeePMD's non-periodic sentinel, not an available
+        # periodic box. Treating it as present would mix PBC and non-PBC
+        # frames under one scalar find_box flag.
+        source_present = not np.allclose(_decode_value(value), 0.0)
     find_key = f"find_{key}"
     if find_key in frame:
         return source_present and bool(
@@ -1970,8 +1980,17 @@ def __init__(
         self,
         lmdb_path: str,
         type_map: list[str],
-        batch_size: int | str = "auto",
+        batch_size: int | str | Sequence[int | str] = "auto",
     ) -> None:
+        if isinstance(batch_size, (Sequence, np.ndarray)) and not isinstance(
+            batch_size, str
+        ):
+            if len(batch_size) != 1:
+                raise ValueError(
+                    "One LMDB path is one training dataset and therefore "
+                    "requires exactly one batch_size value."
+                )
+            batch_size = batch_size[0]
         self.lmdb_path = str(Path(lmdb_path).resolve())
         self._type_map = type_map
         # Read before opening the frame-serving environment, which disables
@@ -2235,14 +2254,8 @@ def get_batch_size_for_nloc(self, nloc: int) -> int:
     def __len__(self) -> int:
         return self.nframes
 
-    def __getitem__(self, index: int) -> dict[str, Any]:
-        """Read frame from LMDB, decode, remap keys, return dict of numpy arrays.
-
-        ``index`` is a dataset-level index in ``[0, len(self))``. Under
-        ``filter:N`` the LMDB key space may have gaps (dropped frames), so
-        we translate through ``self._retained_keys`` before hitting LMDB.
-        """
-        self._data_requirements_frozen = True
+    def _read_frame(self, index: int) -> dict[str, Any]:
+        """Decode one frame without changing requirement-registration state."""
         if index < 0 or index >= self.nframes:
             raise IndexError(f"dataset index {index} out of range [0, {self.nframes})")
         original_key = int(self._retained_keys[index])
@@ -2259,6 +2272,25 @@ def __getitem__(self, index: int) -> dict[str, Any]:
             copy_arrays=True,
         )
 
+    def peek_frame(self, index: int) -> dict[str, Any]:
+        """Inspect one frame without freezing later requirement registration.
+
+        Structural probes such as periodic-boundary detection happen before a
+        model supplies its label requirements. They may inspect a frame, but
+        must not turn that inspection into the first training read.
+        """
+        return self._read_frame(index)
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        """Read and decode one frame, freezing the registered data contract.
+
+        ``index`` is a dataset-level index in ``[0, len(self))``. Under
+        ``filter:N`` the LMDB key space may have gaps, which :meth:`_read_frame`
+        translates through ``self._retained_keys``.
+        """
+        self._data_requirements_frozen = True
+        return self._read_frame(index)
+
     def original_keys(self, indices: Sequence[int]) -> list[int]:
         """Translate dataset indices to original integer LMDB keys."""
         keys: list[int] = []
@@ -2854,17 +2886,17 @@ def compute_block_targets(
         Each element is ``(system_indices_in_block, target_frame_count)``.
         Returns empty list if no expansion is needed (all targets == actual).
     """
-    from deepmd.utils.data_system import (
-        prob_sys_size_ext,
-    )
-
-    # Parse block definitions from the auto_prob string
-    # Format: "prob_sys_size;stt:end:weight;stt:end:weight;..."
-    block_str = auto_prob_style.split(";")[1:]
+    # ``prob_uniform`` is one equal-weight block per original system. The
+    # extended ``prob_sys_size`` form names arbitrary ranges explicitly.
     blocks: list[tuple[int, int, float]] = []
-    for part in block_str:
-        stt, end, weight = part.split(":")
-        blocks.append((int(stt), int(end), float(weight)))
+    if auto_prob_style == "prob_uniform":
+        blocks = [(system_id, system_id + 1, 1.0) for system_id in range(nsystems)]
+    elif auto_prob_style.startswith("prob_sys_size"):
+        for part in auto_prob_style.split(";")[1:]:
+            stt, end, weight = part.split(":")
+            blocks.append((int(stt), int(end), float(weight)))
+    else:
+        raise RuntimeError(f"Unknown auto prob style: {auto_prob_style}")
 
     # A bare ``prob_sys_size`` names no blocks: it asks for a probability
     # proportional to system size, which is what sampling the merged frames
@@ -2905,13 +2937,20 @@ def compute_block_targets(
             f"0 frames, likely after filter:N): {dropped}. Remaining block "
             "weights will be renormalised to sum to 1.0."
         )
-        auto_prob_style = "prob_sys_size;" + ";".join(
-            f"{stt}:{end}:{weight}" for stt, end, weight in nonempty
-        )
         blocks = nonempty
 
-    # Compute per-system probabilities using the standard function
-    sys_probs = prob_sys_size_ext(auto_prob_style, nsystems, system_nframes)
+    block_weights = np.asarray([weight for _, _, weight in blocks], dtype=float)
+    if not np.all(np.isfinite(block_weights)):
+        raise ValueError("block weights must be finite")
+    if np.any(block_weights < 0):
+        raise ValueError("the weight of a block should be no less than 0")
+    total_block_weight = np.sum(block_weights)
+    if total_block_weight <= 0:
+        raise ValueError("the sum of block weights should be greater than 0")
+    keywords = "prob_sys_size;" + ";".join(
+        f"{stt}:{end}:{weight}" for stt, end, weight in blocks
+    )
+    sys_probs = prob_sys_size_ext(keywords, nsystems, system_nframes)
 
     # Group systems by block, compute block-level frames and prob
     block_info: list[tuple[list[int], int, float]] = []  # (sys_ids, frames, prob)
@@ -3673,22 +3712,29 @@ def make_neighbor_stat_data(
     )
 
     reader = LmdbDataReader(lmdb_path, type_map=type_map)
-    nframes = len(reader)
-    rng = np.random.RandomState(42)
-    if nframes > max_frames:
-        indices = np.sort(rng.choice(nframes, max_frames, replace=False))
-    else:
-        indices = np.arange(nframes, dtype=np.int64)
-
-    # Read sampled frames, group by nloc
-    nloc_frames: dict[int, list[tuple[np.ndarray, np.ndarray, np.ndarray | None]]] = {}
-    for idx in indices:
-        frame = reader[int(idx)]
-        atype = frame["atype"]
-        nloc = len(atype)
-        nloc_frames.setdefault(nloc, []).append(
-            (frame["coord"], atype, frame.get("box"))
-        )
+    try:
+        nframes = len(reader)
+        rng = np.random.RandomState(42)
+        if nframes > max_frames:
+            indices = np.sort(rng.choice(nframes, max_frames, replace=False))
+        else:
+            indices = np.arange(nframes, dtype=np.int64)
+
+        # The copied arrays remain valid after the reader is closed, so this
+        # helper does not leave an LMDB transaction or mmap alive in callers.
+        nloc_frames: dict[
+            int, list[tuple[np.ndarray, np.ndarray, np.ndarray | None]]
+        ] = {}
+        for idx in indices:
+            frame = reader[int(idx)]
+            atype = frame["atype"]
+            nloc = len(atype)
+            nloc_frames.setdefault(nloc, []).append(
+                (frame["coord"], atype, frame.get("box"))
+            )
+        ntypes = len(type_map) if type_map else reader._ntypes
+    finally:
+        reader.close()
 
     # Build per-nloc data_system proxies
     data_systems = []
@@ -3710,7 +3756,6 @@ def make_neighbor_stat_data(
         data_systems.append(proxy)
         system_dirs.append(label)
 
-    ntypes = len(type_map) if type_map else reader._ntypes
     return SimpleNamespace(
         system_dirs=system_dirs,
         data_systems=data_systems,
@@ -3888,19 +3933,17 @@ def _read_frames(self, frame_indices: Sequence[int]) -> list[dict[str, Any]]:
                 )
         return frames
 
-    def __del__(self) -> None:
-        """Release the LMDB environment ref-count on garbage collection.
-
-        The count is released only once, and only if construction got as far
-        as taking it: an instance that failed earlier holds no reference, and
-        releasing one it never took would close the environment underneath
-        whichever reader does hold it.
-        """
+    def close(self) -> None:
+        """Release the LMDB environment ref-count idempotently."""
         if getattr(self, "_env", None) is None:
             return
         self._env = None
         _close_lmdb(self.lmdb_path)
 
+    def __del__(self) -> None:
+        """Release the LMDB environment ref-count on garbage collection."""
+        self.close()
+
     @property
     def nloc_groups(self) -> dict[int, np.ndarray]:
         """Nloc → the LMDB frame indices retained for that atom count."""
@@ -4263,14 +4306,33 @@ def __init__(
         lmdb_test_data: "LmdbTestData",
         nloc: int,
         frame_indices: Sequence[int] | None = None,
+        *,
+        pbc: bool | None = None,
+        stat_groups: dict[str, Sequence[int]] | None = None,
     ) -> None:
         self._inner = lmdb_test_data
         self._nloc = nloc
         self._frame_indices = frame_indices
+        self._pbc = pbc
+        self._stat_groups = stat_groups or {}
+        self.dirs = list(self._stat_groups)
 
     def __getattr__(self, name: str) -> Any:
         return getattr(self._inner, name)
 
+    @property
+    def pbc(self) -> bool:
+        """Whether every frame represented by this view is periodic."""
+        return self._inner.pbc if self._pbc is None else self._pbc
+
+    def get_natoms(self) -> int:
+        """Return the fixed atom count of this stack-compatible view."""
+        return self._nloc
+
+    def _load_set(self, set_name: str) -> dict[str, Any]:
+        """Load one bounded neighbor-stat chunk from this view."""
+        return self._inner.get_test_by_indices(self._stat_groups[str(set_name)])
+
     def get_test(self) -> dict[str, Any]:
         if self._frame_indices is not None:
             return self._inner.get_test_by_indices(self._frame_indices)

deepmd/entrypoints/test.py

@@ -111,7 +111,12 @@ def test(
         else:
             systems = [str((root / Path(ss)).resolve()) for ss in systems]
         patterns = data_params.get("rglob_patterns", None)
-        all_sys = process_systems(systems, patterns=patterns)
+        all_sys = process_systems(
+            systems,
+            patterns=patterns,
+            fmt=data_params.get("format"),
+            out_fmt=data_params.get("out_format", data_params.get("output_format")),
+        )
     elif valid_json is not None:
         jdata = j_loader(valid_json)
         jdata = update_deepmd_input(jdata)
@@ -125,7 +130,12 @@ def test(
         else:
             systems = [str((root / Path(ss)).resolve()) for ss in systems]
         patterns = data_params.get("rglob_patterns", None)
-        all_sys = process_systems(systems, patterns=patterns)
+        all_sys = process_systems(
+            systems,
+            patterns=patterns,
+            fmt=data_params.get("format"),
+            out_fmt=data_params.get("out_format", data_params.get("output_format")),
+        )
     elif datafile is not None:
         with open(datafile) as datalist:
             all_sys = datalist.read().splitlines()