Update dependency versions and clean up dataset template methods

- Bump `multidimio` version to `1.2.0rc0` and restrict `zarr` version to `>=3.1.3,<=3.1.6` for compatibility.
- Remove the `include_raw_headers` parameter from dataset building methods in `DatasetFactory` and `AbstractDatasetTemplate`, simplifying the API.
- Refactor coordinate specification methods in various dataset templates to ensure consistent usage of `CoordinateSpec` and improve code clarity.
- Introduce a new function `maybe_add_raw_headers` to conditionally append raw headers to datasets based on experimental settings, enhancing the ingestion pipeline's flexibility.

Author: BrianMichell

src/mdio/builder/templates/base.py

@@ -20,11 +20,11 @@
 from mdio.core.utils_write import MAX_COORDINATES_BYTES
 from mdio.core.utils_write import MAX_SIZE_LIVE_MASK
 from mdio.core.utils_write import get_constrained_chunksize
+from mdio.ingestion.schema_resolver import CoordinateSpec
 
 if TYPE_CHECKING:
     from mdio.builder.schemas.v1.dataset import Dataset
     from mdio.builder.templates.types import SeismicDataDomain
-    from mdio.ingestion.schema_resolver import CoordinateSpec
 
 
 class AbstractDatasetTemplate(ABC):
@@ -73,8 +73,6 @@ def _repr_html_(self) -> str:
 
     def declare_coordinate_specs(self) -> tuple[CoordinateSpec, ...]:
         """Return the non-dimension coordinate specs (name, dims, dtype) for this template."""
-        from mdio.ingestion.schema_resolver import CoordinateSpec
-
         specs = [
             CoordinateSpec(
                 name=coord_name,
@@ -98,16 +96,13 @@ def build_dataset(
         name: str,
         sizes: tuple[int, ...],
         header_dtype: StructuredType = None,
-        include_raw_headers: bool = False,
     ) -> Dataset:
         """Template method that builds the dataset.
 
         Args:
             name: The name of the dataset.
             sizes: The sizes of the dimensions.
             header_dtype: Optional structured headers for the dataset.
-            include_raw_headers: Whether to include raw binary headers variable. Only supported
-                in Zarr v3 format. Defaults to False.
 
         Returns:
             Dataset: The constructed dataset
@@ -125,9 +120,6 @@ def build_dataset(
         if header_dtype is not None:
             self._add_trace_headers(header_dtype)
 
-        if include_raw_headers:
-            self._add_raw_headers()
-
         return self._builder.build()
 
     def add_units(self, units: dict[str, AllUnitModel]) -> None:

src/mdio/builder/templates/seismic_2d_cdp.py

@@ -13,6 +13,7 @@
 from mdio.builder.templates.types import SeismicDataDomain
 from mdio.core.utils_write import MAX_COORDINATES_BYTES
 from mdio.core.utils_write import get_constrained_chunksize
+from mdio.ingestion.schema_resolver import CoordinateSpec
 
 
 class Seismic2DCdpGathersTemplate(AbstractDatasetTemplate):
@@ -39,9 +40,8 @@ def _name(self) -> str:
     def _load_dataset_attributes(self) -> dict[str, Any]:
         return {"surveyType": "2D", "gatherType": "cdp"}
 
-    def declare_coordinate_specs(self) -> tuple[Any, ...]:
-        from mdio.ingestion.schema_resolver import CoordinateSpec
-
+    def declare_coordinate_specs(self) -> tuple[CoordinateSpec, ...]:
+        """Declare CDP-indexed X/Y coordinates for the 2D CDP gathers template."""
         return (
             CoordinateSpec(name="cdp_x", dimensions=("cdp",), dtype=ScalarType.FLOAT64),
             CoordinateSpec(name="cdp_y", dimensions=("cdp",), dtype=ScalarType.FLOAT64),

src/mdio/builder/templates/seismic_3d_cdp.py

@@ -13,6 +13,7 @@
 from mdio.builder.templates.types import SeismicDataDomain
 from mdio.core.utils_write import MAX_COORDINATES_BYTES
 from mdio.core.utils_write import get_constrained_chunksize
+from mdio.ingestion.schema_resolver import CoordinateSpec
 
 
 class Seismic3DCdpGathersTemplate(AbstractDatasetTemplate):
@@ -39,9 +40,8 @@ def _name(self) -> str:
     def _load_dataset_attributes(self) -> dict[str, Any]:
         return {"surveyType": "3D", "gatherType": "cdp"}
 
-    def declare_coordinate_specs(self) -> tuple[Any, ...]:
-        from mdio.ingestion.schema_resolver import CoordinateSpec
-
+    def declare_coordinate_specs(self) -> tuple[CoordinateSpec, ...]:
+        """Declare inline/crossline-indexed X/Y coordinates for the 3D CDP gathers template."""
         return (
             CoordinateSpec(name="cdp_x", dimensions=("inline", "crossline"), dtype=ScalarType.FLOAT64),
             CoordinateSpec(name="cdp_y", dimensions=("inline", "crossline"), dtype=ScalarType.FLOAT64),

src/mdio/builder/templates/seismic_3d_coca.py

@@ -11,6 +11,7 @@
 from mdio.builder.templates.types import SeismicDataDomain
 from mdio.core.utils_write import MAX_COORDINATES_BYTES
 from mdio.core.utils_write import get_constrained_chunksize
+from mdio.ingestion.schema_resolver import CoordinateSpec
 
 
 class Seismic3DCocaGathersTemplate(AbstractDatasetTemplate):
@@ -30,9 +31,8 @@ def _name(self) -> str:
     def _load_dataset_attributes(self) -> dict[str, Any]:
         return {"surveyType": "3D", "gatherType": "common_offset_common_azimuth"}
 
-    def declare_coordinate_specs(self) -> tuple[Any, ...]:
-        from mdio.ingestion.schema_resolver import CoordinateSpec
-
+    def declare_coordinate_specs(self) -> tuple[CoordinateSpec, ...]:
+        """Declare inline/crossline-indexed X/Y coordinates for the 3D CoCA gathers template."""
         return (
             CoordinateSpec(name="cdp_x", dimensions=("inline", "crossline"), dtype=ScalarType.FLOAT64),
             CoordinateSpec(name="cdp_y", dimensions=("inline", "crossline"), dtype=ScalarType.FLOAT64),

src/mdio/builder/templates/seismic_3d_obn.py

@@ -12,6 +12,7 @@
 from mdio.builder.templates.types import SeismicDataDomain
 from mdio.core.utils_write import MAX_COORDINATES_BYTES
 from mdio.core.utils_write import get_constrained_chunksize
+from mdio.ingestion.schema_resolver import CoordinateSpec
 
 
 class Seismic3DObnReceiverGathersTemplate(AbstractDatasetTemplate):
@@ -56,9 +57,8 @@ def _name(self) -> str:
     def _load_dataset_attributes(self) -> dict[str, Any]:
         return {"surveyType": "3D", "gatherType": "common_receiver"}
 
-    def declare_coordinate_specs(self) -> tuple[Any, ...]:
-        from mdio.ingestion.schema_resolver import CoordinateSpec
-
+    def declare_coordinate_specs(self) -> tuple[CoordinateSpec, ...]:
+        """Declare receiver- and shot-indexed coordinates for the 3D OBN receiver gathers template."""
         receiver_dim = ("receiver",)
         shot_dims = ("shot_line", "gun", "shot_index")
         return (

src/mdio/builder/templates/seismic_3d_streamer_field.py

@@ -12,6 +12,7 @@
 from mdio.builder.templates.types import SeismicDataDomain
 from mdio.core.utils_write import MAX_COORDINATES_BYTES
 from mdio.core.utils_write import get_constrained_chunksize
+from mdio.ingestion.schema_resolver import CoordinateSpec
 
 
 class Seismic3DStreamerFieldRecordsTemplate(AbstractDatasetTemplate):
@@ -44,9 +45,8 @@ def _name(self) -> str:
     def _load_dataset_attributes(self) -> dict[str, Any]:
         return {"surveyDimensionality": "3D", "gatherType": "common_source"}
 
-    def declare_coordinate_specs(self) -> tuple[Any, ...]:
-        from mdio.ingestion.schema_resolver import CoordinateSpec
-
+    def declare_coordinate_specs(self) -> tuple[CoordinateSpec, ...]:
+        """Declare shot- and receiver-indexed coordinates for the 3D streamer field records template."""
         shot_dims = ("sail_line", "gun", "shot_index")
         receiver_dims = ("sail_line", "gun", "shot_index", "cable", "channel")
         return (
@@ -91,9 +91,7 @@ def _add_coordinates(self) -> None:
         # Chunk grids for shot-indexed (3D) and receiver-indexed (5D) non-dim coordinates.
         shot_chunk_shape = get_constrained_chunksize(self._dim_sizes[:3], ScalarType.FLOAT64, MAX_COORDINATES_BYTES)
         chunk_grid_3d = RegularChunkGrid(configuration=RegularChunkShape(chunk_shape=shot_chunk_shape))
-        receiver_chunk_shape = get_constrained_chunksize(
-            self._dim_sizes[:5], ScalarType.FLOAT64, MAX_COORDINATES_BYTES
-        )
+        receiver_chunk_shape = get_constrained_chunksize(self._dim_sizes[:5], ScalarType.FLOAT64, MAX_COORDINATES_BYTES)
         chunk_grid_5d = RegularChunkGrid(configuration=RegularChunkShape(chunk_shape=receiver_chunk_shape))
 
         compressor = Blosc(cname=BloscCname.zstd)

src/mdio/ingestion/_raw_headers_experimental.py

@@ -4,20 +4,21 @@
 that removing the feature is a single-file delete plus the (small) handful of call
 sites that import from here. The behaviour:
 
-* ``should_include_raw_headers()`` decides whether the ``raw_headers`` variable
-  should be added to the dataset for this run. Honours the env-driven
-  :class:`mdio.core.config.MDIOSettings` and the active Zarr format.
+* ``maybe_add_raw_headers()`` decides whether the ``raw_headers`` variable should
+  be added to the dataset for this run, and if so appends it via the template's
+  internal builder. Honours the env-driven :class:`mdio.core.config.MDIOSettings`
+  and the active Zarr format.
 * ``attach_raw_binary_header()`` adds the base64-encoded binary file header to a
   segy_file_header attribute dict, if (and only if) the feature is enabled.
 
 Removal plan
 ------------
 Delete this module and:
 
-#. Drop the call to :func:`should_include_raw_headers` in
-   :mod:`mdio.ingestion.pipeline` (and the resulting ``include_raw_headers``
-   kwarg threaded through :class:`~mdio.ingestion.dataset_factory.DatasetFactory`
-   / :meth:`AbstractDatasetTemplate.build_dataset` / ``_add_raw_headers``).
+#. Drop the call to :func:`maybe_add_raw_headers` in
+   :mod:`mdio.ingestion.pipeline`.
+#. Drop the ``_add_raw_headers`` method on
+   :class:`~mdio.builder.templates.base.AbstractDatasetTemplate`.
 #. Drop the call to :func:`attach_raw_binary_header` in
    :mod:`mdio.ingestion.metadata`.
 #. Drop ``MDIOSettings.raw_headers``.
@@ -37,10 +38,13 @@
 if TYPE_CHECKING:
     from typing import Any
 
+    from mdio.builder.schemas.v1.dataset import Dataset
+    from mdio.builder.templates.base import AbstractDatasetTemplate
+
 logger = logging.getLogger(__name__)
 
 
-def should_include_raw_headers() -> bool:
+def _should_include_raw_headers() -> bool:
     """True iff the experimental flag is set and the active Zarr format supports it."""
     if not MDIOSettings().raw_headers:
         return False
@@ -51,6 +55,19 @@ def should_include_raw_headers() -> bool:
     return True
 
 
+def maybe_add_raw_headers(template: AbstractDatasetTemplate, dataset: Dataset) -> Dataset:
+    """Append the raw_headers variable to ``dataset`` if the experimental flag enables it.
+
+    The template's internal builder (still alive after ``build_dataset``) is reused
+    so dimension resolution stays consistent. Returns the original dataset untouched
+    when the feature is disabled.
+    """
+    if not _should_include_raw_headers():
+        return dataset
+    template._add_raw_headers()  # noqa: SLF001 - experimental seam, see removal plan above
+    return template._builder.build()  # noqa: SLF001
+
+
 def attach_raw_binary_header(attrs: dict[str, Any], raw_binary_headers: bytes) -> None:
     """Attach a base64-encoded raw binary header to ``attrs`` if the feature is enabled."""
     if not MDIOSettings().raw_headers:

src/mdio/ingestion/dataset_factory.py

@@ -29,7 +29,6 @@ def build(
         schema: ResolvedSchema,
         dimensions: list[Dimension],
         header_dtype: StructuredType | None = None,
-        include_raw_headers: bool = False,
     ) -> Dataset:
         """Build MDIO dataset from schema and dimensions.
 
@@ -38,7 +37,6 @@ def build(
             schema: Resolved schema specifying dataset structure
             dimensions: List of dimension objects with coordinates
             header_dtype: Optional structured type for trace headers
-            include_raw_headers: Whether to include raw binary headers (Zarr v3 only)
 
         Returns:
             Complete Dataset ready for xarray conversion
@@ -51,5 +49,4 @@ def build(
             name=schema.name,
             sizes=sizes,
             header_dtype=header_dtype,
-            include_raw_headers=include_raw_headers,
         )

src/mdio/ingestion/header_analysis.py

@@ -88,8 +88,7 @@ def analyze_streamer_headers(
                 geom_type = StreamerShotGeometryType.A
                 logger.info("Found overlapping channels, assuming streamer type A")
                 overlap_info = (
-                    "Cable %s index %s with channel range %s overlaps cable %s index %s with "
-                    "channel range %s."
+                    "Cable %s index %s with channel range %s overlaps cable %s index %s with channel range %s."
                 )
                 logger.info(overlap_info, cable1, idx1, cable1_range, cable2, idx2, cable2_range)
                 return unique_cables, cable_chan_min, cable_chan_max, geom_type

src/mdio/ingestion/index_strategies.py

@@ -226,9 +226,7 @@ def __init__(self, line_field: str = "sail_line", always_calculate: bool = False
 
     def transform_headers(self, headers: HeaderArray) -> HeaderArray:
         """Create ``shot_index`` for multi-gun acquisition."""
-        unique_lines, unique_guns_per_line, geom_type = analyze_lines_for_guns(
-            headers, line_field=self.line_field
-        )
+        unique_lines, unique_guns_per_line, geom_type = analyze_lines_for_guns(headers, line_field=self.line_field)
 
         logger.info("Ingesting dataset as shot type: %s (line_field=%s)", geom_type.name, self.line_field)
 
@@ -369,9 +367,7 @@ def create_strategy(
 
             if grid_overrides.auto_shot_wrap:
                 line_field, always_calculate = self._resolve_shot_wrap_params(template)
-                strategies.append(
-                    ShotWrappingStrategy(line_field=line_field, always_calculate=always_calculate)
-                )
+                strategies.append(ShotWrappingStrategy(line_field=line_field, always_calculate=always_calculate))
 
             # NonBinned and HasDuplicates are mutually exclusive.
             if grid_overrides.non_binned: