From 2664660fe67526be780424c394ab86c001bb1403 Mon Sep 17 00:00:00 2001 From: "Andrei.Sorokin" Date: Mon, 27 Jul 2026 12:50:53 +0000 Subject: [PATCH 1/5] Expose chunk size and counter dtype on the `HasDuplicates` path The inserted `trace` dimension was hardcoded to chunk size 1 and an int16 counter, which caps a gather at ~32k duplicates and forces one chunk per trace. Continuous receiver gathers need both knobs: ~2.6e5 segments per receiver, and a chunk size large enough to reach a sane object size. `GridOverrides.chunksize` now applies to the `has_duplicates` path as well as `non_binned`, and a new optional `GridOverrides.trace_dtype` selects the counter dtype. Both default to the previous behaviour (chunk 1, int16), so existing payloads that omit them are unaffected. Also generalizes `_resolve_synthesize_dims` to read the template's `synthesize_missing_dims` attribute instead of isinstance-checking the OBN template, so any template can declare dimensions it synthesizes when they are absent from the SEG-Y spec. The ingestion validator keys off the same attribute. Co-authored-by: Cursor --- docs/guides/grid_overrides.md | 33 ++++++++++++++ src/mdio/ingestion/segy/index_strategies.py | 19 ++++++-- src/mdio/ingestion/segy/validation.py | 9 ++-- src/mdio/segy/geometry.py | 39 +++++++++++----- tests/unit/ingestion/test_schema_effects.py | 8 +++- .../ingestion/test_segy_grid_overrides.py | 25 +++++++++-- .../ingestion/test_segy_index_strategies.py | 44 ++++++++++++++++++- tests/unit/test_grid_overrides_pydantic.py | 20 +++++++++ 8 files changed, 170 insertions(+), 27 deletions(-) diff --git a/docs/guides/grid_overrides.md b/docs/guides/grid_overrides.md index 783afbbd..6a6884ca 100644 --- a/docs/guides/grid_overrides.md +++ b/docs/guides/grid_overrides.md @@ -77,6 +77,39 @@ segy_to_mdio( See [OBN Data Import](obn_data_import.md) for a complete guide on importing OBN data. ``` +## HasDuplicates + +Inserts a `trace` dimension to disambiguate traces that share the same index tuple, by +appending a dense per-tuple counter (in trace order) as the `trace` axis before the vertical +axis. Used for continuous receiver gathers, where many segments share the same +`(component, receiver_line, receiver)` index. + +**Supported Templates:** any template (commonly `ObnContinuousReceiverGathers3D`) + +**Parameters:** + +| Parameter | Default | Description | +| ------------- | ------- | ------------------------------------------------------------------------- | +| `chunksize` | `1` | Chunk size for the inserted `trace` dimension. | +| `trace_dtype` | `int16` | NumPy dtype for the `trace` counter (e.g. `"uint32"` for large gathers). | + +Both parameters are optional and backwards compatible: omitting them reproduces the legacy +behavior (chunk size `1`, `int16` counter). + +**Usage:** + +```python +from mdio import GridOverrides + +# ~64 MiB chunks for a whole-trace (15001-sample) continuous receiver gather. +GridOverrides(has_duplicates=True, chunksize=1024, trace_dtype="uint32") +``` + +```{note} +See [CRG Data Import](crg_data_import.md) for a complete guide on importing continuous +receiver gathers. +``` + ## Special Behaviors Some templates have special behaviors that are applied automatically during import, independent of grid overrides. diff --git a/src/mdio/ingestion/segy/index_strategies.py b/src/mdio/ingestion/segy/index_strategies.py index 3efb9e98..9851e0a4 100644 --- a/src/mdio/ingestion/segy/index_strategies.py +++ b/src/mdio/ingestion/segy/index_strategies.py @@ -127,7 +127,10 @@ class DuplicateHandlingStrategy(IndexStrategy): excluded_fields: Additional fields to exclude from grouping. Used by `NonBinnedStrategy` to keep the explicit `non_binned_dims` from polluting the per-tuple counter. - dtype: NumPy dtype for the appended `trace` counter. + dtype: NumPy dtype for the appended `trace` counter. Defaults to `int16`; + widen to e.g. `uint32` for gathers with more than ~32k traces per index tuple. + chunksize: Chunk size assigned to the inserted `trace` dimension by the schema + effect. Defaults to 1 (one trace per chunk), preserving legacy behavior. """ def __init__( @@ -135,10 +138,12 @@ def __init__( coord_fields: Iterable[str] = (), excluded_fields: Iterable[str] = (), dtype: DTypeLike = np.int16, + chunksize: int = 1, ) -> None: self.coord_fields = frozenset(coord_fields) self.excluded_fields = frozenset(excluded_fields) self.dtype = dtype + self._chunksize = chunksize def _dim_fields(self, headers: HeaderArray) -> list[str]: """Header field names that participate in the duplicate grouping.""" @@ -161,8 +166,8 @@ def transform_headers(self, headers: HeaderArray) -> HeaderArray: return rfn.append_fields(headers, "trace", trace_values, usemask=False) def schema_effect(self) -> SchemaEffect: - """Insert a chunksize-1 ``trace`` dimension to disambiguate duplicate index tuples.""" - return InsertTraceDimEffect(chunksize=1) + """Insert a ``trace`` dimension (chunk sized by ``chunksize``) to disambiguate duplicates.""" + return InsertTraceDimEffect(chunksize=self._chunksize) class NonBinnedStrategy(DuplicateHandlingStrategy): @@ -450,7 +455,13 @@ def create_strategy( ) ) elif grid_overrides.has_duplicates: - strategies.append(DuplicateHandlingStrategy(coord_fields=coord_fields)) + strategies.append( + DuplicateHandlingStrategy( + coord_fields=coord_fields, + chunksize=grid_overrides.chunksize or 1, + dtype=grid_overrides.trace_dtype or np.int16, + ) + ) if not strategies: return RegularGridStrategy() diff --git a/src/mdio/ingestion/segy/validation.py b/src/mdio/ingestion/segy/validation.py index 3828fda3..bd7d83fc 100644 --- a/src/mdio/ingestion/segy/validation.py +++ b/src/mdio/ingestion/segy/validation.py @@ -14,17 +14,14 @@ def validate_spec_in_template(segy_spec: SegySpec, mdio_template: AbstractDatasetTemplate) -> None: """Validate that the SegySpec has all required fields in the MDIO template.""" - # Import here to avoid circular imports at module load time - from mdio.builder.templates.seismic_3d_obn import Seismic3DObnReceiverGathersTemplate # noqa: PLC0415 - header_fields = {field.name for field in segy_spec.trace.header.fields} required_fields = set(mdio_template.spatial_dimension_names) | set(mdio_template.coordinate_names) required_fields = required_fields - set(mdio_template.calculated_dimension_names) - # 'component' is optional for OBN (synthesized if missing) - if isinstance(mdio_template, Seismic3DObnReceiverGathersTemplate): - required_fields.discard("component") + # Dimensions the template can synthesize when absent (e.g. 'component' for + # single-component OBN/CRG data) are optional in the SEG-Y spec. + required_fields -= set(getattr(mdio_template, "synthesize_missing_dims", ())) if any(field in SCALE_COORDINATE_KEYS for field in required_fields): required_fields = required_fields | {"coordinate_scalar"} diff --git a/src/mdio/segy/geometry.py b/src/mdio/segy/geometry.py index e1c56372..5e286d71 100644 --- a/src/mdio/segy/geometry.py +++ b/src/mdio/segy/geometry.py @@ -12,9 +12,11 @@ from typing import TYPE_CHECKING from typing import Any +import numpy as np from pydantic import BaseModel from pydantic import ConfigDict from pydantic import Field +from pydantic import field_validator from pydantic import model_validator from mdio.segy.exceptions import GridOverrideMissingParameterError @@ -54,18 +56,38 @@ class GridOverrides(BaseModel): has_duplicates: bool = Field( default=False, alias="HasDuplicates", - description="Add a trace dimension (chunksize 1) to disambiguate duplicate trace indices.", + description="Add a trace dimension to disambiguate duplicate trace indices " + "(chunk size from `chunksize`, defaulting to 1).", ) chunksize: int | None = Field( default=None, gt=0, - description="Chunk size for the trace dimension when `non_binned` is True.", + description="Chunk size for the inserted trace dimension. Required when `non_binned` " + "is True; optional for `has_duplicates` (defaults to 1 when omitted).", + ) + trace_dtype: str | None = Field( + default=None, + description="NumPy dtype for the inserted trace counter on the `has_duplicates` path " + "(e.g. 'uint32' for gathers exceeding int16's ~32k range). Defaults to int16.", ) non_binned_dims: list[str] | None = Field( default=None, description="Dimension names to collapse into the trace dimension when `non_binned` is True.", ) + @field_validator("trace_dtype") + @classmethod + def _check_trace_dtype(cls, value: str | None) -> str | None: + """Reject a `trace_dtype` string that numpy cannot interpret as a dtype.""" + if value is None: + return value + try: + np.dtype(value) + except TypeError as exc: + msg = f"trace_dtype {value!r} is not a valid numpy dtype" + raise ValueError(msg) from exc + return value + @model_validator(mode="after") def _check_non_binned_parameters(self) -> GridOverrides: """Validate parameters when non_binned is True. @@ -108,18 +130,13 @@ def to_legacy_dict(self) -> dict[str, Any]: def _resolve_synthesize_dims(template: AbstractDatasetTemplate | None) -> tuple[str, ...]: """Return dimension fields to synthesize when missing for a given template. - Only the OBN receiver gathers template currently synthesizes ``component``; every - other template returns ``()`` so the strategy registry skips synthesis entirely. + Templates that accept optional dimensions (e.g. a synthesized ``component`` for + single-component OBN/CRG data) declare them via ``synthesize_missing_dims``; every + other template leaves it empty, so the strategy registry skips synthesis entirely. """ if template is None: return () - # Lazy import: builder templates pull in builder schemas that indirectly import this - # module's ``GridOverrides``, so a top-level import would cycle. - from mdio.builder.templates.seismic_3d_obn import Seismic3DObnReceiverGathersTemplate # noqa: PLC0415 - - if isinstance(template, Seismic3DObnReceiverGathersTemplate): - return ("component",) - return () + return tuple(getattr(template, "synthesize_missing_dims", ()) or ()) def validate_overrides_for_template( diff --git a/tests/unit/ingestion/test_schema_effects.py b/tests/unit/ingestion/test_schema_effects.py index 86345adf..33d10052 100644 --- a/tests/unit/ingestion/test_schema_effects.py +++ b/tests/unit/ingestion/test_schema_effects.py @@ -56,11 +56,17 @@ def test_non_binned_wires_chunksize_and_dims(self) -> None: assert effect.collapse_dims == ("channel",) def test_has_duplicates_inserts_trace(self) -> None: - """HasDuplicates yields a 1-wide InsertTraceDimEffect.""" + """HasDuplicates yields a 1-wide InsertTraceDimEffect by default (backwards compatible).""" effect = IndexStrategyRegistry().schema_effect(GridOverrides(has_duplicates=True)) assert isinstance(effect, InsertTraceDimEffect) assert effect.chunksize == 1 + def test_has_duplicates_honours_chunksize(self) -> None: + """HasDuplicates with a chunksize sizes the inserted trace dim accordingly.""" + effect = IndexStrategyRegistry().schema_effect(GridOverrides(has_duplicates=True, chunksize=1024)) + assert isinstance(effect, InsertTraceDimEffect) + assert effect.chunksize == 1024 + class TestInsertTraceDimEffect: """InsertTraceDimEffect adds a calculated trace dim before the vertical axis.""" diff --git a/tests/unit/ingestion/test_segy_grid_overrides.py b/tests/unit/ingestion/test_segy_grid_overrides.py index d17a27df..bdb8063b 100644 --- a/tests/unit/ingestion/test_segy_grid_overrides.py +++ b/tests/unit/ingestion/test_segy_grid_overrides.py @@ -69,12 +69,12 @@ def run_override( new_names = list(index_names) new_chunks = list(chunksize) if chunksize is not None else None - # Both NonBinned and HasDuplicates add a 'trace' dim at index -1; HasDuplicates - # always uses chunksize 1, NonBinned uses the user-supplied value. + # Both NonBinned and HasDuplicates add a 'trace' dim at index -1. NonBinned requires a + # chunksize; HasDuplicates accepts an optional one, defaulting to 1 when omitted. if config.non_binned or config.has_duplicates: new_names.append("trace") if new_chunks is not None: - inserted_chunk = config.chunksize if config.non_binned else 1 + inserted_chunk = config.chunksize or 1 new_chunks.insert(-1, inserted_chunk) return ( @@ -151,6 +151,25 @@ def test_duplicates(self, mock_streamer_headers: dict[str, npt.NDArray]) -> None assert_array_equal(dims[1].coords, CABLES) assert_array_equal(dims[2].coords, RECEIVERS) + def test_duplicates_with_chunksize(self, mock_streamer_headers: dict[str, npt.NDArray]) -> None: + """HasDuplicates can carry an explicit trace chunksize (CRG tuned path).""" + index_names = ("shot_point", "cable") + grid_overrides = {"HasDuplicates": True, "chunksize": 16} + + streamer_headers = mock_streamer_headers[list(index_names)] + chunksize = (4, 4, 8) + + new_headers, new_names, new_chunks = run_override( + grid_overrides, + index_names, + streamer_headers, + chunksize, + ) + + assert new_names == ("shot_point", "cable", "trace") + # The inserted trace chunk honours the override instead of the legacy hardcoded 1. + assert new_chunks == (4, 4, 16, 8) + def test_non_binned(self, mock_streamer_headers: dict[str, npt.NDArray]) -> None: """Test the NonBinned Grid Override command.""" index_names = ("shot_point", "cable") diff --git a/tests/unit/ingestion/test_segy_index_strategies.py b/tests/unit/ingestion/test_segy_index_strategies.py index 1ffd1ec9..9c411ab3 100644 --- a/tests/unit/ingestion/test_segy_index_strategies.py +++ b/tests/unit/ingestion/test_segy_index_strategies.py @@ -102,9 +102,20 @@ def test_non_binned_only(self) -> None: assert strategy.excluded_fields == frozenset({"channel"}) def test_has_duplicates_only(self) -> None: - """``has_duplicates`` -> DuplicateHandlingStrategy.""" + """``has_duplicates`` -> DuplicateHandlingStrategy with legacy chunk-1/int16 defaults.""" strategy = IndexStrategyRegistry().create_strategy(grid_overrides=GridOverrides(has_duplicates=True)) assert isinstance(strategy, DuplicateHandlingStrategy) + # Backwards compatible: no chunksize/trace_dtype -> chunk 1, int16 (unchanged behavior). + assert strategy.dtype == np.int16 + assert strategy.schema_effect().chunksize == 1 + + def test_has_duplicates_wires_chunksize_and_dtype(self) -> None: + """``has_duplicates`` forwards the override's chunksize and trace_dtype to the strategy.""" + overrides = GridOverrides(has_duplicates=True, chunksize=1024, trace_dtype="uint32") + strategy = IndexStrategyRegistry().create_strategy(grid_overrides=overrides) + assert isinstance(strategy, DuplicateHandlingStrategy) + assert strategy.dtype == "uint32" + assert strategy.schema_effect().chunksize == 1024 def test_non_binned_wins_over_has_duplicates(self) -> None: """Both flags set -> NonBinned wins (matches v1.x semantics).""" @@ -166,6 +177,29 @@ def test_template_coord_names_propagate_to_duplicate_strategy(self) -> None: assert strategy.coord_fields == frozenset(template.coordinate_names) +class TestResolveSynthesizeDims: + """`_resolve_synthesize_dims` reads the template's own `synthesize_missing_dims`.""" + + def test_none_template_returns_empty(self) -> None: + """No template -> nothing to synthesize.""" + assert _resolve_synthesize_dims(None) == () + + def test_obn_template_synthesizes_component(self) -> None: + """The OBN template keeps declaring an optional (synthesized) component.""" + template = TemplateRegistry().get("ObnReceiverGathers3D") + assert _resolve_synthesize_dims(template) == ("component",) + + def test_crg_template_synthesizes_component(self) -> None: + """The CRG template opts into component synthesis via the generalized hook.""" + template = TemplateRegistry().get("ObnContinuousReceiverGathers3D") + assert _resolve_synthesize_dims(template) == ("component",) + + def test_plain_template_synthesizes_nothing(self) -> None: + """A template without optional dims yields no synthesis.""" + template = TemplateRegistry().get("PostStack3DTime") + assert _resolve_synthesize_dims(template) == () + + # --------------------------------------------------------------------------- # RegularGridStrategy # --------------------------------------------------------------------------- @@ -252,11 +286,17 @@ def test_excluded_fields_dropped_from_grouping(self) -> None: np.testing.assert_array_equal(out["trace"], [1, 2, 3]) def test_owns_insert_trace_dim_effect(self) -> None: - """The strategy owns its schema reshape: a chunksize-1 inserted ``trace`` dim.""" + """The strategy owns its schema reshape: a chunksize-1 inserted ``trace`` dim by default.""" effect = DuplicateHandlingStrategy().schema_effect() assert isinstance(effect, InsertTraceDimEffect) assert effect.chunksize == 1 + def test_chunksize_flows_into_schema_effect(self) -> None: + """A caller-supplied chunksize sizes the inserted ``trace`` dimension.""" + effect = DuplicateHandlingStrategy(chunksize=1024).schema_effect() + assert isinstance(effect, InsertTraceDimEffect) + assert effect.chunksize == 1024 + # --------------------------------------------------------------------------- # NonBinnedStrategy diff --git a/tests/unit/test_grid_overrides_pydantic.py b/tests/unit/test_grid_overrides_pydantic.py index 4060c94c..15fa4e65 100644 --- a/tests/unit/test_grid_overrides_pydantic.py +++ b/tests/unit/test_grid_overrides_pydantic.py @@ -24,6 +24,7 @@ def test_grid_overrides_defaults() -> None: assert not overrides.non_binned assert not overrides.has_duplicates assert overrides.chunksize is None + assert overrides.trace_dtype is None assert overrides.non_binned_dims is None assert not bool(overrides) @@ -58,6 +59,25 @@ def test_grid_overrides_rejects_unknown_keys() -> None: GridOverrides.model_validate({"FutureFlag": True}) +def test_grid_overrides_trace_dtype_accepts_valid_dtype() -> None: + """A valid numpy dtype string is accepted on the has_duplicates path.""" + overrides = GridOverrides(has_duplicates=True, chunksize=1024, trace_dtype="uint32") + assert overrides.trace_dtype == "uint32" + + +def test_grid_overrides_trace_dtype_rejects_garbage() -> None: + """A string numpy cannot parse as a dtype is rejected at construction.""" + with pytest.raises(ValidationError): + GridOverrides(has_duplicates=True, trace_dtype="not-a-dtype") + + +def test_grid_overrides_trace_dtype_serialization() -> None: + """``trace_dtype`` round-trips through the legacy dict when set, and is omitted by default.""" + overrides = GridOverrides(has_duplicates=True, chunksize=1024, trace_dtype="uint32") + assert overrides.to_legacy_dict() == {"HasDuplicates": True, "chunksize": 1024, "trace_dtype": "uint32"} + assert GridOverrides(has_duplicates=True).to_legacy_dict() == {"HasDuplicates": True} + + def test_grid_overrides_serialization() -> None: """``model_dump`` round-trips both legacy and modern key shapes.""" overrides = GridOverrides(AutoChannelWrap=True, chunksize=64) From b57bdbe467f48978efb933ed64dfc14b71b91090 Mon Sep 17 00:00:00 2001 From: "Andrei.Sorokin" Date: Mon, 27 Jul 2026 12:51:25 +0000 Subject: [PATCH 2/5] Add OBN Continuous Receiver Gathers (CRG) template Continuous receiver gathers record every receiver continuously rather than per shot, so there is no shot index to grid on: the natural key is (component, receiver_line, receiver), and each receiver holds a long run of time segments that share that key. `ObnContinuousReceiverGathers3D` grids on those three spatial dimensions and leaves the segment axis to the `HasDuplicates` override, which inserts `trace` ahead of the vertical axis. The resulting layout is (component, receiver_line, receiver, trace, time). `component` is declared in `synthesize_missing_dims`, so single-component files that carry no component header ingest against the same template with the dimension synthesized. Co-authored-by: Cursor --- docs/guides/crg_data_import.md | 160 ++++++++++++++ docs/guides/index.md | 12 + src/mdio/builder/template_registry.py | 4 + src/mdio/builder/templates/seismic_3d_crg.py | 102 +++++++++ .../test_import_crg_grid_overrides.py | 208 ++++++++++++++++++ .../unit/v1/templates/test_seismic_3d_crg.py | 155 +++++++++++++ .../v1/templates/test_template_registry.py | 14 +- 7 files changed, 650 insertions(+), 5 deletions(-) create mode 100644 docs/guides/crg_data_import.md create mode 100644 src/mdio/builder/templates/seismic_3d_crg.py create mode 100644 tests/integration/test_import_crg_grid_overrides.py create mode 100644 tests/unit/v1/templates/test_seismic_3d_crg.py diff --git a/docs/guides/crg_data_import.md b/docs/guides/crg_data_import.md new file mode 100644 index 00000000..fc7f94a6 --- /dev/null +++ b/docs/guides/crg_data_import.md @@ -0,0 +1,160 @@ +# CRG Data Import + +This guide covers the `ObnContinuousReceiverGathers3D` template for importing OBN +**Continuous Receiver Gathers** (CRG, clock-time) seismic data into MDIO. + +In OBN acquisition each ocean-bottom node records continuously. A continuous-receiver-gather +delivery bundles many receivers and splits each receiver's long recording into fixed-length +(e.g. 30 s) segments, one per SEG-Y trace. The template stores those segments along a dense +per-receiver axis while preserving each segment's absolute recording time in the headers. + +## Template Overview + +The `ObnContinuousReceiverGathers3D` template organizes data with the following dimensions: + +| Dimension | Description | +| -------------- | ---------------------------------------------------------------------------------- | +| `component` | Sensor component (e.g., 1=X, 2=Y, 3=Z, 4=Hydrophone); synthesized when absent | +| `receiver_line`| Receiver line number | +| `receiver` | Receiver station number | +| `trace` | Dense per-receiver segment index, inserted at ingest (see below) | +| `time`/`depth` | Vertical sample axis (whole trace kept in one chunk) | + +### Coordinates + +- **Physical coordinates**: `group_coord_x`, `group_coord_y` (receiver X/Y), indexed by + `(receiver_line, receiver)`. + +```{note} +The `trace` dimension is **not** declared by the template. It is inserted during ingestion by +the [`HasDuplicates`](grid_overrides.md#hasduplicates) grid override, which appends a dense +per-`(component, receiver_line, receiver)` counter in trace order. This reuses MDIO's existing +duplicate-handling machinery instead of introducing a new calculated dimension. +``` + +## Required Grid Overrides + +### HasDuplicates (Required) + +Because many segments share the same `(component, receiver_line, receiver)` index tuple, the +`HasDuplicates` grid override is **required**. It inserts the `trace` segment dimension: + +```python +from mdio import GridOverrides + +# Whole 30 s trace is 15001 samples; a trace chunk of 1024 gives ~64 MiB chunks. +GridOverrides(has_duplicates=True, chunksize=1024) +``` + +- `chunksize` sizes the inserted `trace` dimension. It is optional and defaults to `1` (the + legacy behavior), but a production CRG ingest should set it so chunks are a sensible size. +- `trace_dtype` sets the counter's dtype (default `int16`). Use `"uint32"` for full-campaign + products where a single receiver has more than ~32,000 segments. + +```{note} +Real recording time is **not** encoded into the segment index. The `trace` axis is a dense +positional index in acquisition order; each segment's absolute time is preserved per trace in +the SEG-Y headers (typically an `epoch` field), which rides through to any cut SEG-Y unchanged. +``` + +## Special Behaviors + +### Component Synthesis + +When the SEG-Y spec does not include a `component` field, MDIO synthesizes it with value `1` +for all traces, so one template ingests both single- and multi-component CRG data. This is +driven by the template's `synthesize_missing_dims` and handled by `ComponentSynthesisStrategy`. + +```{note} +A warning is logged when component is synthesized: +> SEG-Y headers do not contain 'component' field required by template; synthesizing dimension +> with constant value 1 for all traces. +``` + +## Usage + +### Basic Import + +```python +from segy.schema import HeaderField +from segy.standards import get_segy_standard + +from mdio import GridOverrides +from mdio import segy_to_mdio +from mdio.builder.template_registry import get_template + +# CRG SEG-Y trace-header mapping (big-endian; bytes past 179 are vendor-defined). +crg_headers = [ + HeaderField(name="channel", byte=13, format="int32"), + HeaderField(name="coordinate_scalar", byte=71, format="int16"), + HeaderField(name="group_coord_x", byte=81, format="int32"), + HeaderField(name="group_coord_y", byte=85, format="int32"), + HeaderField(name="receiver_line", byte=137, format="int16"), + HeaderField(name="receiver", byte=139, format="int16"), + HeaderField(name="epoch", byte=189, format="int64"), # per-trace recording time +] + +crg_spec = get_segy_standard(1.0).customize(trace_header_fields=crg_headers) + +segy_to_mdio( + input_path="crg_data.sgy", + output_path="crg_data.mdio", + segy_spec=crg_spec, + mdio_template=get_template("ObnContinuousReceiverGathers3D"), + grid_overrides=GridOverrides(has_duplicates=True, chunksize=1024), + overwrite=True, +) +``` + +### Single-Component Data + +For data without a `component` header field, simply omit it from the spec — the template +synthesizes `component = 1`: + +```python +# Same as above, but without a component field. +crg_headers = [ + HeaderField(name="channel", byte=13, format="int32"), + HeaderField(name="coordinate_scalar", byte=71, format="int16"), + HeaderField(name="group_coord_x", byte=81, format="int32"), + HeaderField(name="group_coord_y", byte=85, format="int32"), + HeaderField(name="receiver_line", byte=137, format="int16"), + HeaderField(name="receiver", byte=139, format="int16"), + HeaderField(name="epoch", byte=189, format="int64"), + # component omitted - will be synthesized +] +``` + +### Exploring the Data + +```python +from mdio import open_mdio + +ds = open_mdio("crg_data.mdio") + +# View dimensions +print(ds.sizes) +# {'component': 1, 'receiver_line': 1, 'receiver': 10, 'trace': 8172, 'time': 15001} + +# Select a whole receiver gather +receiver_gather = ds.sel(receiver_line=4871, receiver=5908, component=1) +receiver_gather["amplitude"].plot() +``` + +## Required Header Fields + +| Field | Required | Notes | +| ------------------- | -------- | -------------------------------------------- | +| `receiver_line` | Yes | | +| `receiver` | Yes | Receiver station | +| `coordinate_scalar` | Yes | | +| `group_coord_x` | Yes | Receiver X | +| `group_coord_y` | Yes | Receiver Y | +| `component` | No | Synthesized with value 1 if missing | +| `epoch` | No | Recommended; per-trace time kept in headers | + +## See Also + +- [Grid Overrides](grid_overrides.md) - All available grid overrides +- [OBN Data Import](obn_data_import.md) - Shot-indexed OBN receiver gathers +- [Template Registry](../template_registry.md) diff --git a/docs/guides/index.md b/docs/guides/index.md index 08b83b99..3990d67a 100644 --- a/docs/guides/index.md +++ b/docs/guides/index.md @@ -10,6 +10,7 @@ Welcome to the MDIO guides. This section provides in-depth documentation on adva grid_overrides obn_data_import +crg_data_import ``` ## Overview @@ -29,3 +30,14 @@ Ocean Bottom Node (OBN) data has unique characteristics requiring specialized ha - Component synthesis for single-component data See [OBN Data Import](obn_data_import.md) for the complete guide. + +### CRG Data Import + +Continuous Receiver Gathers (CRG) split each ocean-bottom node's continuous recording into +fixed-length segments. The CRG guide covers: + +- The `ObnContinuousReceiverGathers3D` template +- The `HasDuplicates` grid override that builds the per-receiver segment axis +- Component synthesis for single-component data + +See [CRG Data Import](crg_data_import.md) for the complete guide. diff --git a/src/mdio/builder/template_registry.py b/src/mdio/builder/template_registry.py index 970af1a2..8fdcba14 100644 --- a/src/mdio/builder/template_registry.py +++ b/src/mdio/builder/template_registry.py @@ -25,6 +25,7 @@ from mdio.builder.templates.seismic_2d_streamer_shot import Seismic2DStreamerShotGathersTemplate from mdio.builder.templates.seismic_3d_cdp import Seismic3DCdpGathersTemplate from mdio.builder.templates.seismic_3d_coca import Seismic3DCocaGathersTemplate +from mdio.builder.templates.seismic_3d_crg import Seismic3DCrgReceiverGathersTemplate from mdio.builder.templates.seismic_3d_obn import Seismic3DObnReceiverGathersTemplate from mdio.builder.templates.seismic_3d_offset_tiles import Seismic3DOffsetTilesTemplate from mdio.builder.templates.seismic_3d_poststack import Seismic3DPostStackTemplate @@ -152,6 +153,9 @@ def _register_default_templates(self) -> None: # OBN (Ocean Bottom Node) data self.register(Seismic3DObnReceiverGathersTemplate()) + # OBN Continuous Receiver Gathers (clock time) + self.register(Seismic3DCrgReceiverGathersTemplate()) + # Land/OBC shot-receiver data self.register(Seismic3DShotReceiverLineGathersTemplate()) diff --git a/src/mdio/builder/templates/seismic_3d_crg.py b/src/mdio/builder/templates/seismic_3d_crg.py new file mode 100644 index 00000000..a14c9289 --- /dev/null +++ b/src/mdio/builder/templates/seismic_3d_crg.py @@ -0,0 +1,102 @@ +"""Seismic3DCrgReceiverGathersTemplate MDIO v1 dataset template. + +OBN Continuous Receiver Gathers (CRG), clock-time product. Each ocean-bottom node +records continuously; the delivery bundles many receivers, and every receiver's +long recording is segmented into fixed-length (e.g. 30 s) traces. +""" + +from typing import Any + +from mdio.builder.schemas.dtype import ScalarType +from mdio.builder.schemas.v1.variable import CoordinateMetadata +from mdio.builder.templates.base import AbstractDatasetTemplate +from mdio.builder.templates.types import CoordinateSpec +from mdio.builder.templates.types import DimCoordinateTypes +from mdio.builder.templates.types import SeismicDataDomain + + +class Seismic3DCrgReceiverGathersTemplate(AbstractDatasetTemplate): + """Seismic 3D OBN Continuous Receiver Gathers (clock time) template. + + The template declares the spatial receiver geometry ``(component, receiver_line, + receiver)`` plus the vertical axis. The per-receiver segment axis is **not** declared + here: it is inserted at ingest as a ``trace`` dimension by the ``HasDuplicates`` grid + override, which appends a per-``(component, receiver_line, receiver)`` counter in trace + order. This reuses the shipping duplicate-handling machinery instead of adding a new + calculated dimension, so it stays backwards compatible with the existing ``trace`` + handling downstream. Size the inserted ``trace`` chunk (and, for very long campaigns, + its dtype) through ``GridOverrides(has_duplicates=True, chunksize=..., trace_dtype=...)``. + + Real recording time is preserved per trace in ``headers`` (an ``epoch`` field, when the + SEG-Y spec includes it), not baked into the segment index; the segment axis is a dense + positional index in acquisition order. + + Special handling for the component dimension: + If the SEG-Y spec does not contain a ``component`` field, the ingestion process + synthesizes a ``component`` dimension with constant value 1 for all traces (a + warning is logged). This is driven by ``synthesize_missing_dims`` and handled by + ``ComponentSynthesisStrategy``, so one template ingests both single- and + multi-component CRG data. + """ + + def __init__(self, data_domain: SeismicDataDomain = "time"): + super().__init__(data_domain=data_domain) + + # The segment axis (`trace`) is inserted at ingest by HasDuplicates, so it is + # deliberately absent from the declared dimensions here. + self._spatial_dim_names = ("component", "receiver_line", "receiver") + self.synthesize_missing_dims = ("component",) + self._dim_names = (*self._spatial_dim_names, self._data_domain) + self._physical_coord_names = ("group_coord_x", "group_coord_y") + self._logical_coord_names = () + # 3 spatial + vertical. `time` holds the whole trace in one chunk (>= 15001 + # samples); the HasDuplicates effect inserts the segment chunk between them. + self._var_chunk_shape = (1, 1, 1, 16384) + + @property + def _name(self) -> str: + return "ObnContinuousReceiverGathers3D" + + def _load_dataset_attributes(self) -> dict[str, Any]: + return {"surveyType": "3D", "gatherType": "common_receiver"} + + def declare_dim_coordinate_types(self) -> DimCoordinateTypes: + """Declare the data types for each dimension coordinate in this template.""" + return { + "component": ScalarType.UINT8, + "receiver_line": ScalarType.UINT32, + "receiver": ScalarType.UINT32, + self._data_domain: ScalarType.INT32, + } + + def declare_coordinate_specs(self) -> tuple[CoordinateSpec, ...]: + """Declare the receiver-indexed physical coordinates for this template.""" + receiver_dims = ("receiver_line", "receiver") + return ( + CoordinateSpec(name="group_coord_x", dimensions=receiver_dims, dtype=ScalarType.FLOAT64), + CoordinateSpec(name="group_coord_y", dimensions=receiver_dims, dtype=ScalarType.FLOAT64), + ) + + def _add_coordinates(self) -> None: + # Add dimension coordinates. + for name in ("component", "receiver_line", "receiver"): + self._builder.add_coordinate( + name, + dimensions=(name,), + data_type=self._dim_dtype(name), + ) + self._builder.add_coordinate( + self._data_domain, + dimensions=(self._data_domain,), + data_type=self._dim_dtype(self._data_domain), + metadata=CoordinateMetadata(units_v1=self.get_unit_by_key(self._data_domain)), + ) + + # Add non-dimension coordinates: receiver X/Y indexed by (receiver_line, receiver). + for name in ("group_coord_x", "group_coord_y"): + self._builder.add_coordinate( + name, + dimensions=("receiver_line", "receiver"), + data_type=ScalarType.FLOAT64, + metadata=CoordinateMetadata(units_v1=self.get_unit_by_key(name)), + ) diff --git a/tests/integration/test_import_crg_grid_overrides.py b/tests/integration/test_import_crg_grid_overrides.py new file mode 100644 index 00000000..3819c64e --- /dev/null +++ b/tests/integration/test_import_crg_grid_overrides.py @@ -0,0 +1,208 @@ +"""End to end testing for CRG (Continuous Receiver Gathers) SEG-Y to MDIO conversion. + +The CRG template declares ``(component, receiver_line, receiver, time)`` and relies on the +``HasDuplicates`` grid override to insert the per-receiver segment axis as a ``trace`` +dimension at ingest. These tests exercise the full ``segy_to_mdio`` path for both the +single- and multi-component layouts, proving the template + generalized synthesis hook + +the tuned HasDuplicates chunk/dtype knobs work together. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +import dask +import numpy as np +import pytest +import xarray.testing as xrt +from segy.factory import SegyFactory +from segy.schema import HeaderField +from segy.standards import SegyStandard +from segy.standards import get_segy_standard + +from mdio.api.io import open_mdio +from mdio.builder.template_registry import TemplateRegistry +from mdio.converters.segy import segy_to_mdio + +if TYPE_CHECKING: + from pathlib import Path + + from segy.schema import SegySpec + +dask.config.set(scheduler="synchronous") +os.environ["MDIO__IMPORT__SAVE_SEGY_FILE_HEADER"] = "true" + +CRG_TEMPLATE = "ObnContinuousReceiverGathers3D" + + +def get_segy_mock_crg_spec(include_component: bool = False) -> SegySpec: + """Create a mock CRG SEG-Y specification. + + Args: + include_component: Whether to include a component header field. When omitted, the + template synthesizes a constant ``component = 1``. + + Returns: + SegySpec configured for CRG data. + """ + trace_header_fields = [ + HeaderField(name="orig_field_record_num", byte=9, format="int32"), + HeaderField(name="channel", byte=13, format="int32"), + HeaderField(name="samples_per_trace", byte=115, format="int16"), + HeaderField(name="sample_interval", byte=117, format="int16"), + HeaderField(name="coordinate_scalar", byte=71, format="int16"), + HeaderField(name="group_coord_x", byte=81, format="int32"), + HeaderField(name="group_coord_y", byte=85, format="int32"), + HeaderField(name="receiver_line", byte=137, format="int16"), + HeaderField(name="receiver", byte=139, format="int16"), + ] + if include_component: + trace_header_fields.append(HeaderField(name="component", byte=189, format="int16")) + + rev1_spec = get_segy_standard(1.0) + spec = rev1_spec.customize(trace_header_fields=trace_header_fields) + spec.segy_standard = SegyStandard.REV1 + return spec + + +def create_segy_mock_crg( # noqa: PLR0913 + fake_segy_tmp: Path, + num_samples: int, + receiver_line: int, + receivers: list[int], + segments_per_receiver: int, + components: list[int] | None = None, + filename_suffix: str = "", +) -> Path: + """Create a mock CRG SEG-Y file. + + Each ``(component, receiver_line, receiver)`` tuple gets ``segments_per_receiver`` traces + written in acquisition order, so the HasDuplicates counter produces a dense ``trace`` + axis of that length. + """ + include_component = components is not None + segy_path = fake_segy_tmp / f"crg{'_' + filename_suffix if filename_suffix else ''}.sgy" + + component_list = components if include_component else [None] + trace_count = len(component_list) * len(receivers) * segments_per_receiver + + factory = SegyFactory( + spec=get_segy_mock_crg_spec(include_component=include_component), + sample_interval=2000, + samples_per_trace=num_samples, + ) + headers = factory.create_trace_header_template(trace_count) + samples = factory.create_trace_sample_template(trace_count) + + start_x = 700000 + start_y = 4000000 + step = 100 + + trc_idx = 0 + for component in component_list: + for receiver_idx, receiver in enumerate(receivers): + for _segment in range(segments_per_receiver): + headers["orig_field_record_num"][trc_idx] = receiver + headers["channel"][trc_idx] = trc_idx + 1 + headers["receiver_line"][trc_idx] = receiver_line + headers["receiver"][trc_idx] = receiver + if include_component: + headers["component"][trc_idx] = component + + headers["coordinate_scalar"][trc_idx] = -100 + headers["group_coord_x"][trc_idx] = start_x + step * receiver_idx + headers["group_coord_y"][trc_idx] = start_y + step * receiver_idx + + samples[trc_idx] = np.linspace(start=receiver, stop=receiver + 1, num=num_samples) + trc_idx += 1 + + with segy_path.open(mode="wb") as fp: + fp.write(factory.create_textual_header()) + fp.write(factory.create_binary_header()) + fp.write(factory.create_traces(headers, samples)) + + return segy_path + + +@pytest.fixture +def segy_mock_crg_with_component(fake_segy_tmp: Path) -> Path: + """Generate a mock CRG SEG-Y file with a component header (2 components).""" + return create_segy_mock_crg( + fake_segy_tmp, + num_samples=25, + receiver_line=4871, + receivers=[5908, 5909, 5910], + segments_per_receiver=4, + components=[1, 2], + filename_suffix="with_component", + ) + + +@pytest.fixture +def segy_mock_crg_no_component(fake_segy_tmp: Path) -> Path: + """Generate a mock CRG SEG-Y file without a component header (single component).""" + return create_segy_mock_crg( + fake_segy_tmp, + num_samples=25, + receiver_line=4871, + receivers=[5908, 5909, 5910], + segments_per_receiver=4, + components=None, + filename_suffix="no_component", + ) + + +class TestImportCrgWithComponent: + """CRG import with an explicit component header and tuned HasDuplicates knobs.""" + + def test_import_crg_multicomponent(self, segy_mock_crg_with_component: Path, zarr_tmp: Path) -> None: + """Multi-component CRG ingests to the expected 5-D layout with a tuned trace chunk.""" + segy_spec = get_segy_mock_crg_spec(include_component=True) + grid_override = {"HasDuplicates": True, "chunksize": 2, "trace_dtype": "uint32"} + + segy_to_mdio( + segy_spec=segy_spec, + mdio_template=TemplateRegistry().get(CRG_TEMPLATE), + input_path=segy_mock_crg_with_component, + output_path=zarr_tmp, + overwrite=True, + grid_overrides=grid_override, + ) + + ds = open_mdio(zarr_tmp) + + assert ds.attrs["attributes"]["gridOverrides"] == grid_override + # HasDuplicates inserts `trace` between the spatial dims and the vertical axis. + assert ds["amplitude"].dims == ("component", "receiver_line", "receiver", "trace", "time") + assert ds.sizes == {"component": 2, "receiver_line": 1, "receiver": 3, "trace": 4, "time": 25} + + xrt.assert_duckarray_equal(ds["component"], [1, 2]) + xrt.assert_duckarray_equal(ds["receiver_line"], [4871]) + xrt.assert_duckarray_equal(ds["receiver"], [5908, 5909, 5910]) + + +class TestImportCrgSyntheticComponent: + """CRG import without a component header - component is synthesized.""" + + def test_import_crg_synthetic_component(self, segy_mock_crg_no_component: Path, zarr_tmp: Path) -> None: + """Single-component CRG synthesizes component=1 and still builds the trace axis.""" + segy_spec = get_segy_mock_crg_spec(include_component=False) + grid_override = {"HasDuplicates": True} + + segy_to_mdio( + segy_spec=segy_spec, + mdio_template=TemplateRegistry().get(CRG_TEMPLATE), + input_path=segy_mock_crg_no_component, + output_path=zarr_tmp, + overwrite=True, + grid_overrides=grid_override, + ) + + ds = open_mdio(zarr_tmp) + + assert ds["amplitude"].dims == ("component", "receiver_line", "receiver", "trace", "time") + assert ds.sizes == {"component": 1, "receiver_line": 1, "receiver": 3, "trace": 4, "time": 25} + # Component synthesized with the default constant value 1. + xrt.assert_duckarray_equal(ds["component"], [1]) + xrt.assert_duckarray_equal(ds["receiver"], [5908, 5909, 5910]) diff --git a/tests/unit/v1/templates/test_seismic_3d_crg.py b/tests/unit/v1/templates/test_seismic_3d_crg.py new file mode 100644 index 00000000..470295e5 --- /dev/null +++ b/tests/unit/v1/templates/test_seismic_3d_crg.py @@ -0,0 +1,155 @@ +"""Unit tests for Seismic3DCrgReceiverGathersTemplate.""" + +from tests.unit.v1.helpers import validate_variable + +from mdio.builder.schemas.chunk_grid import RegularChunkGrid +from mdio.builder.schemas.compressors import Blosc +from mdio.builder.schemas.compressors import BloscCname +from mdio.builder.schemas.dtype import ScalarType +from mdio.builder.schemas.dtype import StructuredType +from mdio.builder.schemas.v1.units import LengthUnitEnum +from mdio.builder.schemas.v1.units import LengthUnitModel +from mdio.builder.schemas.v1.units import TimeUnitEnum +from mdio.builder.schemas.v1.units import TimeUnitModel +from mdio.builder.templates.seismic_3d_crg import Seismic3DCrgReceiverGathersTemplate + +UNITS_METER = LengthUnitModel(length=LengthUnitEnum.METER) +UNITS_SECOND = TimeUnitModel(time=TimeUnitEnum.SECOND) + +# The per-receiver `trace` segment axis is inserted at ingest by HasDuplicates, so the +# template itself declares only the spatial receiver geometry plus the vertical axis. +DATASET_SIZE_MAP = {"component": 2, "receiver_line": 1, "receiver": 3, "time": 128} +DATASET_DTYPE_MAP = { + "component": "uint8", + "receiver_line": "uint32", + "receiver": "uint32", + "time": "int32", +} +EXPECTED_COORDINATES = ["group_coord_x", "group_coord_y"] +RECEIVER_DIMS = [("receiver_line", 1), ("receiver", 3)] + + +class TestSeismic3DCrgReceiverGathersTemplate: + """Unit tests for Seismic3DCrgReceiverGathersTemplate.""" + + def test_configuration(self) -> None: + """Test template configuration and attributes.""" + t = Seismic3DCrgReceiverGathersTemplate(data_domain="time") + + assert t.name == "ObnContinuousReceiverGathers3D" + assert t._dim_names == ("component", "receiver_line", "receiver", "time") + assert t.spatial_dimension_names == ("component", "receiver_line", "receiver") + # `trace` is inserted by HasDuplicates at ingest, not declared as a calculated dim. + assert t._calculated_dims == () + assert t.synthesize_missing_dims == ("component",) + assert t._physical_coord_names == ("group_coord_x", "group_coord_y") + assert t._logical_coord_names == () + assert t._var_chunk_shape == (1, 1, 1, 16384) + + assert t._builder is None + assert t._dim_sizes == () + + attrs = t._load_dataset_attributes() + assert attrs == {"surveyType": "3D", "gatherType": "common_receiver"} + assert t.default_variable_name == "amplitude" + + def test_whole_trace_chunk(self) -> None: + """The vertical axis is chunked whole-trace (>= a real 15001-sample record).""" + t = Seismic3DCrgReceiverGathersTemplate(data_domain="time") + assert t.full_chunk_shape[-1] == 16384 + assert t.full_chunk_shape[-1] >= 15001 + + def test_build_dataset(self, structured_headers: StructuredType) -> None: + """Test building a complete dataset with the template.""" + t = Seismic3DCrgReceiverGathersTemplate(data_domain="time") + t.add_units({"group_coord_x": UNITS_METER, "group_coord_y": UNITS_METER}) + t.add_units({"time": UNITS_SECOND}) + + sizes = tuple(DATASET_SIZE_MAP.values()) + dataset = t.build_dataset("CrgSurvey3D", sizes=sizes, header_dtype=structured_headers) + + assert dataset.metadata.name == "CrgSurvey3D" + assert dataset.metadata.attributes["gatherType"] == "common_receiver" + assert dataset.metadata.attributes["defaultVariableName"] == "amplitude" + + # 4 dim coords + 2 non-dim coords + amplitude + trace_mask + headers = 9 variables. + assert len(dataset.variables) == 9 + + # Dimension coordinate variables. + for dim_name, dim_size in DATASET_SIZE_MAP.items(): + validate_variable( + dataset, + name=dim_name, + dims=[(dim_name, dim_size)], + coords=[dim_name], + dtype=ScalarType(DATASET_DTYPE_MAP[dim_name]), + ) + + # Receiver coordinate variables (indexed by receiver_line + receiver). + for coord_name in EXPECTED_COORDINATES: + coord = validate_variable( + dataset, + name=coord_name, + dims=RECEIVER_DIMS, + coords=[coord_name], + dtype=ScalarType.FLOAT64, + ) + assert coord.metadata.units_v1.length == LengthUnitEnum.METER + + # Headers and trace mask span the spatial dims only. + validate_variable( + dataset, + name="headers", + dims=[(k, v) for k, v in DATASET_SIZE_MAP.items() if k != "time"], + coords=EXPECTED_COORDINATES, + dtype=structured_headers, + ) + validate_variable( + dataset, + name="trace_mask", + dims=[(k, v) for k, v in DATASET_SIZE_MAP.items() if k != "time"], + coords=EXPECTED_COORDINATES, + dtype=ScalarType.BOOL, + ) + + # Seismic amplitude variable spans all declared dims and keeps the whole-trace chunk. + seismic = validate_variable( + dataset, + name="amplitude", + dims=list(DATASET_SIZE_MAP.items()), + coords=EXPECTED_COORDINATES, + dtype=ScalarType.FLOAT32, + ) + assert isinstance(seismic.compressor, Blosc) + assert seismic.compressor.cname == BloscCname.zstd + assert isinstance(seismic.metadata.chunk_grid, RegularChunkGrid) + assert seismic.metadata.chunk_grid.configuration.chunk_shape == (1, 1, 1, 16384) + + def test_depth_domain(self, structured_headers: StructuredType) -> None: + """Test building a dataset with depth domain.""" + t = Seismic3DCrgReceiverGathersTemplate(data_domain="depth") + + assert t.trace_domain == "depth" + assert t._dim_names == ("component", "receiver_line", "receiver", "depth") + + sizes = (1, 1, 2, 64) + dataset = t.build_dataset("CrgSurveyDepth", sizes=sizes, header_dtype=structured_headers) + + depth_coord = next((v for v in dataset.variables if v.name == "depth"), None) + assert depth_coord is not None + assert depth_coord.dimensions[0].name == "depth" + assert depth_coord.dimensions[0].size == 64 + + +def _assert_coordinate_specs_match_build(template: Seismic3DCrgReceiverGathersTemplate) -> None: + """declare_coordinate_specs must stay in sync with the built coordinates (base guard).""" + specs = {spec.name: spec for spec in template.declare_coordinate_specs()} + assert set(specs) == {"group_coord_x", "group_coord_y"} + for spec in specs.values(): + assert spec.dimensions == ("receiver_line", "receiver") + assert spec.dtype == ScalarType.FLOAT64 + + +def test_declare_coordinate_specs() -> None: + """The declared coordinate specs describe receiver X/Y over (receiver_line, receiver).""" + _assert_coordinate_specs_match_build(Seismic3DCrgReceiverGathersTemplate()) diff --git a/tests/unit/v1/templates/test_template_registry.py b/tests/unit/v1/templates/test_template_registry.py index 37d63ec2..2576559e 100644 --- a/tests/unit/v1/templates/test_template_registry.py +++ b/tests/unit/v1/templates/test_template_registry.py @@ -38,9 +38,13 @@ "StreamerShotGathers3D", "StreamerFieldRecords3D", "ObnReceiverGathers3D", + "ObnContinuousReceiverGathers3D", "ShotReceiverLineGathers3D", ] +# Number of built-in templates registered by `_register_default_templates`. +NUM_DEFAULT_TEMPLATES = len(EXPECTED_DEFAULT_TEMPLATE_NAMES) + class MockDatasetTemplate(AbstractDatasetTemplate): """Mock template for testing.""" @@ -245,7 +249,7 @@ def test_list_all_templates(self) -> None: registry.register(template2) templates = registry.list_all_templates() - assert len(templates) == 22 + 2 # 22 default + 2 custom + assert len(templates) == NUM_DEFAULT_TEMPLATES + 2 # default + 2 custom assert "Template_One" in templates assert "Template_Two" in templates @@ -255,7 +259,7 @@ def test_clear_templates(self) -> None: # Default templates are always installed templates = list_templates() - assert len(templates) == 22 + assert len(templates) == NUM_DEFAULT_TEMPLATES # Add some templates template1 = MockDatasetTemplate("Template1") @@ -264,7 +268,7 @@ def test_clear_templates(self) -> None: registry.register(template1) registry.register(template2) - assert len(registry.list_all_templates()) == 22 + 2 # 22 default + 2 custom + assert len(registry.list_all_templates()) == NUM_DEFAULT_TEMPLATES + 2 # default + 2 custom # Clear all registry.clear() @@ -397,7 +401,7 @@ def test_list_templates_global(self) -> None: register_template(template2) templates = list_templates() - assert len(templates) == 24 # 22 default + 2 custom + assert len(templates) == NUM_DEFAULT_TEMPLATES + 2 # default + 2 custom assert "template1" in templates assert "template2" in templates @@ -440,7 +444,7 @@ def register_template_worker(template_id: int) -> None: assert len(errors) == 0 assert len(results) == 10 # Including default templates - assert len(registry.list_all_templates()) == 32 # 22 default + 10 registered + assert len(registry.list_all_templates()) == NUM_DEFAULT_TEMPLATES + 10 # default + 10 registered # Check all templates are registered for i in range(10): From 81506ad04ec5753cf6b9178f133f9e367841f3de Mon Sep 17 00:00:00 2001 From: "Andrei.Sorokin" Date: Mon, 27 Jul 2026 13:03:43 +0000 Subject: [PATCH 3/5] Apply prettier formatting to CRG docs tables Co-authored-by: Cursor --- docs/guides/crg_data_import.md | 32 ++++++++++++++++---------------- docs/guides/grid_overrides.md | 8 ++++---- 2 files changed, 20 insertions(+), 20 deletions(-) diff --git a/docs/guides/crg_data_import.md b/docs/guides/crg_data_import.md index fc7f94a6..47c731f1 100644 --- a/docs/guides/crg_data_import.md +++ b/docs/guides/crg_data_import.md @@ -12,13 +12,13 @@ per-receiver axis while preserving each segment's absolute recording time in the The `ObnContinuousReceiverGathers3D` template organizes data with the following dimensions: -| Dimension | Description | -| -------------- | ---------------------------------------------------------------------------------- | -| `component` | Sensor component (e.g., 1=X, 2=Y, 3=Z, 4=Hydrophone); synthesized when absent | -| `receiver_line`| Receiver line number | -| `receiver` | Receiver station number | -| `trace` | Dense per-receiver segment index, inserted at ingest (see below) | -| `time`/`depth` | Vertical sample axis (whole trace kept in one chunk) | +| Dimension | Description | +| --------------- | ----------------------------------------------------------------------------- | +| `component` | Sensor component (e.g., 1=X, 2=Y, 3=Z, 4=Hydrophone); synthesized when absent | +| `receiver_line` | Receiver line number | +| `receiver` | Receiver station number | +| `trace` | Dense per-receiver segment index, inserted at ingest (see below) | +| `time`/`depth` | Vertical sample axis (whole trace kept in one chunk) | ### Coordinates @@ -143,15 +143,15 @@ receiver_gather["amplitude"].plot() ## Required Header Fields -| Field | Required | Notes | -| ------------------- | -------- | -------------------------------------------- | -| `receiver_line` | Yes | | -| `receiver` | Yes | Receiver station | -| `coordinate_scalar` | Yes | | -| `group_coord_x` | Yes | Receiver X | -| `group_coord_y` | Yes | Receiver Y | -| `component` | No | Synthesized with value 1 if missing | -| `epoch` | No | Recommended; per-trace time kept in headers | +| Field | Required | Notes | +| ------------------- | -------- | ------------------------------------------- | +| `receiver_line` | Yes | | +| `receiver` | Yes | Receiver station | +| `coordinate_scalar` | Yes | | +| `group_coord_x` | Yes | Receiver X | +| `group_coord_y` | Yes | Receiver Y | +| `component` | No | Synthesized with value 1 if missing | +| `epoch` | No | Recommended; per-trace time kept in headers | ## See Also diff --git a/docs/guides/grid_overrides.md b/docs/guides/grid_overrides.md index 6a6884ca..7fdfddf4 100644 --- a/docs/guides/grid_overrides.md +++ b/docs/guides/grid_overrides.md @@ -88,10 +88,10 @@ axis. Used for continuous receiver gathers, where many segments share the same **Parameters:** -| Parameter | Default | Description | -| ------------- | ------- | ------------------------------------------------------------------------- | -| `chunksize` | `1` | Chunk size for the inserted `trace` dimension. | -| `trace_dtype` | `int16` | NumPy dtype for the `trace` counter (e.g. `"uint32"` for large gathers). | +| Parameter | Default | Description | +| ------------- | ------- | ------------------------------------------------------------------------ | +| `chunksize` | `1` | Chunk size for the inserted `trace` dimension. | +| `trace_dtype` | `int16` | NumPy dtype for the `trace` counter (e.g. `"uint32"` for large gathers). | Both parameters are optional and backwards compatible: omitting them reproduces the legacy behavior (chunk size `1`, `int16` counter). From 95bd085342a8be2ed667ccfa85f5993d6a671eb1 Mon Sep 17 00:00:00 2001 From: "Andrei.Sorokin" Date: Tue, 28 Jul 2026 08:54:39 +0000 Subject: [PATCH 4/5] Make `trace_dtype` own the whole inserted trace axis The knob only reached the in-flight counter: the stored `trace` coordinate kept the DimensionSpec default, so a caller asking for uint32 still got an int32 axis. Both trace-inserting schema effects now take the dtype and pass it to DimensionSpec explicitly, and the strategy forwards its own. The dtype is also forwarded on the `non_binned` path, which inserts the same `trace` axis and had the same int16 ceiling while silently ignoring the knob. Typing the field as ScalarType and restricting it to integers rejects dtypes that cannot count traces (float32, bool, object, bytes, datetime64) at construction instead of failing deep inside numpy. `None` still means the legacy pair - an int16 counter stored as an int32 coordinate - now named as LEGACY_COUNTER_DTYPE / LEGACY_TRACE_DTYPE and pinned by tests, including one asserting the counter raises OverflowError rather than wrapping when a gather outgrows it. `synthesize_missing_dims` moves to a class attribute on AbstractDatasetTemplate so consumers can read the hook off any template without a getattr fallback; it was previously set in __init__, which hid it from `MagicMock(spec=...)`. Co-authored-by: Cursor --- docs/guides/crg_data_import.md | 8 +-- docs/guides/grid_overrides.md | 13 +++-- src/mdio/builder/templates/base.py | 6 ++- src/mdio/ingestion/segy/index_strategies.py | 34 ++++++++---- src/mdio/ingestion/segy/schema_effects.py | 22 ++++++-- src/mdio/ingestion/segy/validation.py | 2 +- src/mdio/segy/geometry.py | 51 +++++++++++++----- .../test_import_crg_grid_overrides.py | 6 +++ tests/unit/ingestion/test_schema_effects.py | 20 ++++++- .../ingestion/test_segy_index_strategies.py | 52 ++++++++++++++++--- tests/unit/test_grid_overrides_pydantic.py | 12 +++-- 11 files changed, 174 insertions(+), 52 deletions(-) diff --git a/docs/guides/crg_data_import.md b/docs/guides/crg_data_import.md index 47c731f1..b06419d8 100644 --- a/docs/guides/crg_data_import.md +++ b/docs/guides/crg_data_import.md @@ -48,8 +48,10 @@ GridOverrides(has_duplicates=True, chunksize=1024) - `chunksize` sizes the inserted `trace` dimension. It is optional and defaults to `1` (the legacy behavior), but a production CRG ingest should set it so chunks are a sensible size. -- `trace_dtype` sets the counter's dtype (default `int16`). Use `"uint32"` for full-campaign - products where a single receiver has more than ~32,000 segments. +- `trace_dtype` sets the dtype of both the segment counter and the stored `trace` coordinate. + Omitting it keeps the legacy `int16` counter stored as `int32`, which tops out at 32,767 + segments per receiver and raises `OverflowError` past that. Use `"uint32"` for + full-campaign products, where a single receiver holds far more segments. ```{note} Real recording time is **not** encoded into the segment index. The `trace` axis is a dense @@ -101,7 +103,7 @@ segy_to_mdio( output_path="crg_data.mdio", segy_spec=crg_spec, mdio_template=get_template("ObnContinuousReceiverGathers3D"), - grid_overrides=GridOverrides(has_duplicates=True, chunksize=1024), + grid_overrides=GridOverrides(has_duplicates=True, chunksize=1024, trace_dtype="uint32"), overwrite=True, ) ``` diff --git a/docs/guides/grid_overrides.md b/docs/guides/grid_overrides.md index 7fdfddf4..d49950c9 100644 --- a/docs/guides/grid_overrides.md +++ b/docs/guides/grid_overrides.md @@ -88,13 +88,16 @@ axis. Used for continuous receiver gathers, where many segments share the same **Parameters:** -| Parameter | Default | Description | -| ------------- | ------- | ------------------------------------------------------------------------ | -| `chunksize` | `1` | Chunk size for the inserted `trace` dimension. | -| `trace_dtype` | `int16` | NumPy dtype for the `trace` counter (e.g. `"uint32"` for large gathers). | +| Parameter | Default | Description | +| ------------- | ------- | ------------------------------------------------------------------------- | +| `chunksize` | `1` | Chunk size for the inserted `trace` dimension. | +| `trace_dtype` | - | Integer dtype for the `trace` counter and the coordinate it is stored as. | Both parameters are optional and backwards compatible: omitting them reproduces the legacy -behavior (chunk size `1`, `int16` counter). +behavior (chunk size `1`, an `int16` counter stored as an `int32` coordinate). An `int16` +counter tops out at 32,767 traces per index tuple and raises `OverflowError` beyond that, so +set `trace_dtype` for larger gathers. `trace_dtype` applies to `NonBinned` as well, since it +inserts the same `trace` axis. **Usage:** diff --git a/src/mdio/builder/templates/base.py b/src/mdio/builder/templates/base.py index 48d38296..610819ad 100644 --- a/src/mdio/builder/templates/base.py +++ b/src/mdio/builder/templates/base.py @@ -35,6 +35,11 @@ class AbstractDatasetTemplate(ABC): to override specific steps. """ + # Dimensions ingestion may synthesize when the SEG-Y spec does not carry them (e.g. + # `component` for single-component data). Declared on the class, not in `__init__`, so + # ingestion can read the hook off any template; subclasses opt in by setting it. + synthesize_missing_dims: tuple[str, ...] = () + def __init__(self, data_domain: SeismicDataDomain) -> None: self._data_domain = data_domain.lower() @@ -47,7 +52,6 @@ def __init__(self, data_domain: SeismicDataDomain) -> None: self._physical_coord_names: tuple[str, ...] = () self._logical_coord_names: tuple[str, ...] = () self._var_chunk_shape: tuple[int, ...] = () - self.synthesize_missing_dims: tuple[str, ...] = () self._builder: MDIODatasetBuilder | None = None self._dim_sizes: tuple[int, ...] = () diff --git a/src/mdio/ingestion/segy/index_strategies.py b/src/mdio/ingestion/segy/index_strategies.py index 9851e0a4..163fcfaa 100644 --- a/src/mdio/ingestion/segy/index_strategies.py +++ b/src/mdio/ingestion/segy/index_strategies.py @@ -17,6 +17,7 @@ import numpy as np from numpy.lib import recfunctions as rfn +from mdio.builder.schemas.dtype import ScalarType from mdio.core import Dimension from mdio.ingestion.segy.header_analysis import ShotGunGeometryType from mdio.ingestion.segy.header_analysis import StreamerShotGeometryType @@ -30,7 +31,6 @@ if TYPE_CHECKING: from collections.abc import Iterable - from numpy.typing import DTypeLike from segy.arrays import HeaderArray from mdio.builder.templates.base import AbstractDatasetTemplate @@ -39,6 +39,10 @@ logger = logging.getLogger(__name__) +# Dtype of the appended `trace` counter when the caller does not pick one. Counters were +# int16 before the dtype was configurable, so that is what an unspecified dtype keeps. +LEGACY_COUNTER_DTYPE = ScalarType.INT16 + class IndexStrategy(ABC): """Abstract base for header indexing strategies. @@ -127,8 +131,12 @@ class DuplicateHandlingStrategy(IndexStrategy): excluded_fields: Additional fields to exclude from grouping. Used by `NonBinnedStrategy` to keep the explicit `non_binned_dims` from polluting the per-tuple counter. - dtype: NumPy dtype for the appended `trace` counter. Defaults to `int16`; - widen to e.g. `uint32` for gathers with more than ~32k traces per index tuple. + dtype: Integer dtype of the inserted `trace` dimension, used both for the counter + appended here and for the coordinate the schema effect stores. `None` keeps the + legacy pair -- an int16 counter stored as an int32 coordinate -- so ingests + predating this knob are unchanged. Widen it (e.g. `ScalarType.UINT32`) when + traces per index tuple outgrow the counter dtype, which raises `OverflowError` + rather than wrapping. chunksize: Chunk size assigned to the inserted `trace` dimension by the schema effect. Defaults to 1 (one trace per chunk), preserving legacy behavior. """ @@ -137,7 +145,7 @@ def __init__( self, coord_fields: Iterable[str] = (), excluded_fields: Iterable[str] = (), - dtype: DTypeLike = np.int16, + dtype: ScalarType | None = None, chunksize: int = 1, ) -> None: self.coord_fields = frozenset(coord_fields) @@ -157,7 +165,7 @@ def transform_headers(self, headers: HeaderArray) -> HeaderArray: """Append a per-dimension-tuple `trace` counter to headers.""" dim_fields = self._dim_fields(headers) dim_headers = headers[dim_fields] if dim_fields else headers - with_trace = analyze_non_indexed_headers(dim_headers, dtype=self.dtype) + with_trace = analyze_non_indexed_headers(dim_headers, dtype=self.dtype or LEGACY_COUNTER_DTYPE) if with_trace is None or "trace" not in with_trace.dtype.names: return headers @@ -167,7 +175,7 @@ def transform_headers(self, headers: HeaderArray) -> HeaderArray: def schema_effect(self) -> SchemaEffect: """Insert a ``trace`` dimension (chunk sized by ``chunksize``) to disambiguate duplicates.""" - return InsertTraceDimEffect(chunksize=self._chunksize) + return InsertTraceDimEffect(chunksize=self._chunksize, dtype=self.dtype) class NonBinnedStrategy(DuplicateHandlingStrategy): @@ -184,7 +192,8 @@ class NonBinnedStrategy(DuplicateHandlingStrategy): non_binned_dims: Header fields collapsed into `trace`. They are excluded from the duplicate grouping so the counter only varies along the remaining dims. coord_fields: Template coordinate names to exclude from grouping. - dtype: NumPy dtype for the appended `trace` counter. + dtype: Integer dtype of the inserted `trace` dimension; see + `DuplicateHandlingStrategy`. """ def __init__( @@ -192,7 +201,7 @@ def __init__( chunksize: int, non_binned_dims: Iterable[str], coord_fields: Iterable[str] = (), - dtype: DTypeLike = np.int16, + dtype: ScalarType | None = None, ) -> None: collapse_dims = tuple(non_binned_dims) super().__init__( @@ -205,7 +214,11 @@ def __init__( def schema_effect(self) -> SchemaEffect: """Collapse the non-binned dims into a ``trace`` dimension sized by ``chunksize``.""" - return CollapseToTraceEffect(chunksize=self._chunksize, collapse_dims=self._collapse_dims) + return CollapseToTraceEffect( + chunksize=self._chunksize, + collapse_dims=self._collapse_dims, + dtype=self.dtype, + ) class ChannelWrappingStrategy(IndexStrategy): @@ -452,6 +465,7 @@ def create_strategy( chunksize=grid_overrides.chunksize, non_binned_dims=grid_overrides.non_binned_dims or (), coord_fields=coord_fields, + dtype=grid_overrides.trace_dtype, ) ) elif grid_overrides.has_duplicates: @@ -459,7 +473,7 @@ def create_strategy( DuplicateHandlingStrategy( coord_fields=coord_fields, chunksize=grid_overrides.chunksize or 1, - dtype=grid_overrides.trace_dtype or np.int16, + dtype=grid_overrides.trace_dtype, ) ) diff --git a/src/mdio/ingestion/segy/schema_effects.py b/src/mdio/ingestion/segy/schema_effects.py index 1c753891..fc2d77e6 100644 --- a/src/mdio/ingestion/segy/schema_effects.py +++ b/src/mdio/ingestion/segy/schema_effects.py @@ -17,16 +17,23 @@ _TRACE_DIM = "trace" +# Dtype of the inserted `trace` coordinate when the caller does not pick one. Datasets +# written before the dtype was configurable stored it as int32, so that is what an +# unspecified dtype keeps producing. +LEGACY_TRACE_DTYPE = ScalarType.INT32 + class InsertTraceDimEffect(SchemaEffect): """Insert a calculated trace dimension before the vertical axis. Args: chunksize: Chunk size for the trace dimension. + dtype: Dtype of the stored trace coordinate. Defaults to `LEGACY_TRACE_DTYPE`. """ - def __init__(self, chunksize: int = 1) -> None: + def __init__(self, chunksize: int = 1, dtype: ScalarType | None = None) -> None: self.chunksize = chunksize + self.dtype = dtype or LEGACY_TRACE_DTYPE def apply(self, schema: ResolvedSchema) -> ResolvedSchema: """Insert the trace dimension and its chunk before the vertical dimension.""" @@ -35,7 +42,7 @@ def apply(self, schema: ResolvedSchema) -> ResolvedSchema: chunk for dim, chunk in zip(schema.dimensions, schema.chunk_shape, strict=True) if dim.is_spatial ] - trace_dim = DimensionSpec(name=_TRACE_DIM, is_spatial=True, is_calculated=True) + trace_dim = DimensionSpec(name=_TRACE_DIM, is_spatial=True, is_calculated=True, dtype=self.dtype) new_dimensions = [*spatial_dims, trace_dim] new_chunk_shape = [*spatial_chunks, self.chunksize] @@ -54,11 +61,18 @@ class CollapseToTraceEffect(SchemaEffect): chunksize: Chunk size for the trace dimension. collapse_dims: Names of spatial dimensions to collapse. If None, collapses all spatial dimensions except the first. + dtype: Dtype of the stored trace coordinate. Defaults to `LEGACY_TRACE_DTYPE`. """ - def __init__(self, chunksize: int | None, collapse_dims: tuple[str, ...] | None = None) -> None: + def __init__( + self, + chunksize: int | None, + collapse_dims: tuple[str, ...] | None = None, + dtype: ScalarType | None = None, + ) -> None: self.chunksize = chunksize self.collapse_dims = collapse_dims + self.dtype = dtype or LEGACY_TRACE_DTYPE def _resolve_collapse_dims(self, schema: ResolvedSchema) -> tuple[str, ...]: """Resolve the spatial dimensions to collapse.""" @@ -84,7 +98,7 @@ def apply(self, schema: ResolvedSchema) -> ResolvedSchema: new_chunk_shape = [chunk for _, chunk in spatial_dims] if replaced_count > 0: - new_dimensions.append(DimensionSpec(name=_TRACE_DIM, is_spatial=True, is_calculated=True)) + new_dimensions.append(DimensionSpec(name=_TRACE_DIM, is_spatial=True, is_calculated=True, dtype=self.dtype)) new_chunk_shape.append(self.chunksize) for dim, chunk in zip(schema.dimensions, schema.chunk_shape, strict=True): diff --git a/src/mdio/ingestion/segy/validation.py b/src/mdio/ingestion/segy/validation.py index bd7d83fc..3dec2819 100644 --- a/src/mdio/ingestion/segy/validation.py +++ b/src/mdio/ingestion/segy/validation.py @@ -21,7 +21,7 @@ def validate_spec_in_template(segy_spec: SegySpec, mdio_template: AbstractDatase # Dimensions the template can synthesize when absent (e.g. 'component' for # single-component OBN/CRG data) are optional in the SEG-Y spec. - required_fields -= set(getattr(mdio_template, "synthesize_missing_dims", ())) + required_fields -= set(mdio_template.synthesize_missing_dims) if any(field in SCALE_COORDINATE_KEYS for field in required_fields): required_fields = required_fields | {"coordinate_scalar"} diff --git a/src/mdio/segy/geometry.py b/src/mdio/segy/geometry.py index 5e286d71..34ce2cb7 100644 --- a/src/mdio/segy/geometry.py +++ b/src/mdio/segy/geometry.py @@ -12,13 +12,13 @@ from typing import TYPE_CHECKING from typing import Any -import numpy as np from pydantic import BaseModel from pydantic import ConfigDict from pydantic import Field from pydantic import field_validator from pydantic import model_validator +from mdio.builder.schemas.dtype import ScalarType from mdio.segy.exceptions import GridOverrideMissingParameterError if TYPE_CHECKING: @@ -27,6 +27,20 @@ logger = logging.getLogger(__name__) +# The inserted `trace` axis counts traces, so only integer types can carry it. +TRACE_COUNTER_DTYPES = frozenset( + { + ScalarType.INT8, + ScalarType.INT16, + ScalarType.INT32, + ScalarType.INT64, + ScalarType.UINT8, + ScalarType.UINT16, + ScalarType.UINT32, + ScalarType.UINT64, + } +) + class GridOverrides(BaseModel): """Type-safe configuration for grid override operations during SEG-Y ingestion.""" @@ -65,10 +79,12 @@ class GridOverrides(BaseModel): description="Chunk size for the inserted trace dimension. Required when `non_binned` " "is True; optional for `has_duplicates` (defaults to 1 when omitted).", ) - trace_dtype: str | None = Field( + trace_dtype: ScalarType | None = Field( default=None, - description="NumPy dtype for the inserted trace counter on the `has_duplicates` path " - "(e.g. 'uint32' for gathers exceeding int16's ~32k range). Defaults to int16.", + description="Integer dtype of the inserted trace dimension, applied to both the " + "counter and its stored coordinate. Omit it to keep the legacy pair (an int16 " + "counter stored as int32); set e.g. 'uint32' when one index tuple holds more " + "traces than an int16 counter can reach (~32k).", ) non_binned_dims: list[str] | None = Field( default=None, @@ -77,15 +93,22 @@ class GridOverrides(BaseModel): @field_validator("trace_dtype") @classmethod - def _check_trace_dtype(cls, value: str | None) -> str | None: - """Reject a `trace_dtype` string that numpy cannot interpret as a dtype.""" - if value is None: - return value - try: - np.dtype(value) - except TypeError as exc: - msg = f"trace_dtype {value!r} is not a valid numpy dtype" - raise ValueError(msg) from exc + def _check_trace_dtype(cls, value: ScalarType | None) -> ScalarType | None: + """Reject a `trace_dtype` that cannot hold a trace counter. + + Args: + value: Requested trace dtype, or None for the legacy int16/int32 pair. + + Returns: + The validated dtype. + + Raises: + ValueError: If the dtype is not a signed or unsigned integer. + """ + if value is not None and value not in TRACE_COUNTER_DTYPES: + allowed = ", ".join(sorted(TRACE_COUNTER_DTYPES)) + msg = f"trace_dtype must be an integer type ({allowed}), got {value.value!r}" + raise ValueError(msg) return value @model_validator(mode="after") @@ -136,7 +159,7 @@ def _resolve_synthesize_dims(template: AbstractDatasetTemplate | None) -> tuple[ """ if template is None: return () - return tuple(getattr(template, "synthesize_missing_dims", ()) or ()) + return template.synthesize_missing_dims def validate_overrides_for_template( diff --git a/tests/integration/test_import_crg_grid_overrides.py b/tests/integration/test_import_crg_grid_overrides.py index 3819c64e..cc6f01f5 100644 --- a/tests/integration/test_import_crg_grid_overrides.py +++ b/tests/integration/test_import_crg_grid_overrides.py @@ -176,6 +176,9 @@ def test_import_crg_multicomponent(self, segy_mock_crg_with_component: Path, zar # HasDuplicates inserts `trace` between the spatial dims and the vertical axis. assert ds["amplitude"].dims == ("component", "receiver_line", "receiver", "trace", "time") assert ds.sizes == {"component": 2, "receiver_line": 1, "receiver": 3, "trace": 4, "time": 25} + # `trace_dtype` reaches the stored segment axis, not just the in-flight counter. + assert ds["trace"].dtype == np.uint32 + assert ds["amplitude"].encoding["chunks"] == (1, 1, 1, 2, 16384) xrt.assert_duckarray_equal(ds["component"], [1, 2]) xrt.assert_duckarray_equal(ds["receiver_line"], [4871]) @@ -203,6 +206,9 @@ def test_import_crg_synthetic_component(self, segy_mock_crg_no_component: Path, assert ds["amplitude"].dims == ("component", "receiver_line", "receiver", "trace", "time") assert ds.sizes == {"component": 1, "receiver_line": 1, "receiver": 3, "trace": 4, "time": 25} + # No trace_dtype requested -> the legacy int32 segment axis, chunked one trace deep. + assert ds["trace"].dtype == np.int32 + assert ds["amplitude"].encoding["chunks"] == (1, 1, 1, 1, 16384) # Component synthesized with the default constant value 1. xrt.assert_duckarray_equal(ds["component"], [1]) xrt.assert_duckarray_equal(ds["receiver"], [5908, 5909, 5910]) diff --git a/tests/unit/ingestion/test_schema_effects.py b/tests/unit/ingestion/test_schema_effects.py index 33d10052..96cf1f40 100644 --- a/tests/unit/ingestion/test_schema_effects.py +++ b/tests/unit/ingestion/test_schema_effects.py @@ -7,6 +7,7 @@ from mdio.ingestion.schema import DimensionSpec from mdio.ingestion.schema import ResolvedSchema from mdio.ingestion.segy.index_strategies import IndexStrategyRegistry +from mdio.ingestion.segy.schema_effects import LEGACY_TRACE_DTYPE from mdio.ingestion.segy.schema_effects import CollapseToTraceEffect from mdio.ingestion.segy.schema_effects import InsertTraceDimEffect from mdio.segy.geometry import GridOverrides @@ -56,10 +57,11 @@ def test_non_binned_wires_chunksize_and_dims(self) -> None: assert effect.collapse_dims == ("channel",) def test_has_duplicates_inserts_trace(self) -> None: - """HasDuplicates yields a 1-wide InsertTraceDimEffect by default (backwards compatible).""" + """HasDuplicates yields a 1-wide int32 InsertTraceDimEffect by default (backwards compatible).""" effect = IndexStrategyRegistry().schema_effect(GridOverrides(has_duplicates=True)) assert isinstance(effect, InsertTraceDimEffect) assert effect.chunksize == 1 + assert effect.dtype == LEGACY_TRACE_DTYPE def test_has_duplicates_honours_chunksize(self) -> None: """HasDuplicates with a chunksize sizes the inserted trace dim accordingly.""" @@ -78,9 +80,17 @@ def test_inserts_trace_before_vertical(self) -> None: assert result.chunk_shape == (8, 1, 128, 1, 2048) trace = next(d for d in result.dimensions if d.name == "trace") assert trace.is_calculated is True + # No dtype requested -> the trace coordinate keeps the legacy int32 it always had. + assert trace.dtype == LEGACY_TRACE_DTYPE # Coordinates are unchanged by duplicate handling. assert result.coordinates[0].dimensions == ("shot_point", "cable", "channel") + def test_stores_requested_trace_dtype(self) -> None: + """A requested dtype lands on the stored trace coordinate, not just the counter.""" + result = InsertTraceDimEffect(chunksize=1024, dtype=ScalarType.UINT32).apply(_schema()) + trace = next(d for d in result.dimensions if d.name == "trace") + assert trace.dtype == ScalarType.UINT32 + class TestCollapseToTraceEffect: """CollapseToTraceEffect collapses spatial dims into a single trace dim.""" @@ -90,6 +100,14 @@ def test_default_collapses_all_but_first(self) -> None: result = CollapseToTraceEffect(chunksize=64, collapse_dims=None).apply(_schema()) assert [d.name for d in result.dimensions] == ["shot_point", "trace", "time"] assert result.chunk_shape == (8, 64, 2048) + trace = next(d for d in result.dimensions if d.name == "trace") + assert trace.dtype == LEGACY_TRACE_DTYPE + + def test_stores_requested_trace_dtype(self) -> None: + """The collapsed trace axis honours a requested dtype, like the inserted one.""" + result = CollapseToTraceEffect(chunksize=64, collapse_dims=None, dtype=ScalarType.UINT32).apply(_schema()) + trace = next(d for d in result.dimensions if d.name == "trace") + assert trace.dtype == ScalarType.UINT32 def test_explicit_collapse_dims(self) -> None: """Only the named dims collapse; the coordinate is rewritten onto trace.""" diff --git a/tests/unit/ingestion/test_segy_index_strategies.py b/tests/unit/ingestion/test_segy_index_strategies.py index 9c411ab3..3e957a2b 100644 --- a/tests/unit/ingestion/test_segy_index_strategies.py +++ b/tests/unit/ingestion/test_segy_index_strategies.py @@ -8,17 +8,15 @@ from __future__ import annotations -from typing import TYPE_CHECKING from typing import Any import numpy as np import pytest +from mdio.builder.schemas.dtype import ScalarType from mdio.builder.template_registry import TemplateRegistry +from mdio.builder.templates.base import AbstractDatasetTemplate from mdio.ingestion.segy.index_strategies import ChannelWrappingStrategy - -if TYPE_CHECKING: - from mdio.builder.templates.base import AbstractDatasetTemplate from mdio.ingestion.segy.index_strategies import ComponentSynthesisStrategy from mdio.ingestion.segy.index_strategies import CompositeStrategy from mdio.ingestion.segy.index_strategies import DuplicateHandlingStrategy @@ -26,6 +24,7 @@ from mdio.ingestion.segy.index_strategies import NonBinnedStrategy from mdio.ingestion.segy.index_strategies import RegularGridStrategy from mdio.ingestion.segy.index_strategies import ShotWrappingStrategy +from mdio.ingestion.segy.schema_effects import LEGACY_TRACE_DTYPE from mdio.ingestion.segy.schema_effects import CollapseToTraceEffect from mdio.ingestion.segy.schema_effects import InsertTraceDimEffect from mdio.segy.exceptions import GridOverrideKeysError @@ -102,20 +101,35 @@ def test_non_binned_only(self) -> None: assert strategy.excluded_fields == frozenset({"channel"}) def test_has_duplicates_only(self) -> None: - """``has_duplicates`` -> DuplicateHandlingStrategy with legacy chunk-1/int16 defaults.""" + """``has_duplicates`` -> DuplicateHandlingStrategy with the legacy chunk-1/int32 axis.""" strategy = IndexStrategyRegistry().create_strategy(grid_overrides=GridOverrides(has_duplicates=True)) assert isinstance(strategy, DuplicateHandlingStrategy) - # Backwards compatible: no chunksize/trace_dtype -> chunk 1, int16 (unchanged behavior). - assert strategy.dtype == np.int16 + # Backwards compatible: no chunksize/trace_dtype -> chunk 1, int32 stored trace axis. + assert strategy.dtype is None assert strategy.schema_effect().chunksize == 1 + assert strategy.schema_effect().dtype == LEGACY_TRACE_DTYPE def test_has_duplicates_wires_chunksize_and_dtype(self) -> None: """``has_duplicates`` forwards the override's chunksize and trace_dtype to the strategy.""" overrides = GridOverrides(has_duplicates=True, chunksize=1024, trace_dtype="uint32") strategy = IndexStrategyRegistry().create_strategy(grid_overrides=overrides) assert isinstance(strategy, DuplicateHandlingStrategy) - assert strategy.dtype == "uint32" + assert strategy.dtype == ScalarType.UINT32 assert strategy.schema_effect().chunksize == 1024 + assert strategy.schema_effect().dtype == ScalarType.UINT32 + + def test_non_binned_wires_dtype(self) -> None: + """``trace_dtype`` reaches the NonBinned path too; both insert the same ``trace`` axis.""" + overrides = GridOverrides( + non_binned=True, + chunksize=64, + non_binned_dims=["channel"], + trace_dtype="uint32", + ) + strategy = IndexStrategyRegistry().create_strategy(grid_overrides=overrides) + assert isinstance(strategy, NonBinnedStrategy) + assert strategy.dtype == ScalarType.UINT32 + assert strategy.schema_effect().dtype == ScalarType.UINT32 def test_non_binned_wins_over_has_duplicates(self) -> None: """Both flags set -> NonBinned wins (matches v1.x semantics).""" @@ -199,6 +213,10 @@ def test_plain_template_synthesizes_nothing(self) -> None: template = TemplateRegistry().get("PostStack3DTime") assert _resolve_synthesize_dims(template) == () + def test_hook_is_part_of_the_base_class_contract(self) -> None: + """The hook is a class attribute, so consumers can read it off any template.""" + assert AbstractDatasetTemplate.synthesize_missing_dims == () + # --------------------------------------------------------------------------- # RegularGridStrategy @@ -297,6 +315,24 @@ def test_chunksize_flows_into_schema_effect(self) -> None: assert isinstance(effect, InsertTraceDimEffect) assert effect.chunksize == 1024 + def test_default_counter_dtype_is_legacy_int16(self) -> None: + """Without a dtype the appended counter stays ``int16``, as before the knob existed.""" + headers = _make_struct({"inline": np.array([1, 1, 2], dtype=np.int32)}) + out = DuplicateHandlingStrategy().transform_headers(headers) + assert out["trace"].dtype == np.int16 + + def test_dtype_widens_the_counter(self) -> None: + """A requested dtype is applied to the counter, lifting the int16 trace-count ceiling.""" + headers = _make_struct({"inline": np.array([1, 1, 2], dtype=np.int32)}) + out = DuplicateHandlingStrategy(dtype=ScalarType.UINT32).transform_headers(headers) + assert out["trace"].dtype == np.uint32 + + def test_int16_counter_overflows_loudly(self) -> None: + """Exceeding the counter dtype raises instead of silently wrapping to a wrong index.""" + headers = _make_struct({"inline": np.ones(np.iinfo(np.int16).max + 1, dtype=np.int32)}) + with pytest.raises(OverflowError): + DuplicateHandlingStrategy().transform_headers(headers) + # --------------------------------------------------------------------------- # NonBinnedStrategy diff --git a/tests/unit/test_grid_overrides_pydantic.py b/tests/unit/test_grid_overrides_pydantic.py index 15fa4e65..cae10036 100644 --- a/tests/unit/test_grid_overrides_pydantic.py +++ b/tests/unit/test_grid_overrides_pydantic.py @@ -8,6 +8,7 @@ import pytest from pydantic import ValidationError +from mdio.builder.schemas.dtype import ScalarType from mdio.converters.segy import _coerce_grid_overrides from mdio.segy.geometry import GridOverrides @@ -60,15 +61,16 @@ def test_grid_overrides_rejects_unknown_keys() -> None: def test_grid_overrides_trace_dtype_accepts_valid_dtype() -> None: - """A valid numpy dtype string is accepted on the has_duplicates path.""" + """A wider integer dtype is accepted for the inserted trace dimension.""" overrides = GridOverrides(has_duplicates=True, chunksize=1024, trace_dtype="uint32") - assert overrides.trace_dtype == "uint32" + assert overrides.trace_dtype == ScalarType.UINT32 -def test_grid_overrides_trace_dtype_rejects_garbage() -> None: - """A string numpy cannot parse as a dtype is rejected at construction.""" +@pytest.mark.parametrize("dtype", ["not-a-dtype", "float32", "bool"]) +def test_grid_overrides_trace_dtype_rejects_non_integer(dtype: str) -> None: + """Anything that cannot count traces is rejected at construction, not at ingest.""" with pytest.raises(ValidationError): - GridOverrides(has_duplicates=True, trace_dtype="not-a-dtype") + GridOverrides(has_duplicates=True, trace_dtype=dtype) def test_grid_overrides_trace_dtype_serialization() -> None: From 5ac13c0eb3d6c7b68375854d9a07fd22073d1e67 Mon Sep 17 00:00:00 2001 From: "Andrei.Sorokin" Date: Tue, 28 Jul 2026 11:29:11 +0000 Subject: [PATCH 5/5] Require HasDuplicates for the CRG template Fail early when CRG imports omit the override that supplies their trace dimension, and trim redundant tests and documentation. Co-authored-by: Cursor --- docs/guides/crg_data_import.md | 21 +------------------ docs/guides/grid_overrides.md | 12 +++++++---- src/mdio/commands/segy.py | 7 ++++--- src/mdio/segy/geometry.py | 12 +++++++++-- .../ingestion/test_segy_grid_overrides.py | 19 ----------------- tests/unit/ingestion/test_segy_pipeline.py | 12 +++++++++++ 6 files changed, 35 insertions(+), 48 deletions(-) diff --git a/docs/guides/crg_data_import.md b/docs/guides/crg_data_import.md index b06419d8..35f4eebf 100644 --- a/docs/guides/crg_data_import.md +++ b/docs/guides/crg_data_import.md @@ -75,7 +75,7 @@ A warning is logged when component is synthesized: ## Usage -### Basic Import +### Single-Component Import ```python from segy.schema import HeaderField @@ -108,25 +108,6 @@ segy_to_mdio( ) ``` -### Single-Component Data - -For data without a `component` header field, simply omit it from the spec — the template -synthesizes `component = 1`: - -```python -# Same as above, but without a component field. -crg_headers = [ - HeaderField(name="channel", byte=13, format="int32"), - HeaderField(name="coordinate_scalar", byte=71, format="int16"), - HeaderField(name="group_coord_x", byte=81, format="int32"), - HeaderField(name="group_coord_y", byte=85, format="int32"), - HeaderField(name="receiver_line", byte=137, format="int16"), - HeaderField(name="receiver", byte=139, format="int16"), - HeaderField(name="epoch", byte=189, format="int64"), - # component omitted - will be synthesized -] -``` - ### Exploring the Data ```python diff --git a/docs/guides/grid_overrides.md b/docs/guides/grid_overrides.md index d49950c9..2b6039cf 100644 --- a/docs/guides/grid_overrides.md +++ b/docs/guides/grid_overrides.md @@ -117,15 +117,19 @@ receiver gathers. Some templates have special behaviors that are applied automatically during import, independent of grid overrides. -### Component Synthesis (OBN) +### Component Synthesis -When using the `ObnReceiverGathers3D` template, if the SEG-Y specification does not include a `component` field, MDIO automatically synthesizes it with value `1` for all traces. This allows single-component data (e.g., hydrophone-only) to use the same template without modification. +Templates can declare optional dimensions through `synthesize_missing_dims`. The +`ObnReceiverGathers3D` and `ObnContinuousReceiverGathers3D` templates use this hook for +`component`: if the SEG-Y specification omits that field, MDIO synthesizes it with value +`1` for all traces. This allows single-component data (e.g., hydrophone-only) to use the +same template without modification. ```{note} A warning is logged when component is synthesized: -> SEG-Y headers do not contain 'component' field required by template 'ObnReceiverGathers3D'. -> Synthesizing 'component' dimension with constant value 1 for all traces. +> SEG-Y headers do not contain 'component' field required by template; synthesizing dimension +> with constant value 1 for all traces. ``` ## Error Handling diff --git a/src/mdio/commands/segy.py b/src/mdio/commands/segy.py index e234d9e7..8087c15c 100644 --- a/src/mdio/commands/segy.py +++ b/src/mdio/commands/segy.py @@ -296,14 +296,15 @@ def segy_import( # noqa: PLR0913 \b For dataset with expected duplicate traces we have the following parameterization. This - will use the same logic as NonBinned with a fixed chunksize of 1. The other keys are still - important. The below example allows multiple traces per receiver (i.e. reshoot). + uses the same duplicate-counter logic as NonBinned. The inserted trace dimension defaults + to a chunksize of 1, which can be configured as shown below. The other keys are still + important. This example allows multiple traces per receiver (i.e. reshoot). \b --header-locations 9,213,13 --header-names shot,cable,chan --header-types int32,int16,int32 --chunk-size 8,2,256,512 - --grid-overrides '{"HasDuplicates": True}' + --grid-overrides '{"HasDuplicates": True, "chunksize": 16}' """ # Lazy import to reduce CLI startup time from mdio import segy_to_mdio # noqa: PLC0415 diff --git a/src/mdio/segy/geometry.py b/src/mdio/segy/geometry.py index 34ce2cb7..c699bfe8 100644 --- a/src/mdio/segy/geometry.py +++ b/src/mdio/segy/geometry.py @@ -166,13 +166,14 @@ def validate_overrides_for_template( config: GridOverrides | None, template: AbstractDatasetTemplate | None, ) -> None: - """Reject grid override / template pairings that v1.1 forbade. + """Reject invalid grid override / template pairings. ``auto_shot_wrap`` is streamer-only and ``calculate_shot_index`` is OBN-only; using either with the wrong template silently produced wrong shot indices in v1.1 unless the per-command validator caught it. This is the one guard the :class:`GridOverrides` model cannot enforce on its own (it depends on the chosen template), so the ingestion - pipeline calls it before any header parsing. + pipeline calls it before any header parsing. The CRG template requires + ``has_duplicates`` because that override supplies its undeclared ``trace`` dimension. Args: config: Typed grid overrides, or ``None`` when no overrides were requested. @@ -181,7 +182,14 @@ def validate_overrides_for_template( Raises: TypeError: When ``auto_shot_wrap`` is set without a streamer template, or ``calculate_shot_index`` is set without an OBN receiver-gathers template. + ValueError: When the CRG template is used without ``has_duplicates``. """ + from mdio.builder.templates.seismic_3d_crg import Seismic3DCrgReceiverGathersTemplate # noqa: PLC0415 + + if isinstance(template, Seismic3DCrgReceiverGathersTemplate) and (config is None or not config.has_duplicates): + msg = "ObnContinuousReceiverGathers3D requires the HasDuplicates grid override." + raise ValueError(msg) + if not config: return diff --git a/tests/unit/ingestion/test_segy_grid_overrides.py b/tests/unit/ingestion/test_segy_grid_overrides.py index bdb8063b..0f93aca8 100644 --- a/tests/unit/ingestion/test_segy_grid_overrides.py +++ b/tests/unit/ingestion/test_segy_grid_overrides.py @@ -151,25 +151,6 @@ def test_duplicates(self, mock_streamer_headers: dict[str, npt.NDArray]) -> None assert_array_equal(dims[1].coords, CABLES) assert_array_equal(dims[2].coords, RECEIVERS) - def test_duplicates_with_chunksize(self, mock_streamer_headers: dict[str, npt.NDArray]) -> None: - """HasDuplicates can carry an explicit trace chunksize (CRG tuned path).""" - index_names = ("shot_point", "cable") - grid_overrides = {"HasDuplicates": True, "chunksize": 16} - - streamer_headers = mock_streamer_headers[list(index_names)] - chunksize = (4, 4, 8) - - new_headers, new_names, new_chunks = run_override( - grid_overrides, - index_names, - streamer_headers, - chunksize, - ) - - assert new_names == ("shot_point", "cable", "trace") - # The inserted trace chunk honours the override instead of the legacy hardcoded 1. - assert new_chunks == (4, 4, 16, 8) - def test_non_binned(self, mock_streamer_headers: dict[str, npt.NDArray]) -> None: """Test the NonBinned Grid Override command.""" index_names = ("shot_point", "cable") diff --git a/tests/unit/ingestion/test_segy_pipeline.py b/tests/unit/ingestion/test_segy_pipeline.py index 6b7e0697..368df5d3 100644 --- a/tests/unit/ingestion/test_segy_pipeline.py +++ b/tests/unit/ingestion/test_segy_pipeline.py @@ -118,6 +118,18 @@ def test_non_binned_without_parameters_raises(self) -> None: grid_overrides={"NonBinned": True}, ) + @pytest.mark.parametrize("grid_overrides", [None, {}]) + def test_crg_without_has_duplicates_raises(self, grid_overrides: dict | None) -> None: + """The CRG template requires ``HasDuplicates`` to supply its trace dimension.""" + with pytest.raises(ValueError, match="requires the HasDuplicates grid override"): + segy_to_mdio( + segy_spec=None, + mdio_template=TemplateRegistry().get("ObnContinuousReceiverGathers3D"), + input_path="in.segy", + output_path="out.mdio", + grid_overrides=grid_overrides, + ) + class TestBuildRawHeaderVariables: """Tests for the isolated experimental raw-headers feature."""