diff --git a/docs/guides/crg_data_import.md b/docs/guides/crg_data_import.md new file mode 100644 index 000000000..35f4eebfc --- /dev/null +++ b/docs/guides/crg_data_import.md @@ -0,0 +1,143 @@ +# 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 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 +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 + +### Single-Component 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, trace_dtype="uint32"), + overwrite=True, +) +``` + +### 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/grid_overrides.md b/docs/guides/grid_overrides.md index 783afbbd6..2b6039cfc 100644 --- a/docs/guides/grid_overrides.md +++ b/docs/guides/grid_overrides.md @@ -77,19 +77,59 @@ 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` | - | 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`, 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:** + +```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. -### 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/docs/guides/index.md b/docs/guides/index.md index 08b83b991..3990d67a0 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 970af1a2f..8fdcba141 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/base.py b/src/mdio/builder/templates/base.py index 48d382969..610819add 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/builder/templates/seismic_3d_crg.py b/src/mdio/builder/templates/seismic_3d_crg.py new file mode 100644 index 000000000..a14c92893 --- /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/src/mdio/commands/segy.py b/src/mdio/commands/segy.py index e234d9e78..8087c15cc 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/ingestion/segy/index_strategies.py b/src/mdio/ingestion/segy/index_strategies.py index 3efb9e98c..163fcfaa3 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,18 +131,27 @@ 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: 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. """ 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) 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.""" @@ -152,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 @@ -161,8 +174,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, dtype=self.dtype) class NonBinnedStrategy(DuplicateHandlingStrategy): @@ -179,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__( @@ -187,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__( @@ -200,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): @@ -447,10 +465,17 @@ 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: - 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, + ) + ) if not strategies: return RegularGridStrategy() diff --git a/src/mdio/ingestion/segy/schema_effects.py b/src/mdio/ingestion/segy/schema_effects.py index 1c7538918..fc2d77e6e 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 3828fda36..3dec2819e 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(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 e1c56372d..c699bfe89 100644 --- a/src/mdio/segy/geometry.py +++ b/src/mdio/segy/geometry.py @@ -15,8 +15,10 @@ 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: @@ -25,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.""" @@ -54,18 +70,47 @@ 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: ScalarType | None = Field( + default=None, + 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, 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: 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") def _check_non_binned_parameters(self) -> GridOverrides: """Validate parameters when non_binned is True. @@ -108,31 +153,27 @@ 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 template.synthesize_missing_dims 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. @@ -141,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/integration/test_import_crg_grid_overrides.py b/tests/integration/test_import_crg_grid_overrides.py new file mode 100644 index 000000000..cc6f01f5a --- /dev/null +++ b/tests/integration/test_import_crg_grid_overrides.py @@ -0,0 +1,214 @@ +"""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} + # `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]) + 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} + # 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 86345adfc..96cf1f407 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,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 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.""" + effect = IndexStrategyRegistry().schema_effect(GridOverrides(has_duplicates=True, chunksize=1024)) + assert isinstance(effect, InsertTraceDimEffect) + assert effect.chunksize == 1024 class TestInsertTraceDimEffect: @@ -72,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.""" @@ -84,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_grid_overrides.py b/tests/unit/ingestion/test_segy_grid_overrides.py index d17a27df4..0f93aca83 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 ( diff --git a/tests/unit/ingestion/test_segy_index_strategies.py b/tests/unit/ingestion/test_segy_index_strategies.py index 1ffd1ec96..3e957a2b5 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,9 +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.""" + """``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, 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 == 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).""" @@ -166,6 +191,33 @@ 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) == () + + 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 # --------------------------------------------------------------------------- @@ -252,11 +304,35 @@ 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 + + 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/ingestion/test_segy_pipeline.py b/tests/unit/ingestion/test_segy_pipeline.py index 6b7e0697c..368df5d3c 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.""" diff --git a/tests/unit/test_grid_overrides_pydantic.py b/tests/unit/test_grid_overrides_pydantic.py index 4060c94c8..cae10036c 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 @@ -24,6 +25,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 +60,26 @@ def test_grid_overrides_rejects_unknown_keys() -> None: GridOverrides.model_validate({"FutureFlag": True}) +def test_grid_overrides_trace_dtype_accepts_valid_dtype() -> None: + """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 == ScalarType.UINT32 + + +@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=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) 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 000000000..470295e5d --- /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 37d63ec2b..2576559e5 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):