feat(data): support recursive LMDB inputs and weighted training groups

deepmd/common.py

@@ -29,8 +29,9 @@
 from deepmd.env import (
     GLOBAL_NP_FLOAT_PRECISION,
 )
-from deepmd.utils.path import (
-    DPPath,
+from deepmd.utils.data_path import (
+    expand_sys_str,
+    rglob_sys_str,
 )
 
 __all__ = [
@@ -41,6 +42,7 @@
     "get_np_precision",
     "j_loader",
     "make_default_mesh",
+    "rglob_sys_str",
     "select_idx_map",
 ]
 
@@ -188,50 +190,6 @@ def j_loader(filename: str | Path) -> dict[str, Any]:
         raise TypeError("config file must be json, or yaml/yml")
 
 
-def expand_sys_str(root_dir: str | Path) -> list[str]:
-    """Recursively iterate over directories taking those that contain `type.raw` file.
-
-    Parameters
-    ----------
-    root_dir : Union[str, Path]
-        starting directory
-
-    Returns
-    -------
-    list[str]
-        list of string pointing to system directories
-    """
-    root_dir = DPPath(root_dir)
-    matches = [str(d) for d in root_dir.rglob("*") if (d / "type.raw").is_file()]
-    if (root_dir / "type.raw").is_file():
-        matches.append(str(root_dir))
-    return matches
-
-
-def rglob_sys_str(root_dir: str, patterns: list[str]) -> list[str]:
-    """Recursively iterate over directories taking those that contain `type.raw` file.
-
-    Parameters
-    ----------
-    root_dir : str, Path
-        starting directory
-    patterns : list[str]
-        list of glob patterns to match directories
-
-    Returns
-    -------
-    list[str]
-        list of string pointing to system directories
-    """
-    root_dir = Path(root_dir)
-    matches = []
-    for pattern in patterns:
-        matches.extend(
-            [str(d) for d in root_dir.rglob(pattern) if (d / "type.raw").is_file()]
-        )
-    return list(set(matches))  # remove duplicates
-
-
 def get_np_precision(precision: "_PRECISION") -> np.dtype:
     """Get numpy precision constant from string.
 

deepmd/dpmodel/train/data.py

@@ -14,6 +14,14 @@
     Any,
 )
 
+from deepmd.utils.data import (
+    DataRequirementItem,
+    min_pair_dist_requirement,
+)
+from deepmd.utils.data_config import (
+    DataConfig,
+    data_group_configs,
+)
 from deepmd.utils.stat_file import (
     StatFileMode,
     StatFileSpec,
@@ -37,7 +45,7 @@ class TrainingTaskConfig:
 
     key: str
     model_params: Mapping[str, Any]
-    training_data_params: Mapping[str, Any]
+    training_data_params: DataConfig
     validation_data_params: Mapping[str, Any] | None
     stat_file: str | None
     valid_numb_batch: int
@@ -110,6 +118,33 @@ def make_task_maps(
     return training_data, validation_data, stat_data
 
 
+def add_training_data_requirements(
+    data: Any,
+    config: DataConfig,
+    model_params: Mapping[str, Any],
+    requirements: list[DataRequirementItem],
+    register: Callable[[Any, list[DataRequirementItem]], None],
+) -> bool:
+    """Register each group's labels and its own derived pair-clearance margin.
+
+    The returned flag enables margin filtering in the trainer and statistics
+    collector. Distances belong to individual groups' requirements, not to
+    that task-wide enable flag.
+    """
+    groups = data_group_configs(config)
+    datasets = data.children if len(groups) > 1 else (data,)
+    filter_enabled = []
+    for dataset, group in zip(datasets, groups, strict=True):
+        pair_filter = min_pair_dist_requirement(
+            model_params, float(group.get("min_pair_dist", 0.0))
+        )
+        register(dataset, requirements + ([] if pair_filter is None else [pair_filter]))
+        filter_enabled.append(pair_filter is not None)
+    if len(groups) > 1:
+        data.set_filter_enabled(filter_enabled)
+    return any(filter_enabled)
+
+
 def print_data_summaries(
     training_data: Mapping[str, Any],
     validation_data: Mapping[str, Any | None],

deepmd/dpmodel/train/schedule.py

@@ -88,8 +88,9 @@ def resolve_step_schedule(
         Task keys in the order used by the trainer.  Single-task runs pass the
         one key under which their data is registered.
     training_data : Mapping[str, Any]
-        Training data of every task, keyed like ``model_keys``.  Only consulted
-        when multi-task sampling weights fall back to the per-task data size.
+        Training data of every task, keyed like ``model_keys``. Only consulted
+        when multi-task sampling weights default to native data sizes:
+        NumPy/HDF5 system counts or retained LMDB frame counts.
     epoch_length : Callable[[str], int]
         Maps a task key to the number of steps *this rank* performs during one
         epoch of that task.  Backends whose data pipeline is sharded across