euler-loading 2.6.0__tar.gz → 2.8.0__tar.gz

{euler_loading-2.6.0 → euler_loading-2.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: euler-loading
-Version: 2.6.0
+Version: 2.8.0
 Summary: Multi-modal PyTorch dataloader using ds-crawler indices
 Requires-Python: >=3.9
 Requires-Dist: ds-crawler>=2.3.0

{euler_loading-2.6.0 → euler_loading-2.8.0}/README.md

@@ -4,7 +4,7 @@ Multi-modal PyTorch `Dataset` that synchronises files across arbitrary dataset m
 
 Each modality points at a directory (or `.zip` archive) that carries its own `ds-crawler.config` (or cached `output.json`).
 ds-crawler indexes the directory tree, discovers files, and exposes hierarchical metadata (path properties, calibration files, …).
-euler-loading then **intersects file IDs** across all modalities so that every sample contains exactly one file per modality. Additional hierarchical data (e.g. per-scene calibration files) can be loaded via `hierarchical_modalities`.
+euler-loading then **intersects file IDs** across all modalities so that every sample contains exactly one file per modality. Additional hierarchical data (e.g. per-scene calibration files) can be loaded via `hierarchical_modalities`. For augmentation-style datasets — many augmented files per ground-truth sample — `euler_layout` metadata lets `MultiModalDataset.from_layout(...)` load shared targets as collapsed hierarchical modalities.
 How a file is actually **loaded** (image, depth map, point cloud, …) is configurable per modality — either supply a `Callable` or let euler-loading resolve a built-in loader automatically from the ds-crawler config.
 Writer functions can be resolved the same way, so inference outputs can be written back in dataset-native formats.
 
@@ -65,6 +65,8 @@ Frozen dataclass describing one data modality.
 | `hierarchy_scope` | `str \| None` | Optional scope label for hierarchical modalities (e.g. `scene_camera`). |
 | `applies_to` | `list[str] \| None` | Optional list of regular modality names a hierarchical modality applies to. |
 | `split` | `str \| None` | Optional inline split name. Loads `.ds_crawler/split_<name>.json` from the modality root (directory or zip) and overlays it on the normal ds-crawler metadata. |
+| `cache` | `bool \| None` | Opt-in/out for in-memory caching of decoded values (only meaningful for hierarchical modalities; regular modalities are never cached). `True` keeps every distinct loaded file in a process-lifetime dict. `False` re-reads on every access. `None` (default): hierarchical → `True` because small shared calibration files benefit. |
+| `collapse_single` | `bool` | For hierarchical modalities only. When `True`, a sample that matches exactly one hierarchical file returns that loaded value directly instead of `{file_id: value}`. This is useful for shared GT tensors used by multiple augmentations of the same source sample. |
 | `metadata` | `dict[str, Any]` | Optional arbitrary metadata. Keys under `metadata["euler_loading"]` are treated as euler-loading defaults. |
 
 The loader is the **only** place where domain-specific I/O happens.
@@ -168,14 +170,22 @@ PyTorch `Dataset`. On construction it:
     ...
     "id":          str,                   # file ID (leaf only, shared across modalities)
     "full_id":     str,                   # full hierarchical path including file ID (e.g. "/scene/camera/frame")
-    "meta":        {                      # per-modality ds-crawler file entries (regular modalities only)
-        "<modality_name>": {"id": ..., "path": ..., "path_properties": ..., "basename_properties": ...},
+    "meta":        {                      # per-modality ds-crawler file entries
+        "<modality_name>": {"id": ..., "path": ..., "path_properties": ..., "basename_properties": ..., "attributes": ...},
+        ...
+    },
+    "attributes":  {                      # per-modality top-level surface for file_entry["attributes"]
+        "<modality_name>": {...},                # for regular modalities (single dict)
+        "<hierarchical_modality_name>": {        # for hierarchical modalities (one dict per matched file)
+            "<file_id>": {...},
+            ...
+        },
         ...
     },
 }
 ```
 
-Hierarchical modality results are cached so shared files are parsed only once.
+Hierarchical modality results are cached by default. This can be flipped via `Modality(..., cache=True|False)` — see the [`cache` row in the `Modality` table](#modalitypath--loadernone-metadatanone).
 
 ### `FileRecord`
 
@@ -315,6 +325,44 @@ Files across modalities are matched by these IDs, so **the directory structure m
 
 Calibration files or other per-scene/per-sequence metadata can be loaded via `hierarchical_modalities`. These files are matched to samples based on their position in the hierarchy — all files at or above a sample's hierarchy level are included and cached for efficiency.
 
+## Layout-aware augmentation loading
+
+Augmentation datasets often contain many files for the same source sample, while a target modality such as GT depth should be reused for every augmentation of that source. The recommended layout is to index the source sample as a hierarchy level and the augmentation variant as the file ID:
+
+```
+augmented_rgb/
+  scene_000000/CS_FRONT/file_id:000000000000000025/mor_10m.png
+  scene_000000/CS_FRONT/file_id:000000000000000025/mor_20m.png
+  scene_000000/CS_FRONT/file_id:000000000000000026/mor_10m.png
+  ...
+gt_depth/
+  scene_000000/CS_FRONT/file_id:000000000000000025/depth.png
+  scene_000000/CS_FRONT/file_id:000000000000000026/depth.png
+  ...
+```
+
+Both modality indexes should include an `euler_layout` addon. The augmented RGB layout declares a `sample_axis` such as `file_id` at `location="hierarchy"` and a `variant_axis` at `location="file_id"`. The GT depth layout declares the same `sample_axis` and no variant axis.
+
+```python
+from euler_loading import Modality, MultiModalDataset
+
+dataset = MultiModalDataset.from_layout(
+    {
+        "rgb_aug": Modality("/data/augmented_rgb", loader=load_rgb),
+        "depth": Modality("/data/gt_depth", loader=load_depth),
+    },
+    primary="rgb_aug",
+)
+
+sample = dataset[0]
+# sample["rgb_aug"]   – the per-augmentation RGB
+# sample["depth"]     – the GT depth for this file_id, shared across variants
+# sample["full_id"]   – e.g. "/scene_000000/CS_FRONT/file_id:000…025/mor_10m"
+# sample["id"]        – the augmentation's leaf id, e.g. "mor_10m"
+```
+
+`from_layout(...)` plans the GT modality as `hierarchical_modalities={"depth": Modality(..., collapse_single=True)}`. The hierarchical match uses the regular sample's hierarchy path, so source samples remain disambiguated by their parent hierarchy as well as by `file_id`. If a shared modality declares the same `sample_axis` but its index does not expose that axis as hierarchy, construction raises and the dataset should be re-indexed with the source sample axis in hierarchy.
+
 ## DenseDepthLoader protocol
 
 `euler_loading.DenseDepthLoader` is a `runtime_checkable` Protocol defining the loader contract for dense-depth datasets. A conforming module must expose:

{euler_loading-2.6.0 → euler_loading-2.8.0}/euler_loading/_writing.py

@@ -55,10 +55,11 @@ def _wrap_writer_save_index(
                     result["meta"] = dict(meta)
 
             if isinstance(addons, Mapping):
-                for addon_name in ("euler_train", "euler_loading"):
-                    addon = addons.get(addon_name)
+                for addon_name, addon in addons.items():
+                    if not str(addon_name).startswith("euler_"):
+                        continue
                     if isinstance(addon, Mapping):
-                        result[addon_name] = dict(addon)
+                        result[str(addon_name)] = dict(addon)
 
         index_tree = output.get("index")
         if isinstance(index_tree, Mapping):

{euler_loading-2.6.0 → euler_loading-2.8.0}/euler_loading/dataset.py

@@ -7,7 +7,7 @@ import os
 from pathlib import Path
 import zipfile
 from collections.abc import Mapping
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from typing import Any, Callable
 
 try:
@@ -25,6 +25,11 @@ from ds_crawler import (
 )
 from ds_crawler.zip_utils import get_zip_root_prefix, is_zip_path
 
+try:
+    from ds_crawler import get_layout_addon as _ds_crawler_get_layout_addon
+except ImportError:  # pragma: no cover - compatibility with older ds-crawler
+    _ds_crawler_get_layout_addon = None
+
 from ._ds_crawler_utils import load_index_output, parse_path_with_split
 from ._metadata import _build_runlog_entry, _get_ds_crawler_descriptor
 from ._resolution import (
@@ -120,6 +125,112 @@ def _get_index_meta(index_output: Mapping[str, Any]) -> dict[str, Any] | None:
     return None
 
 
+def _get_index_layout(index_output: Mapping[str, Any]) -> dict[str, Any] | None:
+    """Return the optional ``euler_layout`` addon from an index payload."""
+    if _ds_crawler_get_layout_addon is not None:
+        try:
+            return _ds_crawler_get_layout_addon(dict(index_output))
+        except Exception:
+            logger.debug("Failed to parse euler_layout addon.", exc_info=True)
+
+    head = index_output.get("head")
+    if isinstance(head, Mapping):
+        addons = head.get("addons")
+        if isinstance(addons, Mapping):
+            layout = addons.get("euler_layout")
+            if isinstance(layout, Mapping):
+                return dict(layout)
+
+    layout = index_output.get("euler_layout")
+    if isinstance(layout, Mapping):
+        return dict(layout)
+    return None
+
+
+def _layout_axis(layout: Mapping[str, Any] | None, axis: str) -> dict[str, Any] | None:
+    if not isinstance(layout, Mapping):
+        return None
+    value = layout.get(axis)
+    return dict(value) if isinstance(value, Mapping) else None
+
+
+def _layout_axis_name(layout: Mapping[str, Any] | None, axis: str) -> str | None:
+    value = _layout_axis(layout, axis)
+    name = value.get("name") if value else None
+    return name if isinstance(name, str) and name else None
+
+
+def _layout_has_variant_axis(layout: Mapping[str, Any] | None) -> bool:
+    return _layout_axis(layout, "variant_axis") is not None
+
+
+def _layout_family(layout: Mapping[str, Any] | None) -> str | None:
+    if not isinstance(layout, Mapping):
+        return None
+    family = layout.get("family")
+    return family if isinstance(family, str) and family else None
+
+
+def _layout_families_compatible(
+    left: Mapping[str, Any] | None,
+    right: Mapping[str, Any] | None,
+) -> bool:
+    left_family = _layout_family(left)
+    right_family = _layout_family(right)
+    return (
+        left_family is None
+        or right_family is None
+        or left_family == right_family
+    )
+
+
+def _index_hierarchy_separator(index_output: Mapping[str, Any]) -> str | None:
+    indexing = index_output.get("indexing")
+    hierarchy = indexing.get("hierarchy") if isinstance(indexing, Mapping) else None
+    if isinstance(hierarchy, Mapping):
+        separator = hierarchy.get("separator")
+        if isinstance(separator, str) and separator:
+            return separator
+    return None
+
+
+def _tree_has_hierarchy_axis(
+    node: Mapping[str, Any],
+    *,
+    axis_name: str,
+    separator: str,
+) -> bool:
+    prefix = f"{axis_name}{separator}"
+    children = node.get("children")
+    if not isinstance(children, Mapping):
+        return False
+    for key, child in children.items():
+        if isinstance(key, str) and key.startswith(prefix):
+            return True
+        if isinstance(child, Mapping) and _tree_has_hierarchy_axis(
+            child, axis_name=axis_name, separator=separator,
+        ):
+            return True
+    return False
+
+
+def _index_has_hierarchy_axis(
+    index_output: Mapping[str, Any],
+    *,
+    axis_name: str,
+) -> bool:
+    separator = _index_hierarchy_separator(index_output)
+    if not separator:
+        return False
+    try:
+        tree = _get_index_tree(index_output)
+    except KeyError:
+        return False
+    return _tree_has_hierarchy_axis(
+        tree, axis_name=axis_name, separator=separator,
+    )
+
+
 def _get_file_attributes(file_entry: Mapping[str, Any]) -> dict[str, Any] | None:
     attributes = file_entry.get("attributes")
     if isinstance(attributes, Mapping):
@@ -217,6 +328,22 @@ class Modality:
         split: Optional inline split name. When set, euler-loading reads the
                full ds-crawler metadata from the modality root and overlays the
                dataset payload from ``.ds_crawler/split_<name>.json``.
+        cache: Optional opt-in/out for in-memory caching of loaded values.
+               Only meaningful for hierarchical modalities; regular
+               modalities load once per sample and never cache.
+
+               - ``True``: every distinct file is decoded at most once
+                 per dataset lifetime, then served from memory.
+               - ``False``: every access re-reads and re-decodes the
+                 file.  Use this for large per-sample tensors where the
+                 cache would balloon.
+               - ``None`` (default): hierarchical modalities default to
+                 ``True`` (small calibration files benefit).
+        collapse_single: When ``True`` for a hierarchical modality, return the
+                 single matched hierarchical value directly instead of a
+                 ``{file_id: value}`` dict.  This is useful for shared GT
+                 tensors loaded once per source sample across augmentations.
+                 A sample with more than one matched file raises.
         metadata: Optional arbitrary metadata. Keys under
                   ``metadata["euler_loading"]`` are treated as
                   euler-loading-specific defaults.
@@ -232,6 +359,8 @@ class Modality:
     hierarchy_scope: str | None = None
     applies_to: list[str] | None = None
     split: str | None = None
+    cache: bool | None = None
+    collapse_single: bool = False
     metadata: dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
@@ -403,6 +532,100 @@ class MultiModalDataset(_BaseDataset):
             "hierarchical_modalities": hierarchical_modalities,
         }
 
+    @classmethod
+    def from_layout(
+        cls,
+        modalities: dict[str, Modality],
+        *,
+        primary: str | None = None,
+        transforms: list[Callable[[dict[str, Any]], dict[str, Any]]] | None = None,
+    ) -> "MultiModalDataset":
+        """Create a dataset by planning modality roles from ``euler_layout``.
+
+        When the primary modality has a variant axis (for example fog
+        augmentations grouped under ``file_id:<source>``) and an auxiliary
+        modality has the same sample axis but no variant axis, the auxiliary is
+        loaded as shared per-source data.  If the auxiliary index already
+        exposes that sample axis as a hierarchy key, it is wired as a collapsed
+        hierarchical modality.  Otherwise construction fails because the shared
+        sample key cannot be addressed unambiguously.
+
+        Existing constructor semantics remain unchanged; this helper is the
+        opt-in path for layout-aware loading.
+        """
+        if not modalities:
+            raise ValueError("At least one modality must be provided.")
+
+        indexed: dict[str, dict[str, Any]] = {}
+        layouts: dict[str, dict[str, Any] | None] = {}
+        for name, modality in modalities.items():
+            index = load_index_output(
+                modality.path,
+                split=modality.split,
+                index_dataset_from_path_fn=index_dataset_from_path,
+            )
+            indexed[name] = index
+            layouts[name] = _get_index_layout(index)
+
+        if primary is None:
+            primary = next(
+                (
+                    name for name, layout in layouts.items()
+                    if _layout_has_variant_axis(layout)
+                ),
+                next(iter(modalities)),
+            )
+        if primary not in modalities:
+            raise ValueError(
+                f"primary modality {primary!r} is not in modalities; "
+                f"known: {sorted(modalities)}"
+            )
+
+        primary_layout = layouts[primary]
+        primary_sample_axis = _layout_axis_name(primary_layout, "sample_axis")
+        primary_has_variants = _layout_has_variant_axis(primary_layout)
+
+        regular_modalities: dict[str, Modality] = {primary: modalities[primary]}
+        hierarchical_modalities: dict[str, Modality] = {}
+
+        for name, modality in modalities.items():
+            if name == primary:
+                continue
+
+            layout = layouts[name]
+            sample_axis = _layout_axis_name(layout, "sample_axis")
+            shared_with_primary = (
+                primary_has_variants
+                and primary_sample_axis is not None
+                and sample_axis == primary_sample_axis
+                and not _layout_has_variant_axis(layout)
+                and _layout_families_compatible(primary_layout, layout)
+            )
+
+            if not shared_with_primary:
+                regular_modalities[name] = modality
+                continue
+
+            if _index_has_hierarchy_axis(indexed[name], axis_name=sample_axis):
+                hierarchical_modalities[name] = replace(
+                    modality, collapse_single=True,
+                )
+                continue
+
+            raise ValueError(
+                f"Modality {name!r} shares sample_axis={sample_axis!r} with "
+                f"variant primary {primary!r}, but its ds-crawler index has "
+                "no matching hierarchy axis. Re-index it with that axis as "
+                "hierarchy so it can be loaded as a collapsed hierarchical "
+                "modality."
+            )
+
+        return cls(
+            modalities=regular_modalities,
+            hierarchical_modalities=hierarchical_modalities or None,
+            transforms=transforms,
+        )
+
     def __init__(
         self,
         modalities: dict[str, Modality],
@@ -498,10 +721,11 @@ class MultiModalDataset(_BaseDataset):
             modality_counts = {
                 name: len(lookup) for name, lookup in self._lookups.items()
             }
-            raise ValueError(
+            message = (
                 f"No common IDs found across modalities. "
                 f"Per-modality file counts: {modality_counts}"
             )
+            raise ValueError(message)
 
         for name, lookup in self._lookups.items():
             total = len(lookup)
@@ -521,7 +745,10 @@ class MultiModalDataset(_BaseDataset):
         self._zip_prefixes: dict[str, str] = {}
         self._zip_handles: dict[tuple[str, int], zipfile.ZipFile] = {}
 
-        for name, modality in list(modalities.items()) + list(self._hierarchical_modalities.items()):
+        for name, modality in (
+            list(modalities.items())
+            + list(self._hierarchical_modalities.items())
+        ):
             if is_zip_path(modality.path):
                 self._zip_modalities.add(name)
                 self._zip_prefixes[name] = get_zip_root_prefix(
@@ -529,7 +756,13 @@ class MultiModalDataset(_BaseDataset):
                 )
 
         # -- Caches ----------------------------------------------------------
-        self._hierarchical_cache: dict[str, Any] = {}
+        # Per-modality caches; ``None`` disables caching for that modality.
+        # Hierarchical default: enabled (calibration-style small files).
+        self._caches: dict[str, dict[str, Any] | None] = {}
+        for name, modality in self._hierarchical_modalities.items():
+            self._caches[name] = {} if (
+                modality.cache if modality.cache is not None else True
+            ) else None
 
         # -- Optional transform binding -------------------------------------
         for transform in self._transforms:
@@ -699,18 +932,21 @@ class MultiModalDataset(_BaseDataset):
         written_paths: dict[str, str] = {}
 
         for modality_name, value in outputs.items():
-            if modality_name not in self._lookups:
+            if modality_name in self._lookups:
+                record = self._lookups[modality_name][sample_id]
+                modality_meta = _get_index_meta(
+                    self._index_outputs[modality_name]
+                )
+            else:
                 raise KeyError(
-                    f"Modality {modality_name!r} is not a regular modality and cannot "
-                    "be addressed by sample index."
+                    f"Modality {modality_name!r} is not a regular "
+                    "modality and cannot be addressed by sample index."
                 )
 
             writer = self.get_writer(modality_name)
             destination = _resolve_output_destination(
                 output_root=output_root, modality_name=modality_name
             )
-            record = self._lookups[modality_name][sample_id]
-            modality_meta = _get_index_meta(self._index_outputs[modality_name])
             basename = Path(record.file_entry["path"]).name
             relative_path = record.file_entry["path"]
             full_id = _build_writer_full_id(
@@ -848,36 +1084,59 @@ class MultiModalDataset(_BaseDataset):
             files_by_level = self._hierarchical_lookups[name]
             matched = match_hierarchical_files(hierarchy_path, files_by_level)
             modality_meta = _get_index_meta(self._hierarchical_index_outputs[name])
+            cache = self._caches[name]
             loaded: dict[str, Any] = {}
             attrs_for_modality: dict[str, dict[str, Any]] = {}
+            entries_by_id: dict[str, dict[str, Any]] = {}
             for entry in matched:
                 entry_attributes = _get_file_attributes(entry)
-                cache_key = f"{modality.path}/{entry['path']}"
                 accepts_attributes = self._loaders_accept_attributes[name]
+                cache_key = f"{modality.path}/{entry['path']}"
                 if accepts_attributes:
                     cache_key = (
                         f"{cache_key}::attributes="
                         f"{_attributes_cache_key(entry_attributes)}"
                     )
-                if cache_key not in self._hierarchical_cache:
+                if cache is not None and cache_key in cache:
+                    value = cache[cache_key]
+                else:
                     if name in self._zip_modalities:
                         file_or_path = self._open_from_zip(
                             name, modality.path, entry["path"],
                         )
                     else:
                         file_or_path = f"{modality.path}/{entry['path']}"
-                    self._hierarchical_cache[cache_key] = _load_with_optional_attributes(
+                    value = _load_with_optional_attributes(
                         self._resolved_loaders[name],
                         file_or_path,
                         modality_meta,
                         entry_attributes,
                         accepts_attributes=accepts_attributes,
                     )
-                loaded[entry["id"]] = self._hierarchical_cache[cache_key]
+                    if cache is not None:
+                        cache[cache_key] = value
+                loaded[entry["id"]] = value
+                entries_by_id[entry["id"]] = entry
                 if entry_attributes:
                     attrs_for_modality[entry["id"]] = entry_attributes
-            sample[name] = loaded
-            attributes[name] = attrs_for_modality
+            if modality.collapse_single:
+                if len(loaded) > 1:
+                    raise ValueError(
+                        f"Hierarchical modality {name!r} was configured with "
+                        f"collapse_single=True, but sample {sample_id!r} "
+                        f"matched {len(loaded)} files."
+                    )
+                if loaded:
+                    only_id, only_value = next(iter(loaded.items()))
+                    sample[name] = only_value
+                    meta[name] = entries_by_id[only_id]
+                    attributes[name] = attrs_for_modality.get(only_id, {})
+                else:
+                    sample[name] = None
+                    attributes[name] = {}
+            else:
+                sample[name] = loaded
+                attributes[name] = attrs_for_modality
 
         sample["id"] = file_id
         sample["full_id"] = "/" + "/".join(hierarchy_path + (file_id,))

{euler_loading-2.6.0 → euler_loading-2.8.0}/package-lock.json

@@ -1,5 +1,5 @@
 {
-  "name": "commit-ok",
+  "name": "euler-lo",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {}

{euler_loading-2.6.0 → euler_loading-2.8.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "euler-loading"
-version = "2.6.0"
+version = "2.8.0"
 description = "Multi-modal PyTorch dataloader using ds-crawler indices"
 requires-python = ">=3.9"
 dependencies = [

{euler_loading-2.6.0 → euler_loading-2.8.0}/tests/test_dataset.py

@@ -586,7 +586,7 @@ class TestHierarchicalModalities:
         assert s0 == s1
 
     def test_hierarchical_cached(self):
-        """Shared hierarchical files are loaded only once."""
+        """Shared hierarchical files are loaded only once (default behavior)."""
         loader = MagicMock(return_value="parsed")
         rgb_index = _deep_regular_index(["f001", "f002", "f003"])
         hier_index = _hierarchical_intrinsics_index()
@@ -614,6 +614,36 @@ class TestHierarchicalModalities:
         # The intrinsics file should have been loaded exactly once.
         loader.assert_called_once()
 
+    def test_hierarchical_cache_false_reloads_every_access(self):
+        """Opt-out via cache=False reloads even shared hierarchical files."""
+        loader = MagicMock(return_value="parsed")
+        rgb_index = _deep_regular_index(["f001", "f002", "f003"])
+        hier_index = _hierarchical_intrinsics_index()
+
+        def mock_index(path, **kw):
+            return rgb_index if "rgb" in path else hier_index
+
+        with patch(
+            "euler_loading.dataset.index_dataset_from_path",
+            side_effect=mock_index,
+        ):
+            ds = MultiModalDataset(
+                modalities={
+                    "rgb": Modality("/data/rgb", loader=dummy_loader),
+                },
+                hierarchical_modalities={
+                    "cam_intrinsics": Modality(
+                        "/data/intrinsics", loader=loader, cache=False,
+                    ),
+                },
+            )
+
+        for i in range(3):
+            _ = ds[i]
+
+        # 3 accesses with caching disabled → 3 loads of the same file.
+        assert loader.call_count == 3
+
     def test_no_hierarchy_overlap_returns_empty_dict(self):
         """When hierarchical modality has no matching ancestors, result is {}."""
         rgb_index = _deep_regular_index(["f001"])
@@ -801,7 +831,7 @@ class TestAugmentedRgbWithHierarchicalDepth:
         # 2 file-ids x 2 augs per file-id
         assert len(ds) == 4
 
-    def test_each_sample_has_depth_keyed_by_file_id(self):
+    def test_each_sample_has_depth_grouped_by_file_id(self):
         ds = self._make()
         for i in range(len(ds)):
             sample = ds[i]
@@ -888,6 +918,117 @@ class TestAugmentedRgbWithHierarchicalDepth:
         assert loader.call_count == 2
 
 
+# ---------------------------------------------------------------------------
+# Layout-aware augmentation planning
+# ---------------------------------------------------------------------------
+
+
+def _augmented_variant_index(*, separator: str = ":") -> dict[str, Any]:
+    """Augmented index with ``children[file_id:<id>]`` source-sample level."""
+    return {
+        "indexing": {
+            "hierarchy": {"separator": separator},
+            "id": {"join_char": "+"},
+        },
+        "dataset": {
+            "children": {
+                f"file_id{separator}abc": {
+                    "files": [
+                        _make_file("aug-aug_1", "abc/aug_1.png"),
+                        _make_file("aug-aug_2", "abc/aug_2.png"),
+                    ]
+                },
+                f"file_id{separator}xyz": {
+                    "files": [
+                        _make_file("aug-aug_1", "xyz/aug_1.png"),
+                        _make_file("aug-aug_2", "xyz/aug_2.png"),
+                    ]
+                },
+            }
+        },
+    }
+
+
+def _with_layout(
+    index: dict[str, Any],
+    *,
+    variant: bool,
+    family: str = "demo",
+) -> dict[str, Any]:
+    result = json.loads(json.dumps(index))
+    layout: dict[str, Any] = {
+        "version": "1.0",
+        "family": family,
+        "sample_axis": {"name": "file_id", "location": "hierarchy"},
+    }
+    if variant:
+        layout["variant_axis"] = {
+            "name": "fog_augmentation",
+            "location": "file_id",
+        }
+    result["head"] = {
+        "contract": {"kind": "dataset_head", "version": "1.0"},
+        "dataset": {"id": "demo", "name": "Demo"},
+        "modality": {"key": "rgb" if variant else "depth", "meta": {}},
+        "addons": {"euler_layout": layout},
+    }
+    result.setdefault(
+        "indexing",
+        {"hierarchy": {"separator": ":"}, "id": {"join_char": "+"}},
+    )
+    return result
+
+
+class TestLayoutAwareDatasetPlanning:
+    def test_from_layout_collapses_shared_hierarchical_modality(self):
+        rgb_index = _with_layout(_augmented_variant_index(), variant=True)
+        depth_index = _with_layout(
+            _per_id_depth_hierarchical_index(), variant=False,
+        )
+
+        def mock_index(path, **kw):
+            return rgb_index if "rgb" in path else depth_index
+
+        with patch(
+            "euler_loading.dataset.index_dataset_from_path",
+            side_effect=mock_index,
+        ):
+            ds = MultiModalDataset.from_layout(
+                {
+                    "rgb_aug": Modality("/data/rgb_aug", loader=dummy_loader),
+                    "depth": Modality("/data/depth", loader=dummy_loader),
+                },
+                primary="rgb_aug",
+            )
+
+        sample = ds[0]
+
+        assert len(ds) == 4
+        assert isinstance(sample["depth"], str)
+        assert sample["depth"].startswith("loaded:/data/depth/")
+        assert isinstance(sample["meta"]["depth"], dict)
+
+    def test_from_layout_rejects_flat_shared_modality(self):
+        rgb_index = _with_layout(_augmented_variant_index(), variant=True)
+        depth_index = _with_layout(_flat_index("depth", ["abc", "xyz"]), variant=False)
+
+        def mock_index(path, **kw):
+            return rgb_index if "rgb" in path else depth_index
+
+        with patch(
+            "euler_loading.dataset.index_dataset_from_path",
+            side_effect=mock_index,
+        ):
+            with pytest.raises(ValueError, match="no matching hierarchy axis"):
+                MultiModalDataset.from_layout(
+                    {
+                        "rgb_aug": Modality("/data/rgb_aug", loader=dummy_loader),
+                        "depth": Modality("/data/depth", loader=dummy_loader),
+                    },
+                    primary="rgb_aug",
+                )
+
+
 # ---------------------------------------------------------------------------
 # Per-file ``attributes`` field exposed on samples
 # ---------------------------------------------------------------------------
@@ -939,7 +1080,7 @@ class TestSampleAttributes:
         assert "attributes" in sample
         assert isinstance(sample["attributes"], dict)
 
-    def test_top_level_attributes_keyed_by_modality(self):
+    def test_top_level_attributes_grouped_by_modality(self):
         ds = self._make()
         sample = ds[0]
         assert sample["attributes"]["rgb"] == {"weight": 0.42, "src": "blender"}
@@ -1159,7 +1300,7 @@ class TestHierarchicalModalityAttributes:
                 },
             )
 
-    def test_hierarchical_attributes_keyed_by_file_id(self):
+    def test_hierarchical_attributes_grouped_by_file_id(self):
         ds = self._make()
         sample = ds[0]
         assert sample["attributes"]["cam_intrinsics"] == {