diff --git a/docs/guides/continuous_receiver_gathers.md b/docs/guides/continuous_receiver_gathers.md new file mode 100644 index 00000000..7c3ff3d6 --- /dev/null +++ b/docs/guides/continuous_receiver_gathers.md @@ -0,0 +1,165 @@ +# Continuous Receiver Gathers + +This guide covers the `NodalContinuousReceiverGathers3D` template for importing continuously recorded receiver data — land nodes and OBN — into MDIO. + +## What makes this data different + +Continuously recording receivers are powered on, left to record for weeks or months, and recovered later. There are no shots to index on. A delivered SEG-Y trace is one fixed-length segment of a single receiver's recording, and the only header that reliably tells segments apart is a 64-bit microsecond epoch timestamp. + +This gives the data two distinct time axes: absolute wall-clock time, which selects a segment, and the sample axis within a segment. + +## Template Overview + +The `NodalContinuousReceiverGathers3D` template organizes data with the following dimensions: + +Dimensions are ordered `receiver_line`, `receiver`, `component`, `segment_index`, `time`, following the field hierarchy: line, then station, then the sensors at that station, then position in time. Keeping `component` next to `segment_index` means all components of one station are adjacent on disk, so a multi-component read of a station stays local. + +| Dimension | Description | +| --------------- | --------------------------------------------------------------------------------------------- | +| `receiver_line` | Receiver line identifier | +| `receiver` | Receiver (station) identifier | +| `component` | Sensor component (e.g., 1=X, 2=Y, 3=Z, 4=Hydrophone) | +| `segment_index` | Calculated dense position in time (see [Calculated segment_index](#calculated-segment_index)) | +| `time` | Sample axis within one segment | + +### Coordinates + +- **Logical coordinates**: `epoch` — exact segment start time in microseconds, spanning the full spatial key `(receiver_line, receiver, component, segment_index)` +- **Physical coordinates**: `group_coord_x`, `group_coord_y` — indexed by `(receiver_line, receiver)`, since a receiver does not move during its deployment + +## Calculated segment_index + +`segment_index` is a calculated dimension, declared the same way `ObnReceiverGathers3D` declares `shot_index`. It is filled at ingest by the `CalculateSegmentIndex` grid override, which ranks each trace's `epoch` among the sorted unique epochs of its own `(receiver_line, receiver, component)` group. The override is mandatory: without it, ingestion fails because nothing else can supply the dimension. + +### Why epoch cannot be the dimension + +Receivers are powered on by hand, so their segment boundaries generally do not line up. No two receivers in a delivery are expected to share an epoch value, which means an `epoch` axis has one entry per trace and the grid becomes `receivers × traces` — a sparsity ratio equal to the receiver count. A 200-receiver deliverable holding four million traces produces a billion-cell grid, which will not build. + +Ranking epochs per receiver collapses this to a dense grid: + +``` +Before (absolute epoch, µs, staggered starts): + Receiver 101: 1556582074780000, 1556582104780000, 1556582134780000, ... + Receiver 102: 1556582081780000, 1556582111780000, 1556582141780000, ... + +After (dense segment_index): + Receiver 101: 0, 1, 2, ... + Receiver 102: 0, 1, 2, ... +``` + +### Why the ranking uses the header, not file order + +A receiver's segments are commonly split across several field records that are not written in time order, so a counter that numbers traces by their position in the file will place segments at the wrong index. Ranking sorts on the `epoch` value itself, so the resulting axis is ordered by recording time regardless of how the vendor laid out the file. + +### Epoch values are preserved exactly + +`segment_index` alone cannot recover wall-clock time, because segment 0 of one receiver may start seconds after segment 0 of another. Ingestion therefore never modifies, rounds, or snaps the epoch: it is written as an `int64` microsecond coordinate spanning every spatial dimension, so a live segment's true recording time remains recoverable and receivers stay comparable across files. + +### Components are ranked independently + +Each component is its own ranking group. If two components of one receiver recorded different segment counts, a given `segment_index` may hold different wall-clock times for each of them. Because `epoch` spans `component`, every cell still reports the true start time of the segment stored there — read `epoch` rather than assuming the axis is aligned across components. + +### Ragged receivers and the epoch fill value + +Receivers in one deliverable often record slightly different segment counts, so the `segment_index` axis is sized to the longest receiver and shorter ones leave empty cells. Those cells have `trace_mask == False`. For `epoch`, empty cells hold the int64 fill sentinel (`np.iinfo(np.int64).max`). Always gate wall-clock queries on `trace_mask` (or an equivalent live-cell test); a one-sided `epoch >= T` filter alone will include dead cells. + +## Chunking + +The default chunk shape is `(1, 1, 1, 180, 15001)` over `(receiver_line, receiver, component, segment_index, time)`: about 10.3 MiB of float32, a little over the 8 MiB shot-template budget. Compression shrinks the stored size further anyway. + +- **One receiver and one component per chunk.** A chunk never mixes recordings from different stations or sensors. +- **180 segments along `segment_index`.** 90 minutes of 30-second data. Wall-clock window reads on one receiver are consecutive on this axis, so packing segments cuts object-store request count without mixing receivers. +- **15,001 samples along `time`.** Matches a 30-second slice at 2 ms sample interval, so a typical segment fits one time chunk with no leftover samples. Longer sub-sampled records at or under that length also fit. + +To change chunking for a different access pattern, set it on the template: + +```python +template = get_template("NodalContinuousReceiverGathers3D") +template.full_chunk_shape = (1, 1, 1, 64, 8192) +``` + +## Special Behaviors + +### Component Synthesis + +When the SEG-Y spec does not include a `component` field, MDIO automatically synthesizes it with value `1` for all traces, so single-component deliverables use the same template unmodified. + +## 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_headers = [ + HeaderField(name="orig_field_record_num", byte=9, format="int32"), + 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"), + HeaderField(name="component", byte=237, format="int16"), +] + +crg_spec = get_segy_standard(1.0).customize(trace_header_fields=crg_headers) + +segy_to_mdio( + input_path="continuous_receiver_gathers.sgy", + output_path="continuous_receiver_gathers.mdio", + segy_spec=crg_spec, + mdio_template=get_template("NodalContinuousReceiverGathers3D"), + grid_overrides=GridOverrides(calculate_segment_index=True), + overwrite=True, +) +``` + +```{note} +Some vendors document the epoch as two 32-bit words rather than one 64-bit field. Declaring a single `int64` field at the first byte reads the same value, so no header arithmetic is needed during ingestion. +``` + +### Exploring the Data + +```python +import numpy as np + +from mdio import open_mdio + +ds = open_mdio("continuous_receiver_gathers.mdio") + +print(ds.sizes) +# {'receiver_line': 1, 'receiver': 232, 'component': 1, 'segment_index': 27108, 'time': 15001} + +# Absolute start time of every segment +print(ds["epoch"].dims) # ('receiver_line', 'receiver', 'component', 'segment_index') + +# Find the live segments of one receiver covering a wall-clock window +receiver = ds.sel(component=1, receiver_line=10, receiver=101) +live = receiver["trace_mask"] +window = live & (receiver["epoch"] >= 1556582074780000) & (receiver["epoch"] < 1556582674780000) +receiver["amplitude"].isel(segment_index=np.flatnonzero(window)).plot() +``` + +## Required Header Fields + +| Field | Required | Notes | +| ------------------- | -------- | ------------------------------------- | +| `receiver_line` | Yes | | +| `receiver` | Yes | | +| `epoch` | Yes | 64-bit microsecond segment start time | +| `component` | No | Synthesized with value 1 if missing | +| `coordinate_scalar` | Yes | | +| `group_coord_x` | Yes | | +| `group_coord_y` | Yes | | + +## 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 783afbbd..0eb68afb 100644 --- a/docs/guides/grid_overrides.md +++ b/docs/guides/grid_overrides.md @@ -77,13 +77,62 @@ segy_to_mdio( See [OBN Data Import](obn_data_import.md) for a complete guide on importing OBN data. ``` +### CalculateSegmentIndex + +**Purpose:** Build the dense `segment_index` axis for continuously recording receivers by ranking each trace's `epoch` header. + +**Supported Templates:** `NodalContinuousReceiverGathers3D` + +**Required Headers:** `epoch`, plus the template's receiver dimensions (`receiver_line`, `receiver`, `component`) + +Each trace's index is the position of its `epoch` among the sorted unique epochs of its own `(receiver_line, receiver, component)` group, so the axis is ordered by recording time however the vendor laid out the file. Epoch values are never modified; they are preserved as an `int64` microsecond coordinate. + +``` +Before (absolute epoch, µs, staggered starts): + Receiver 101: 1556582074780000, 1556582104780000, 1556582134780000, ... + Receiver 102: 1556582081780000, 1556582111780000, 1556582141780000, ... + +After (dense segment_index): + Receiver 101: 0, 1, 2, ... + Receiver 102: 0, 1, 2, ... +``` + +**Usage:** + +```python +from mdio import GridOverrides +from mdio import segy_to_mdio + +segy_to_mdio( + input_path="continuous_receivers.sgy", + output_path="continuous_receivers.mdio", + segy_spec=crg_spec, + mdio_template=get_template("NodalContinuousReceiverGathers3D"), + grid_overrides=GridOverrides(calculate_segment_index=True), +) +``` + +The override is mandatory for this template: nothing else can supply `segment_index`, so ingestion fails without it. It cannot be combined with `HasDuplicates` or `NonBinned`, whose trace counters would group on the freshly computed index. + +**Chunking:** unlike `HasDuplicates`, this override takes no chunking arguments. `segment_index` is a dimension declared by the template, so the template owns its chunk shape and the override only fills in the index values. To change chunking, set it on the template instead: + +```python +template = get_template("NodalContinuousReceiverGathers3D") +template.full_chunk_shape = (1, 1, 1, 64, 8192) # receiver_line, receiver, component, segment_index, time +``` + +```{note} +See [Continuous Receiver Gathers](continuous_receiver_gathers.md) for the full guide, including +chunking rationale, the epoch coordinate, and ragged-grid fill semantics. +``` + ## 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. +The `ObnReceiverGathers3D` and `NodalContinuousReceiverGathers3D` templates declare `component` in `synthesize_missing_dims`. 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. ```{note} A warning is logged when component is synthesized: @@ -111,6 +160,7 @@ GridOverrideKeysError: Grid override 'CalculateShotIndex' requires keys: {'shot_ ## See Also +- [Continuous Receiver Gathers](continuous_receiver_gathers.md) - Template-owned segment ranking - [OBN Data Import](obn_data_import.md) - Complete guide for OBN data - [Template Registry](../template_registry.md) - Available templates - [Tutorials](../tutorials/index.md) - Hands-on guides diff --git a/docs/guides/index.md b/docs/guides/index.md index 08b83b99..536235d6 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 +continuous_receiver_gathers ``` ## Overview @@ -29,3 +30,13 @@ 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. + +### Continuous Receiver Gathers + +Continuously recording land nodes and OBN receivers have no shots to index on, only a timestamp per fixed-length segment. This guide covers: + +- The `NodalContinuousReceiverGathers3D` template +- Why absolute time cannot be a dimension, and how the `CalculateSegmentIndex` override fixes it +- How exact epoch values are preserved, including ragged-receiver fill semantics + +See [Continuous Receiver Gathers](continuous_receiver_gathers.md) for the complete guide. diff --git a/src/mdio/builder/template_registry.py b/src/mdio/builder/template_registry.py index 970af1a2..730453c2 100644 --- a/src/mdio/builder/template_registry.py +++ b/src/mdio/builder/template_registry.py @@ -25,6 +25,9 @@ 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_nodal_continuous_receiver_gathers import ( + Seismic3DNodalContinuousReceiverGathersTemplate, +) 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 +155,9 @@ def _register_default_templates(self) -> None: # OBN (Ocean Bottom Node) data self.register(Seismic3DObnReceiverGathersTemplate()) + # Continuously recording nodal receivers (land nodes and OBN) + self.register(Seismic3DNodalContinuousReceiverGathersTemplate()) + # 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 48d38296..1409204f 100644 --- a/src/mdio/builder/templates/base.py +++ b/src/mdio/builder/templates/base.py @@ -35,6 +35,9 @@ class AbstractDatasetTemplate(ABC): to override specific steps. """ + # Opt-in dims synthesized with a constant when absent from the SEG-Y spec. + synthesize_missing_dims: tuple[str, ...] = () + def __init__(self, data_domain: SeismicDataDomain) -> None: self._data_domain = data_domain.lower() @@ -47,7 +50,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_nodal_continuous_receiver_gathers.py b/src/mdio/builder/templates/seismic_3d_nodal_continuous_receiver_gathers.py new file mode 100644 index 00000000..92b642d0 --- /dev/null +++ b/src/mdio/builder/templates/seismic_3d_nodal_continuous_receiver_gathers.py @@ -0,0 +1,98 @@ +"""Seismic3DNodalContinuousReceiverGathersTemplate MDIO v1 dataset template.""" + +from typing import Any + +from mdio.builder.schemas.dtype import ScalarType +from mdio.builder.schemas.v1.units import TimeUnitEnum +from mdio.builder.schemas.v1.units import TimeUnitModel +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 + +EPOCH_UNIT = TimeUnitModel(time=TimeUnitEnum.MICROSECOND) + + +class Seismic3DNodalContinuousReceiverGathersTemplate(AbstractDatasetTemplate): + """Seismic 3D continuous receiver gathers template (land nodes and ocean-bottom nodes). + + Continuously recording receivers have no shots to grid on. Each SEG-Y trace is a + fixed-length segment of one receiver's recording. ``segment_index`` is a calculated + dimension (like ``shot_index`` on ``ObnReceiverGathers3D``), filled at ingest by + ``GridOverrides(calculate_segment_index=True)``. The override ranks each trace's + ``epoch`` within its ``(receiver_line, receiver, component)`` group. + + Original ``epoch`` values are preserved as an ``int64`` microsecond coordinate spanning + the full spatial key. Empty cells (receivers with fewer segments than the grid axis) + hold the int64 fill sentinel and must be excluded via ``trace_mask``. + + Special handling for the component dimension: + If the SEG-Y spec does not contain a ``component`` field, ingestion synthesizes the + dimension with constant value 1 for all traces and logs a warning. This is driven by + ``synthesize_missing_dims`` and handled by ``ComponentSynthesisStrategy``. + """ + + synthesize_missing_dims = ("component",) + + def __init__(self, data_domain: SeismicDataDomain = "time"): + if data_domain != "time": + msg = "NodalContinuousReceiverGathers3D only supports the time domain, got {data_domain!r}" + raise ValueError(msg.format(data_domain=data_domain)) + + super().__init__(data_domain=data_domain) + + self._spatial_dim_names = ("receiver_line", "receiver", "component", "segment_index") + self._calculated_dims = ("segment_index",) + self._dim_names = (*self._spatial_dim_names, self._data_domain) + self._physical_coord_names = ("group_coord_x", "group_coord_y") + self._logical_coord_names = ("epoch",) + self._var_chunk_shape = (1, 1, 1, 180, 15001) + self.add_units({"epoch": EPOCH_UNIT}) + + @property + def _name(self) -> str: + return "NodalContinuousReceiverGathers3D" + + def _load_dataset_attributes(self) -> dict[str, Any]: + return {"surveyType": "3D", "gatherType": "continuous_receiver"} + + def declare_coordinate_specs(self) -> tuple[CoordinateSpec, ...]: + """Declare receiver- and segment-indexed coordinates for continuous receiver gathers.""" + 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), + CoordinateSpec(name="epoch", dimensions=self.spatial_dimension_names, dtype=ScalarType.INT64), + ) + + def declare_dim_coordinate_types(self) -> DimCoordinateTypes: + """Declare the data types for each dimension coordinate in this template.""" + return { + "receiver_line": ScalarType.UINT32, + "receiver": ScalarType.UINT32, + "component": ScalarType.UINT8, + self._data_domain: ScalarType.INT32, + } + + def _add_coordinates(self) -> None: + # Add dimension coordinates + # EXCLUDE: `segment_index` since it's 0-N (calculated dimension) + for name in ("receiver_line", "receiver", "component", self._data_domain): + self._add_dimension_coordinate(name) + + # Add non-dimension coordinates + 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)), + ) + + self._builder.add_coordinate( + "epoch", + dimensions=self.spatial_dimension_names, + data_type=ScalarType.INT64, + metadata=CoordinateMetadata(units_v1=self.get_unit_by_key("epoch")), + ) diff --git a/src/mdio/builder/templates/seismic_3d_obn.py b/src/mdio/builder/templates/seismic_3d_obn.py index 4ed6c334..697bcf34 100644 --- a/src/mdio/builder/templates/seismic_3d_obn.py +++ b/src/mdio/builder/templates/seismic_3d_obn.py @@ -29,12 +29,13 @@ class Seismic3DObnReceiverGathersTemplate(AbstractDatasetTemplate): ``synthesize_missing_dims`` and handled by ``ComponentSynthesisStrategy``. """ + synthesize_missing_dims = ("component",) + def __init__(self, data_domain: SeismicDataDomain = "time"): super().__init__(data_domain=data_domain) self._spatial_dim_names = ("component", "receiver", "shot_line", "gun", "shot_index") self._calculated_dims = ("shot_index",) - self.synthesize_missing_dims = ("component",) self._dim_names = (*self._spatial_dim_names, self._data_domain) self._physical_coord_names = ( "group_coord_x", diff --git a/src/mdio/ingestion/segy/index_strategies.py b/src/mdio/ingestion/segy/index_strategies.py index 3efb9e98..2b4c7921 100644 --- a/src/mdio/ingestion/segy/index_strategies.py +++ b/src/mdio/ingestion/segy/index_strategies.py @@ -29,6 +29,7 @@ if TYPE_CHECKING: from collections.abc import Iterable + from collections.abc import Sequence from numpy.typing import DTypeLike from segy.arrays import HeaderArray @@ -40,6 +41,57 @@ logger = logging.getLogger(__name__) +def append_header_field(headers: HeaderArray, name: str, values: np.ndarray) -> HeaderArray: + """Append a per-trace field to a header array. + + ``.base`` is None for non-view arrays; fall back to the array itself. + """ + base = headers.base if headers.base is not None else headers + return rfn.append_fields(base, name, values, usemask=False) + + +def rank_within_groups(values: np.ndarray, group_keys: Sequence[np.ndarray]) -> tuple[np.ndarray, np.ndarray]: + """Return the 0-based rank of each value within its group, plus a duplicate mask. + + Ranks are positional offsets within each group after sorting by group then value. + They are dense only when the duplicate mask is all ``False``. + + Args: + values: Per-trace values to rank. + group_keys: Per-trace arrays whose combination identifies a trace's group. + + Returns: + Ranks and a boolean duplicate mask, both aligned with ``values``. + """ + num_traces = len(values) + order = np.lexsort((values, *reversed(group_keys))) + + starts_group = np.zeros(num_traces, dtype=bool) + if num_traces: + starts_group[0] = True + for key in group_keys: + sorted_key = key[order] + starts_group[1:] |= sorted_key[1:] != sorted_key[:-1] + + positions = np.arange(num_traces) + group_start = np.maximum.accumulate(np.where(starts_group, positions, 0)) + ranks = np.empty(num_traces, dtype=np.uint32) + ranks[order] = positions - group_start + + sorted_values = values[order] + duplicates = np.zeros(num_traces, dtype=bool) + duplicates[order[1:]] = ~starts_group[1:] & (sorted_values[1:] == sorted_values[:-1]) + return ranks, duplicates + + +def _binned_spatial_dims(template: AbstractDatasetTemplate | None) -> tuple[str, ...]: + """Spatial dimensions that come from headers, excluding calculated ones.""" + if template is None: + return () + calculated = set(template.calculated_dimension_names) + return tuple(name for name in template.spatial_dimension_names if name not in calculated) + + class IndexStrategy(ABC): """Abstract base for header indexing strategies. @@ -280,9 +332,7 @@ def transform_headers(self, headers: HeaderArray) -> HeaderArray: return headers shot_index = np.empty(len(headers), dtype="uint32") - # `.base` is None for non-view arrays; fall back to the array itself. - base_array = headers.base if headers.base is not None else headers - headers = rfn.append_fields(base_array, "shot_index", shot_index, usemask=False) + headers = append_header_field(headers, "shot_index", shot_index) if geom_type == ShotGunGeometryType.B: for line_val in unique_lines: @@ -299,6 +349,87 @@ def transform_headers(self, headers: HeaderArray) -> HeaderArray: return headers +class HeaderRankingStrategy(IndexStrategy): + """Derive a dense 0-based dimension by ranking a header field within groups of traces. + + Within each ``group_fields`` group, a trace's index is the position of its + ``value_field`` in that group's sorted unique values. The source values are never + modified; templates typically keep the original header as a coordinate. + + Args: + value_field: Header field to rank (e.g. ``epoch``). + index_name: Name of the appended dimension field (e.g. ``segment_index``). + group_fields: Header fields identifying the ranking group. Must be non-empty. + + Raises: + ValueError: If ``group_fields`` is empty. + """ + + def __init__( + self, + value_field: str, + index_name: str, + group_fields: Iterable[str], + ) -> None: + self.value_field = value_field + self.index_name = index_name + self.group_fields = tuple(group_fields) + if not self.group_fields: + msg = ( + f"HeaderRankingStrategy for {index_name!r} requires non-empty group_fields; " + "global ranking is not supported." + ) + raise ValueError(msg) + + @property + def required_keys(self) -> frozenset[str]: + """The ranked field plus every field that defines a ranking group.""" + return frozenset({self.value_field, *self.group_fields}) + + def transform_headers(self, headers: HeaderArray) -> HeaderArray: + """Append `index_name`, the rank of `value_field` within each `group_fields` group.""" + values = np.asarray(headers[self.value_field]) + group_keys = tuple(np.asarray(headers[name]) for name in self.group_fields) + ranks, duplicates = rank_within_groups(values, group_keys) + + if duplicates.any(): + raise ValueError(self._duplicate_message(values, group_keys, duplicates)) + + headers = append_header_field(headers, self.index_name, ranks) + + logger.info( + "Ranked '%s' into dense '%s' across %d group(s) keyed by %s", + self.value_field, + self.index_name, + int(np.count_nonzero(ranks == 0)), + self.group_fields, + ) + return headers + + def _duplicate_message( + self, + values: np.ndarray, + group_keys: tuple[np.ndarray, ...], + duplicates: np.ndarray, + ) -> str: + """Describe the first offending key. Only reached on the failure path, so an extra scan is fine.""" + offender = int(np.flatnonzero(duplicates)[0]) + key_fields = (*self.group_fields, self.value_field) + + matches = values == values[offender] + for key in group_keys: + matches &= key == key[offender] + + key_desc = ", ".join( + f"{name}={key[offender]}" for name, key in zip(key_fields, (*group_keys, values), strict=True) + ) + return ( + f"Duplicate {self.value_field!r} for ranking into {self.index_name!r}: " + f"{key_desc} appears {int(np.count_nonzero(matches))} times. " + f"Each ({', '.join(key_fields)}) combination must be unique." + ) + + class ComponentSynthesisStrategy(IndexStrategy): """Synthesize template-required dimension fields that are absent from the headers. @@ -323,8 +454,7 @@ def transform_headers(self, headers: HeaderArray) -> HeaderArray: dim, ) comp_array = np.ones(len(headers), dtype=np.uint8) - base_array = headers.base if headers.base is not None else headers - headers = rfn.append_fields(base_array, dim, comp_array, usemask=False) + headers = append_header_field(headers, dim, comp_array) return headers @@ -377,22 +507,27 @@ class IndexStrategyRegistry: and the schema view of an override cannot drift. """ - def schema_effect(self, grid_overrides: GridOverrides | None) -> SchemaEffect | None: + def schema_effect( + self, + grid_overrides: GridOverrides | None, + template: AbstractDatasetTemplate | None = None, + ) -> SchemaEffect | None: """Return the schema reshaping implied by `grid_overrides`, if any. Derived from the same strategy that will transform headers, so layout changes stay in - lock-step with the header transform. Template and synthesis hints do not affect the - reshape, so they are omitted when building the strategy here. + lock-step with the header transform. Args: grid_overrides: Typed grid override configuration, or `None`. + template: Template the overrides were validated against. Required by overrides + that read dimension names off the template, such as `calculate_segment_index`. Returns: The matching `SchemaEffect`, or `None` when no layout change applies. """ if not grid_overrides: return None - return self.create_strategy(grid_overrides).schema_effect() + return self.create_strategy(grid_overrides, template=template).schema_effect() def create_strategy( self, @@ -402,7 +537,7 @@ def create_strategy( ) -> IndexStrategy: """Build a strategy (possibly composite) for the given config. - Strategy ordering, when multiple flags are set, mirrors previous behavior: + Strategy ordering, when multiple flags are set: 1. `ComponentSynthesisStrategy` (so later strategies can rely on the synthesized field being present). @@ -410,7 +545,9 @@ def create_strategy( 3. `ShotWrappingStrategy` for `auto_shot_wrap` (streamer; `sail_line`). 4. `ShotWrappingStrategy` for `calculate_shot_index` (OBN; `shot_line`, `always_calculate=True`). - 5. `NonBinnedStrategy` or `DuplicateHandlingStrategy` (mutually exclusive; + 5. `HeaderRankingStrategy` for `calculate_segment_index` (ranks `epoch` into + `segment_index`). + 6. `NonBinnedStrategy` or `DuplicateHandlingStrategy` (mutually exclusive; `non_binned` wins when both are set). Args: @@ -418,7 +555,8 @@ def create_strategy( user-driven overrides. synthesize_dims: Dimensions to synthesize if missing (e.g., `component`). template: Optional dataset template; used to look up coordinate names so - duplicate-handling counters group on dimension fields only. + duplicate-handling counters group on dimension fields only, and to derive + ranking groups for `calculate_segment_index`. Returns: A single `IndexStrategy` instance. Returns `RegularGridStrategy` when no @@ -441,6 +579,15 @@ def create_strategy( if grid_overrides.calculate_shot_index: strategies.append(ShotWrappingStrategy(line_field="shot_line", always_calculate=True)) + if grid_overrides.calculate_segment_index: + strategies.append( + HeaderRankingStrategy( + value_field="epoch", + index_name="segment_index", + group_fields=_binned_spatial_dims(template), + ) + ) + if grid_overrides.non_binned: strategies.append( NonBinnedStrategy( diff --git a/src/mdio/ingestion/segy/pipeline.py b/src/mdio/ingestion/segy/pipeline.py index b2c54fab..2943eed3 100644 --- a/src/mdio/ingestion/segy/pipeline.py +++ b/src/mdio/ingestion/segy/pipeline.py @@ -158,7 +158,7 @@ def segy_to_mdio( # noqa: PLR0913 # Grid overrides are SEG-Y specific: the registry maps them to a schema reshape that the # format-agnostic resolver then applies. - schema_effect = IndexStrategyRegistry().schema_effect(grid_overrides) + schema_effect = IndexStrategyRegistry().schema_effect(grid_overrides, mdio_template) schema = SchemaResolver().resolve(mdio_template, schema_effect) indexed_headers, dimensions = read_index_headers( diff --git a/src/mdio/ingestion/segy/validation.py b/src/mdio/ingestion/segy/validation.py index 3828fda3..f708dcc5 100644 --- a/src/mdio/ingestion/segy/validation.py +++ b/src/mdio/ingestion/segy/validation.py @@ -14,17 +14,13 @@ 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") + # Synthesized dims (e.g. component) need not be in the 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 e1c56372..a168b311 100644 --- a/src/mdio/segy/geometry.py +++ b/src/mdio/segy/geometry.py @@ -17,6 +17,7 @@ from pydantic import Field from pydantic import model_validator +from mdio.segy.exceptions import GridOverrideIncompatibleError from mdio.segy.exceptions import GridOverrideMissingParameterError if TYPE_CHECKING: @@ -46,6 +47,11 @@ class GridOverrides(BaseModel): alias="CalculateShotIndex", description="OBN: derive dense shot_index from sparse shot_point values per shot_line.", ) + calculate_segment_index: bool = Field( + default=False, + alias="CalculateSegmentIndex", + description="Continuous recording: derive dense segment_index by ranking epoch per receiver.", + ) non_binned: bool = Field( default=False, alias="NonBinned", @@ -90,12 +96,32 @@ def _check_non_binned_parameters(self) -> GridOverrides: raise GridOverrideMissingParameterError(command, missing) return self + @model_validator(mode="after") + def _check_segment_index_exclusivity(self) -> GridOverrides: + """Reject pairing ``calculate_segment_index`` with ``non_binned`` or ``has_duplicates``. + + Raises: + GridOverrideIncompatibleError: If incompatible overrides are combined. + + Returns: + The validated GridOverrides instance. + """ + if not self.calculate_segment_index: + return self + + this_command = "CalculateSegmentIndex" + for flag, other_command in ((self.non_binned, "NonBinned"), (self.has_duplicates, "HasDuplicates")): + if flag: + raise GridOverrideIncompatibleError(this_command, other_command) + return self + def __bool__(self) -> bool: """Return True if any override flag is enabled.""" return ( self.auto_channel_wrap or self.auto_shot_wrap or self.calculate_shot_index + or self.calculate_segment_index or self.non_binned or self.has_duplicates ) @@ -105,48 +131,30 @@ def to_legacy_dict(self) -> dict[str, Any]: return self.model_dump(by_alias=True, exclude_defaults=True) -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. - """ - 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 () - - def validate_overrides_for_template( config: GridOverrides | None, template: AbstractDatasetTemplate | None, ) -> None: """Reject grid override / template pairings that v1.1 forbade. - ``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. + ``auto_shot_wrap`` is streamer-only, ``calculate_shot_index`` is OBN-only, and + ``calculate_segment_index`` is continuous-recording-only. These are the guards + :class:`GridOverrides` cannot enforce on its own (they depend on the chosen template), + so the ingestion pipeline calls it before any header parsing. Args: config: Typed grid overrides, or ``None`` when no overrides were requested. template: Template chosen by the caller, or ``None`` if omitted. Raises: - TypeError: When ``auto_shot_wrap`` is set without a streamer template, or - ``calculate_shot_index`` is set without an OBN receiver-gathers template. + TypeError: When an override is paired with an unsupported template. """ if not config: return if config.auto_shot_wrap: - # Lazy import: see ``_resolve_synthesize_dims`` for the cycle rationale. + # 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_streamer_field import ( # noqa: PLC0415 Seismic3DStreamerFieldRecordsTemplate, ) @@ -166,3 +174,16 @@ def validate_overrides_for_template( actual = type(template).__name__ if template is not None else "None" msg = f"calculate_shot_index only supports {Seismic3DObnReceiverGathersTemplate.__name__}, got {actual}." raise TypeError(msg) + + if config.calculate_segment_index: + from mdio.builder.templates.seismic_3d_nodal_continuous_receiver_gathers import ( # noqa: PLC0415 + Seismic3DNodalContinuousReceiverGathersTemplate, + ) + + if not isinstance(template, Seismic3DNodalContinuousReceiverGathersTemplate): + actual = type(template).__name__ if template is not None else "None" + msg = ( + f"calculate_segment_index only supports " + f"{Seismic3DNodalContinuousReceiverGathersTemplate.__name__}, got {actual}." + ) + raise TypeError(msg) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 2f17624e..35af8d24 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -19,6 +19,15 @@ from segy.schema import SegySpec +def _write_segy_factory_file(path: Path, factory: SegyFactory, headers: np.ndarray, samples: np.ndarray) -> Path: + """Write a SegyFactory-built header/sample pair to disk.""" + with 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 path + + def get_segy_mock_4d_spec() -> SegySpec: """Create a mock 4D SEG-Y specification.""" trace_header_fields = [ @@ -283,12 +292,7 @@ def create_segy_mock_obn( # noqa: PLR0913 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 + return _write_segy_factory_file(segy_path, factory, headers, samples) @pytest.fixture(scope="module") @@ -384,6 +388,131 @@ def segy_mock_obn_multiline_type_a(fake_segy_tmp: Path) -> Path: ) +def get_segy_mock_crg_spec(include_component: bool = True) -> SegySpec: + """Create a mock continuous receiver gather SEG-Y specification. + + Args: + include_component: Whether to include the component header field. + + Returns: + SegySpec configured for continuous receiver gather data. + """ + trace_header_fields = [ + HeaderField(name="orig_field_record_num", byte=9, format="int32"), + 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="samples_per_trace", byte=115, format="int16"), + HeaderField(name="sample_interval", byte=117, format="int16"), + HeaderField(name="receiver_line", byte=137, format="int16"), + HeaderField(name="receiver", byte=139, format="int16"), + HeaderField(name="epoch", byte=189, format="int64"), + ] + + if include_component: + trace_header_fields.append(HeaderField(name="component", byte=237, 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 + + +# Mock continuous-recording epochs (microseconds) with staggered receiver starts. +CRG_START_EPOCH = 1_556_582_074_780_000 +CRG_SEGMENT_DURATION = 30_000_000 +CRG_RECEIVER_STAGGER = 7_000_000 +CRG_SEGMENTS_PER_RECORD = 2 + + +def crg_expected_epochs(receiver_index: int, num_segments: int) -> list[int]: + """Return the epochs a mock receiver records, in time order.""" + start = CRG_START_EPOCH + receiver_index * CRG_RECEIVER_STAGGER + return [start + segment * CRG_SEGMENT_DURATION for segment in range(num_segments)] + + +def create_segy_mock_crg( # noqa: PLR0913 + fake_segy_tmp: Path, + num_samples: int, + receiver_line: int, + receivers: list[int], + num_segments: int | list[int], + components: list[int] | None = None, + filename_suffix: str = "", +) -> Path: + """Create a mock continuous receiver gather SEG-Y file for use in tests. + + Args: + fake_segy_tmp: Temporary directory for SEG-Y files. + num_samples: Number of samples per trace. + receiver_line: Receiver line ID shared by all receivers. + receivers: List of receiver IDs. + num_segments: Segments per receiver. A scalar applies to every receiver; a sequence + sets a per-receiver count (ragged grids). + components: List of component IDs. If None, no component header is written. + filename_suffix: Optional suffix for the filename. + + Returns: + Path to the created SEG-Y file. + + Raises: + ValueError: If a ``num_segments`` sequence length does not match ``receivers``. + """ + include_component = components is not None + segy_path = fake_segy_tmp / f"crg{'_' + filename_suffix if filename_suffix else ''}.sgy" + + if isinstance(num_segments, int): + segments_per_receiver = [num_segments] * len(receivers) + else: + segments_per_receiver = list(num_segments) + if len(segments_per_receiver) != len(receivers): + msg = "num_segments sequence length must match receivers" + raise ValueError(msg) + + component_list = components if include_component else [None] + traces = [ + (component, receiver_idx, receiver, segment) + for component in component_list + for receiver_idx, receiver in enumerate(receivers) + for segment in range(segments_per_receiver[receiver_idx]) + ] + np.random.default_rng(seed=42).shuffle(traces) + + 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(len(traces)) + samples = factory.create_trace_sample_template(len(traces)) + + start_x = 700000 + start_y = 4000000 + step_x = 100 + + for trc_idx, (component, receiver_idx, receiver, segment) in enumerate(traces): + n_segments = segments_per_receiver[receiver_idx] + headers["orig_field_record_num"][trc_idx] = 1 + segment // CRG_SEGMENTS_PER_RECORD + headers["channel"][trc_idx] = 1 + segment % CRG_SEGMENTS_PER_RECORD + headers["receiver_line"][trc_idx] = receiver_line + headers["receiver"][trc_idx] = receiver + headers["epoch"][trc_idx] = crg_expected_epochs(receiver_idx, n_segments)[segment] + + if include_component: + headers["component"][trc_idx] = component + + headers["coordinate_scalar"][trc_idx] = -100 + headers["group_coord_x"][trc_idx] = start_x + step_x * receiver_idx + headers["group_coord_y"][trc_idx] = start_y + + # Sample data encodes receiver * 1000 + segment for placement checks + samples[trc_idx] = receiver * 1000 + segment + + return _write_segy_factory_file(segy_path, factory, headers, samples) + + @pytest.fixture(scope="module") def segy_mock_obn_multiline_type_a_sparse(fake_segy_tmp: Path) -> Path: """Generate mock OBN SEG-Y file with Type A geometry and sparse shot points. @@ -415,3 +544,87 @@ def segy_mock_obn_multiline_type_a_sparse(fake_segy_tmp: Path) -> Path: components=components, filename_suffix="multiline_type_a_sparse", ) + + +@pytest.fixture(scope="module") +def segy_mock_crg_with_component(fake_segy_tmp: Path) -> Path: + """Generate mock continuous receiver gather SEG-Y file written out of time order.""" + return create_segy_mock_crg( + fake_segy_tmp, + num_samples=25, + receiver_line=10, + receivers=[101, 102, 103], + num_segments=4, + components=[1, 2], + filename_suffix="with_component", + ) + + +@pytest.fixture(scope="module") +def segy_mock_crg_no_component(fake_segy_tmp: Path) -> Path: + """Generate mock continuous receiver gather SEG-Y file without a component header.""" + return create_segy_mock_crg( + fake_segy_tmp, + num_samples=25, + receiver_line=10, + receivers=[101, 102, 103], + num_segments=4, + components=None, + filename_suffix="no_component", + ) + + +@pytest.fixture(scope="module") +def segy_mock_crg_ragged(fake_segy_tmp: Path) -> Path: + """Generate mock CRG SEG-Y file with unequal segment counts per receiver.""" + return create_segy_mock_crg( + fake_segy_tmp, + num_samples=25, + receiver_line=10, + receivers=[101, 102, 103], + num_segments=[4, 3, 2], + components=[1], + filename_suffix="ragged", + ) + + +@pytest.fixture(scope="module") +def segy_mock_crg_uneven_components(fake_segy_tmp: Path) -> Path: + """Generate mock CRG SEG-Y file with uneven segment coverage across components.""" + include_component = True + receivers = [101] + receiver_line = 10 + num_samples = 25 + # component 1: segments 0,1,2; component 2: segments 0,2 + traces = [ + (1, 0), + (1, 1), + (1, 2), + (2, 0), + (2, 2), + ] + np.random.default_rng(seed=7).shuffle(traces) + + segy_path = fake_segy_tmp / "crg_uneven_components.sgy" + 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(len(traces)) + samples = factory.create_trace_sample_template(len(traces)) + epochs = crg_expected_epochs(0, 3) + + for trc_idx, (component, segment) in enumerate(traces): + headers["orig_field_record_num"][trc_idx] = 1 + headers["channel"][trc_idx] = 1 + headers["receiver_line"][trc_idx] = receiver_line + headers["receiver"][trc_idx] = receivers[0] + headers["epoch"][trc_idx] = epochs[segment] + headers["component"][trc_idx] = component + headers["coordinate_scalar"][trc_idx] = -100 + headers["group_coord_x"][trc_idx] = 700000 + headers["group_coord_y"][trc_idx] = 4000000 + samples[trc_idx] = component * 1000 + segment + + return _write_segy_factory_file(segy_path, factory, headers, samples) diff --git a/tests/integration/test_import_nodal_continuous_receiver_gathers.py b/tests/integration/test_import_nodal_continuous_receiver_gathers.py new file mode 100644 index 00000000..155cae0e --- /dev/null +++ b/tests/integration/test_import_nodal_continuous_receiver_gathers.py @@ -0,0 +1,209 @@ +"""End to end testing for continuous receiver gather SEG-Y to MDIO conversion.""" + +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 tests.integration.conftest import crg_expected_epochs +from tests.integration.conftest import get_segy_mock_crg_spec + +from mdio.api.io import open_mdio +from mdio.builder.template_registry import TemplateRegistry +from mdio.converters.segy import segy_to_mdio +from mdio.segy.geometry import GridOverrides + +if TYPE_CHECKING: + from pathlib import Path + + from xarray import Dataset + +dask.config.set(scheduler="synchronous") +os.environ["MDIO__IMPORT__SAVE_SEGY_FILE_HEADER"] = "true" + +# Mirrors the segy_mock_crg_* fixtures +NUM_SAMPLES = 25 +RECEIVER_LINE = 10 +RECEIVERS = [101, 102, 103] +COMPONENTS = [1, 2] +NUM_SEGMENTS = 4 +TEMPLATE_NAME = "NodalContinuousReceiverGathers3D" +INT64_FILL = np.iinfo(np.int64).max + + +def _ingest(input_path: Path, output_path: Path, include_component: bool) -> Dataset: + """Ingest a mock CRG file with CalculateSegmentIndex enabled.""" + segy_to_mdio( + segy_spec=get_segy_mock_crg_spec(include_component=include_component), + mdio_template=TemplateRegistry().get(TEMPLATE_NAME), + input_path=input_path, + output_path=output_path, + overwrite=True, + grid_overrides=GridOverrides(calculate_segment_index=True), + ) + return open_mdio(output_path) + + +@pytest.fixture(scope="module") +def crg_dataset(segy_mock_crg_with_component: Path, zarr_tmp: Path) -> Dataset: + """Ingest the mock CRG deliverable once for structural assertions.""" + return _ingest(segy_mock_crg_with_component, zarr_tmp, include_component=True) + + +class TestImportContinuousReceiverGathers: + """End-to-end import of continuous receiver gathers with CalculateSegmentIndex.""" + + def test_dimensions(self, crg_dataset: Dataset) -> None: + """Assert amplitude dimension order and coordinate values.""" + assert crg_dataset["amplitude"].dims == ( + "receiver_line", + "receiver", + "component", + "segment_index", + "time", + ) + + xrt.assert_duckarray_equal(crg_dataset["component"], COMPONENTS) + xrt.assert_duckarray_equal(crg_dataset["receiver_line"], [RECEIVER_LINE]) + xrt.assert_duckarray_equal(crg_dataset["receiver"], RECEIVERS) + xrt.assert_duckarray_equal(crg_dataset["time"], list(range(0, NUM_SAMPLES * 2, 2))) + + def test_segment_index_is_dense_and_positional(self, crg_dataset: Dataset) -> None: + """segment_index is a dense 0-N axis; no trace dimension is introduced.""" + assert crg_dataset.sizes["segment_index"] == NUM_SEGMENTS + assert "trace" not in crg_dataset.dims + + def test_grid_is_dense(self, crg_dataset: Dataset) -> None: + """Every grid cell holds a live trace for the rectangular mock.""" + num_traces = len(COMPONENTS) * len(RECEIVERS) * NUM_SEGMENTS + assert int(crg_dataset["trace_mask"].sum()) == num_traces + assert bool(crg_dataset["trace_mask"].all()) + + def test_override_recorded_in_metadata(self, crg_dataset: Dataset) -> None: + """CalculateSegmentIndex is recorded in dataset metadata.""" + assert crg_dataset.attrs["attributes"]["gridOverrides"] == {"CalculateSegmentIndex": True} + assert crg_dataset["segy_file_header"].attrs["binaryHeader"]["samples_per_trace"] == NUM_SAMPLES + + def test_epoch_preserved_exactly(self, crg_dataset: Dataset) -> None: + """Epoch values survive unmodified and span the full spatial key.""" + assert "epoch" in crg_dataset.coords + assert crg_dataset["epoch"].dims == ("receiver_line", "receiver", "component", "segment_index") + assert crg_dataset["epoch"].dtype == np.int64 + + for component in COMPONENTS: + for receiver_idx, receiver in enumerate(RECEIVERS): + expected = crg_expected_epochs(receiver_idx, NUM_SEGMENTS) + actual = crg_dataset["epoch"].sel( + component=component, + receiver_line=RECEIVER_LINE, + receiver=receiver, + ) + xrt.assert_duckarray_equal(actual, expected) + + def test_segments_ordered_by_epoch_not_file_order(self, crg_dataset: Dataset) -> None: + """Segments land by epoch order despite shuffled file order.""" + for receiver in RECEIVERS: + trace = crg_dataset["amplitude"].sel(component=1, receiver_line=RECEIVER_LINE, receiver=receiver) + recorded = trace.values[:, 0] # first sample of each segment + xrt.assert_duckarray_equal(recorded, [receiver * 1000 + s for s in range(NUM_SEGMENTS)]) + + epochs = crg_dataset["epoch"].sel(component=1, receiver_line=RECEIVER_LINE).values + assert np.all(np.diff(epochs, axis=-1) > 0) + + def test_receiver_coordinates_do_not_span_time(self, crg_dataset: Dataset) -> None: + """Receiver positions are indexed by receiver only.""" + for coord_name in ("group_coord_x", "group_coord_y"): + assert crg_dataset[coord_name].dims == ("receiver_line", "receiver") + + # Coordinate scalar of -100 divides the stored header values + xrt.assert_duckarray_equal(crg_dataset["group_coord_x"].squeeze(), [7000.0, 7001.0, 7002.0]) + xrt.assert_duckarray_equal(crg_dataset["group_coord_y"].squeeze(), [40000.0] * len(RECEIVERS)) + + +class TestImportContinuousReceiverGathersSyntheticComponent: + """Import when the SEG-Y spec omits the component header.""" + + def test_import_synthetic_component(self, segy_mock_crg_no_component: Path, zarr_tmp2: Path) -> None: + """Component is synthesized with constant value 1.""" + ds = _ingest(segy_mock_crg_no_component, zarr_tmp2, include_component=False) + + assert "component" in ds.dims + xrt.assert_duckarray_equal(ds["component"], [1]) + xrt.assert_duckarray_equal(ds["receiver"], RECEIVERS) + assert ds.sizes["segment_index"] == NUM_SEGMENTS + assert bool(ds["trace_mask"].all()) + + +class TestImportContinuousReceiverGathersRagged: + """Import when receivers have unequal segment counts.""" + + def test_ragged_grid_and_epoch_sentinel(self, segy_mock_crg_ragged: Path, tmp_path: Path) -> None: + """Short receivers leave empty cells filled with the int64 sentinel.""" + ds = _ingest(segy_mock_crg_ragged, tmp_path / "ragged.mdio", include_component=True) + + assert ds.sizes["segment_index"] == 4 # max across receivers + live = int(ds["trace_mask"].sum()) + cells = int(np.prod([ds.sizes[k] for k in ("receiver_line", "receiver", "component", "segment_index")])) + assert live == 4 + 3 + 2 + assert live < cells + + # Receiver 103 recorded only 2 segments; later cells are dead with the int64 fill. + short = ds.sel(component=1, receiver_line=RECEIVER_LINE, receiver=103) + assert short["trace_mask"].values.tolist() == [True, True, False, False] + assert int(short["epoch"].values[2]) == INT64_FILL + assert int(short["epoch"].values[3]) == INT64_FILL + live_epochs = short["epoch"].values[short["trace_mask"].values] + np.testing.assert_array_equal(live_epochs, crg_expected_epochs(2, 2)) + + +class TestImportContinuousReceiverGathersUnevenComponents: + """Import when components of one receiver miss different segments.""" + + def test_each_component_indexed_from_its_own_epochs( + self, + segy_mock_crg_uneven_components: Path, + tmp_path: Path, + ) -> None: + """Each component ranks independently; missing segments shift only that component.""" + ds = _ingest(segy_mock_crg_uneven_components, tmp_path / "uneven.mdio", include_component=True) + + assert ds.sizes["segment_index"] == 3 + mask = ds["trace_mask"].sel(receiver_line=RECEIVER_LINE, receiver=101) + assert mask.sel(component=1).values.tolist() == [True, True, True] + assert mask.sel(component=2).values.tolist() == [True, True, False] + + all_epochs = crg_expected_epochs(0, 3) + epochs = ds["epoch"].sel(receiver_line=RECEIVER_LINE, receiver=101) + np.testing.assert_array_equal(epochs.sel(component=1).values, all_epochs) + # Component 2 holds epochs 0 and 60 packed at index 0 and 1. + np.testing.assert_array_equal( + epochs.sel(component=2).values, + [all_epochs[0], all_epochs[2], INT64_FILL], + ) + + amp = ds["amplitude"].sel(receiver_line=RECEIVER_LINE, receiver=101) + assert amp.sel(component=1).values[2, 0] == 1000 + 2 + assert amp.sel(component=2).values[1, 0] == 2000 + 2 + + +class TestContinuousReceiverGathersOverrideGuards: + """Guards for misconfigured continuous receiver gather imports.""" + + def test_ingest_without_override_fails( + self, + segy_mock_crg_with_component: Path, + tmp_path: Path, + ) -> None: + """Ingestion fails when CalculateSegmentIndex is omitted.""" + with pytest.raises(ValueError, match=r"Required computed fields \['segment_index'\]"): + segy_to_mdio( + segy_spec=get_segy_mock_crg_spec(include_component=True), + mdio_template=TemplateRegistry().get(TEMPLATE_NAME), + input_path=segy_mock_crg_with_component, + output_path=tmp_path / "no_override.mdio", + overwrite=True, + ) diff --git a/tests/unit/ingestion/test_segy_grid_overrides.py b/tests/unit/ingestion/test_segy_grid_overrides.py index d17a27df..0e9b52ab 100644 --- a/tests/unit/ingestion/test_segy_grid_overrides.py +++ b/tests/unit/ingestion/test_segy_grid_overrides.py @@ -14,12 +14,13 @@ from numpy import unique from numpy.testing import assert_array_equal +from mdio.builder.template_registry import TemplateRegistry from mdio.core import Dimension from mdio.ingestion.segy.index_strategies import IndexStrategyRegistry +from mdio.segy.exceptions import GridOverrideIncompatibleError from mdio.segy.exceptions import GridOverrideMissingParameterError from mdio.segy.exceptions import GridOverrideUnknownError from mdio.segy.geometry import GridOverrides -from mdio.segy.geometry import _resolve_synthesize_dims from mdio.segy.geometry import validate_overrides_for_template if TYPE_CHECKING: @@ -55,7 +56,7 @@ def run_override( # production pipeline so this helper cannot drift from what `segy_to_mdio` runs. validate_overrides_for_template(config, template) - synthesize_dims = _resolve_synthesize_dims(template) + synthesize_dims = template.synthesize_missing_dims if template is not None else () registry = IndexStrategyRegistry() strategy = registry.create_strategy( grid_overrides=config, @@ -283,3 +284,59 @@ def test_calculate_shot_index_obn(self) -> None: gun2_shot_indices = np.unique(new_headers["shot_index"][gun2_mask]) assert_array_equal(gun1_shot_indices, [0, 1, 2]) assert_array_equal(gun2_shot_indices, [1, 2, 3]) + + +class TestCalculateSegmentIndex: + """Tests for the CalculateSegmentIndex grid override.""" + + CRG_TEMPLATE_NAME = "NodalContinuousReceiverGathers3D" + + @staticmethod + def _crg_headers() -> npt.NDArray: + """Two receivers, three segments each, written out of time order.""" + hdr_dtype = np.dtype( + { + "names": ["receiver_line", "receiver", "component", "epoch"], + "formats": ["uint32", "uint32", "uint8", "int64"], + } + ) + start = 1_556_582_074_780_000 + step = 30_000_000 + records = [ + (10, receiver, 1, start + receiver_idx * 7_000_000 + segment * step) + for receiver_idx, receiver in enumerate((101, 102)) + for segment in (2, 0, 1) + ] + return np.array(records, dtype=hdr_dtype) + + def test_ranks_epoch_into_segment_index(self) -> None: + """Ranks each receiver's epochs into a dense segment_index.""" + headers = self._crg_headers() + new_headers, _, _ = run_override( + {"CalculateSegmentIndex": True}, + ("receiver_line", "receiver", "component", "epoch"), + headers, + template=TemplateRegistry().get(self.CRG_TEMPLATE_NAME), + ) + + assert "segment_index" in new_headers.dtype.names + assert_array_equal(new_headers["segment_index"], [2, 0, 1, 2, 0, 1]) + assert_array_equal(new_headers["epoch"], headers["epoch"]) + + def test_requires_continuous_receiver_template(self) -> None: + """Rejects CalculateSegmentIndex with a non-CRG template.""" + with pytest.raises(TypeError, match="calculate_segment_index only supports"): + validate_overrides_for_template( + GridOverrides(calculate_segment_index=True), + TemplateRegistry().get("ObnReceiverGathers3D"), + ) + + @pytest.mark.parametrize("other", ["has_duplicates", "non_binned"]) + def test_rejects_trace_dimension_overrides(self, other: str) -> None: + """Rejects pairing CalculateSegmentIndex with NonBinned or HasDuplicates.""" + extra: dict[str, Any] = {other: True} + if other == "non_binned": + extra |= {"chunksize": 8, "non_binned_dims": ["receiver"]} + + with pytest.raises(GridOverrideIncompatibleError, match="CalculateSegmentIndex"): + GridOverrides(calculate_segment_index=True, **extra) diff --git a/tests/unit/ingestion/test_segy_index_strategies.py b/tests/unit/ingestion/test_segy_index_strategies.py index 1ffd1ec9..f104db9b 100644 --- a/tests/unit/ingestion/test_segy_index_strategies.py +++ b/tests/unit/ingestion/test_segy_index_strategies.py @@ -22,15 +22,16 @@ from mdio.ingestion.segy.index_strategies import ComponentSynthesisStrategy from mdio.ingestion.segy.index_strategies import CompositeStrategy from mdio.ingestion.segy.index_strategies import DuplicateHandlingStrategy +from mdio.ingestion.segy.index_strategies import HeaderRankingStrategy from mdio.ingestion.segy.index_strategies import IndexStrategyRegistry 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.index_strategies import append_header_field from mdio.ingestion.segy.schema_effects import CollapseToTraceEffect from mdio.ingestion.segy.schema_effects import InsertTraceDimEffect from mdio.segy.exceptions import GridOverrideKeysError from mdio.segy.geometry import GridOverrides -from mdio.segy.geometry import _resolve_synthesize_dims from mdio.segy.geometry import validate_overrides_for_template @@ -47,7 +48,7 @@ def run( """Run mock strategy validation and transform.""" config = GridOverrides.model_validate(grid_overrides) validate_overrides_for_template(config, template) - synthesize_dims = _resolve_synthesize_dims(template) + synthesize_dims = template.synthesize_missing_dims if template is not None else () registry = IndexStrategyRegistry() strategy = registry.create_strategy( grid_overrides=config, @@ -155,6 +156,21 @@ def test_calculate_shot_index_uses_shot_line(self) -> None: assert strategy.line_field == "shot_line" assert strategy.always_calculate is True + def test_calculate_segment_index_ranks_epoch_per_receiver(self) -> None: + """calculate_segment_index builds HeaderRankingStrategy over template spatial dims.""" + template = TemplateRegistry().get("NodalContinuousReceiverGathers3D") + strategy = IndexStrategyRegistry().create_strategy( + grid_overrides=GridOverrides(calculate_segment_index=True), + synthesize_dims=template.synthesize_missing_dims, + template=template, + ) + assert isinstance(strategy, CompositeStrategy) + ranking = strategy.strategies[-1] + assert isinstance(ranking, HeaderRankingStrategy) + assert ranking.value_field == "epoch" + assert ranking.index_name == "segment_index" + assert ranking.group_fields == ("receiver_line", "receiver", "component") + def test_template_coord_names_propagate_to_duplicate_strategy(self) -> None: """Template coordinates flow into the duplicate-handling strategy as exclusions.""" template = TemplateRegistry().get("StreamerShotGathers3D") @@ -465,6 +481,124 @@ def test_existing_field_left_alone(self) -> None: np.testing.assert_array_equal(out["component"], [2, 3, 4]) +# --------------------------------------------------------------------------- +# HeaderRankingStrategy +# --------------------------------------------------------------------------- + + +def _continuous_headers(receivers: np.ndarray, epochs: np.ndarray) -> np.ndarray: + """Build continuous-recording headers for one component and one receiver line.""" + return _make_struct( + { + "component": np.ones(len(receivers), dtype=np.uint8), + "receiver_line": np.full(len(receivers), 10, dtype=np.uint32), + "receiver": receivers.astype(np.uint32), + "epoch": epochs.astype(np.int64), + } + ) + + +class TestHeaderRankingStrategy: + """Tests for HeaderRankingStrategy.""" + + GROUP_FIELDS = ("receiver_line", "receiver", "component") + + def _strategy(self) -> HeaderRankingStrategy: + return HeaderRankingStrategy( + value_field="epoch", + index_name="segment_index", + group_fields=self.GROUP_FIELDS, + ) + + def test_ranks_by_header_not_file_order(self) -> None: + """Ranks by header value, not by file order.""" + headers = _continuous_headers( + receivers=np.array([101, 101, 101, 101]), + epochs=np.array([90, 30, 60, 0]), + ) + out = self._strategy().transform_headers(headers) + np.testing.assert_array_equal(out["segment_index"], [3, 1, 2, 0]) + + def test_ranks_independently_per_group(self) -> None: + """Each group gets its own dense 0-based index.""" + headers = _continuous_headers( + receivers=np.array([101, 101, 101, 102, 102, 102]), + epochs=np.array([0, 30, 60, 7, 37, 67]), + ) + out = self._strategy().transform_headers(headers) + np.testing.assert_array_equal(out["segment_index"], [0, 1, 2, 0, 1, 2]) + + def test_ranks_each_component_independently(self) -> None: + """Uneven coverage per component ranks each component independently.""" + headers = _make_struct( + { + "component": np.array([1, 1, 1, 2, 2], dtype=np.uint8), + "receiver_line": np.full(5, 10, dtype=np.uint32), + "receiver": np.full(5, 101, dtype=np.uint32), + "epoch": np.array([0, 30, 60, 0, 60], dtype=np.int64), + } + ) + out = self._strategy().transform_headers(headers) + np.testing.assert_array_equal(out["segment_index"], [0, 1, 2, 0, 1]) + + def test_source_values_are_untouched(self) -> None: + """Ranked header values are left unmodified.""" + epochs = np.array([1_556_582_074_780_000, 1_556_582_044_780_000], dtype=np.int64) + headers = _continuous_headers(receivers=np.array([101, 101]), epochs=epochs) + out = self._strategy().transform_headers(headers) + np.testing.assert_array_equal(out["epoch"], epochs) + + def test_gaps_do_not_break_density(self) -> None: + """Gaps in values still produce contiguous ranks.""" + headers = _continuous_headers( + receivers=np.array([101, 101, 101]), + epochs=np.array([0, 30, 6000]), # long gap before the last segment + ) + out = self._strategy().transform_headers(headers) + np.testing.assert_array_equal(out["segment_index"], [0, 1, 2]) + + def test_index_dtype_spans_long_deployments(self) -> None: + """Index dtype is uint32.""" + headers = _continuous_headers(receivers=np.array([101]), epochs=np.array([0])) + out = self._strategy().transform_headers(headers) + assert out.dtype["segment_index"] == np.uint32 + + def test_empty_group_fields_raises(self) -> None: + """Empty group_fields raises ValueError.""" + with pytest.raises(ValueError, match="non-empty group_fields"): + HeaderRankingStrategy(value_field="epoch", index_name="segment_index", group_fields=()) + + def test_duplicate_epoch_raises(self) -> None: + """Duplicate value_field within a group raises ValueError.""" + headers = _continuous_headers( + receivers=np.array([101, 101]), + epochs=np.array([100, 100]), + ) + with pytest.raises(ValueError, match=r"Duplicate 'epoch'.*receiver=101"): + self._strategy().transform_headers(headers) + + def test_missing_required_field_raises(self) -> None: + """Missing required fields raise GridOverrideKeysError.""" + headers = _make_struct({"receiver": np.array([1, 2], dtype=np.uint32)}) + with pytest.raises(GridOverrideKeysError, match="HeaderRankingStrategy"): + self._strategy().validate_headers(headers) + + def test_no_schema_effect(self) -> None: + """HeaderRankingStrategy does not reshape the schema.""" + assert self._strategy().schema_effect() is None + + +class TestAppendHeaderField: + """Tests for append_header_field.""" + + def test_appends_to_full_array(self) -> None: + """Appends a new field aligned with existing rows.""" + headers = _make_struct({"receiver": np.array([1, 2, 3], dtype=np.uint32)}) + out = append_header_field(headers, "segment_index", np.array([0, 1, 2], dtype=np.uint32)) + np.testing.assert_array_equal(out["segment_index"], [0, 1, 2]) + assert len(out) == 3 + + # --------------------------------------------------------------------------- # CompositeStrategy # --------------------------------------------------------------------------- diff --git a/tests/unit/v1/templates/test_seismic_3d_nodal_continuous_receiver_gathers.py b/tests/unit/v1/templates/test_seismic_3d_nodal_continuous_receiver_gathers.py new file mode 100644 index 00000000..ec600bbe --- /dev/null +++ b/tests/unit/v1/templates/test_seismic_3d_nodal_continuous_receiver_gathers.py @@ -0,0 +1,215 @@ +"""Unit tests for Seismic3DNodalContinuousReceiverGathersTemplate.""" + +import numpy as np +import pytest +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.dataset import Dataset +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_nodal_continuous_receiver_gathers import ( + Seismic3DNodalContinuousReceiverGathersTemplate, +) +from mdio.core.utils_write import MAX_COORDINATES_BYTES +from mdio.core.utils_write import get_constrained_chunksize +from mdio.ingestion.dataset_factory import build_mdio_dataset +from mdio.ingestion.schema.resolver import SchemaResolver + +UNITS_METER = LengthUnitModel(length=LengthUnitEnum.METER) +UNITS_SECOND = TimeUnitModel(time=TimeUnitEnum.SECOND) + +# Typical continuous receiver deliverable scale used by chunking tests. +DATASET_SIZE_MAP = { + "receiver_line": 1, + "receiver": 232, + "component": 1, + "segment_index": 27108, + "time": 15001, +} +DATASET_DTYPE_MAP = { + "receiver_line": "uint32", + "receiver": "uint32", + "component": "uint8", + "time": "int32", +} +EXPECTED_COORDINATES = ["group_coord_x", "group_coord_y", "epoch"] +RECEIVER_DIMS = ("receiver_line", "receiver") +EPOCH_DIMS = ("receiver_line", "receiver", "component", "segment_index") +EXPECTED_CHUNK_SHAPE = (1, 1, 1, 180, 15001) + + +def _validate_coordinates_headers_trace_mask(dataset: Dataset, headers: StructuredType, domain: str) -> None: + """Validate the coordinate, headers, trace_mask variables in the dataset.""" + # 4 dim coords (excl. segment_index which is 0-N) + 3 non-dim coords + 1 data + 1 trace mask + # + 1 headers = 10 variables + assert len(dataset.variables) == 10 + + validate_variable( + dataset, + name="headers", + dims=[(k, v) for k, v in DATASET_SIZE_MAP.items() if k != domain], + coords=EXPECTED_COORDINATES, + dtype=headers, + ) + + validate_variable( + dataset, + name="trace_mask", + dims=[(k, v) for k, v in DATASET_SIZE_MAP.items() if k != domain], + coords=EXPECTED_COORDINATES, + dtype=ScalarType.BOOL, + ) + + # Verify dimension coordinate variables (excluding segment_index which is calculated 0-N) + for dim_name, dim_size in DATASET_SIZE_MAP.items(): + if dim_name == "segment_index": + continue + validate_variable( + dataset, + name=dim_name, + dims=[(dim_name, dim_size)], + coords=[dim_name], + dtype=ScalarType(DATASET_DTYPE_MAP[dim_name]), + ) + + # Verify receiver coordinate variables (indexed by receiver only) + for coord_name in ["group_coord_x", "group_coord_y"]: + coord = validate_variable( + dataset, + name=coord_name, + dims=[(k, DATASET_SIZE_MAP[k]) for k in RECEIVER_DIMS], + coords=[coord_name], + dtype=ScalarType.FLOAT64, + ) + assert coord.metadata.units_v1.length == LengthUnitEnum.METER + + # Verify epoch coordinate (indexed by full spatial key) + epoch = validate_variable( + dataset, + name="epoch", + dims=[(k, DATASET_SIZE_MAP[k]) for k in EPOCH_DIMS], + coords=["epoch"], + dtype=ScalarType.INT64, + ) + assert epoch.metadata.units_v1.time == TimeUnitEnum.MICROSECOND + + +class TestSeismic3DNodalContinuousReceiverGathersTemplate: + """Unit tests for Seismic3DNodalContinuousReceiverGathersTemplate.""" + + def test_configuration(self) -> None: + """Test template configuration and attributes.""" + t = Seismic3DNodalContinuousReceiverGathersTemplate() + + assert t.name == "NodalContinuousReceiverGathers3D" + assert t._dim_names == ("receiver_line", "receiver", "component", "segment_index", "time") + assert t._calculated_dims == ("segment_index",) + assert t._physical_coord_names == ("group_coord_x", "group_coord_y") + assert t._logical_coord_names == ("epoch",) + assert t._var_chunk_shape == (1, 1, 1, 180, 15001) + + # Variables instantiated when build_dataset() is called + assert t._builder is None + assert t._dim_sizes == () + + attrs = t._load_dataset_attributes() + assert attrs == {"surveyType": "3D", "gatherType": "continuous_receiver"} + assert t.default_variable_name == "amplitude" + + def test_component_is_synthesized_when_missing(self) -> None: + """Component is listed in synthesize_missing_dims.""" + t = Seismic3DNodalContinuousReceiverGathersTemplate() + + assert t.synthesize_missing_dims == ("component",) + + def test_segment_index_is_calculated(self) -> None: + """segment_index is a calculated dimension with no dim coordinate type.""" + t = Seismic3DNodalContinuousReceiverGathersTemplate() + + assert t.calculated_dimension_names == ("segment_index",) + assert "segment_index" not in t.declare_dim_coordinate_types() + + def test_chunk_matches_ninety_minute_segment_budget(self) -> None: + """Default chunk shape matches EXPECTED_CHUNK_SHAPE and exceeds 8 MiB float32.""" + t = Seismic3DNodalContinuousReceiverGathersTemplate() + t._dim_sizes = tuple(DATASET_SIZE_MAP.values()) + + chunk_shape = t.full_chunk_shape + assert chunk_shape == EXPECTED_CHUNK_SHAPE + + chunk_bytes = int(np.prod(chunk_shape)) * 4 + assert chunk_bytes == 180 * 15001 * 4 + assert chunk_bytes > 8 * 1024 * 1024 + + def test_chunking_keeps_one_receiver_packs_segments(self) -> None: + """Chunk keeps one receiver/component; packs 180 segments and 15001 samples.""" + t = Seismic3DNodalContinuousReceiverGathersTemplate() + + chunk_shape = t._var_chunk_shape + assert chunk_shape[0] == 1 # receiver_line + assert chunk_shape[1] == 1 # receiver + assert chunk_shape[2] == 1 # component + assert chunk_shape[3] == 180 # segment_index + assert chunk_shape[4] == 15001 # time + + def test_build_dataset(self, structured_headers: StructuredType) -> None: + """Test building a complete dataset with the template.""" + t = Seismic3DNodalContinuousReceiverGathersTemplate() + + 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("ContinuousSurvey3D", sizes=sizes, header_dtype=structured_headers) + + assert dataset.metadata.name == "ContinuousSurvey3D" + assert dataset.metadata.attributes["surveyType"] == "3D" + assert dataset.metadata.attributes["gatherType"] == "continuous_receiver" + assert dataset.metadata.attributes["defaultVariableName"] == "amplitude" + + _validate_coordinates_headers_trace_mask(dataset, structured_headers, "time") + + 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 == EXPECTED_CHUNK_SHAPE + assert seismic.metadata.stats_v1 is None + + def test_depth_domain_rejected(self) -> None: + """Depth domain is rejected.""" + with pytest.raises(ValueError, match="only supports the time domain"): + Seismic3DNodalContinuousReceiverGathersTemplate(data_domain="depth") + + def test_epoch_chunked_under_coordinate_cap_at_advertised_scale(self) -> None: + """Epoch chunking stays under MAX_COORDINATES_BYTES at production scale.""" + template = Seismic3DNodalContinuousReceiverGathersTemplate() + schema = SchemaResolver().resolve(template) + sizes = tuple(DATASET_SIZE_MAP.values()) + dataset = build_mdio_dataset(schema=schema, sizes=sizes) + + epoch = next(v for v in dataset.variables if v.name == "epoch") + assert epoch.metadata is not None + assert epoch.metadata.chunk_grid is not None + chunk_shape = epoch.metadata.chunk_grid.configuration.chunk_shape + + expected = get_constrained_chunksize( + shape=tuple(DATASET_SIZE_MAP[k] for k in EPOCH_DIMS), + dtype="int64", + max_bytes=MAX_COORDINATES_BYTES, + ) + assert chunk_shape == expected + assert int(np.prod(chunk_shape)) * 8 <= MAX_COORDINATES_BYTES diff --git a/tests/unit/v1/templates/test_template_registry.py b/tests/unit/v1/templates/test_template_registry.py index 37d63ec2..c9304a29 100644 --- a/tests/unit/v1/templates/test_template_registry.py +++ b/tests/unit/v1/templates/test_template_registry.py @@ -38,9 +38,12 @@ "StreamerShotGathers3D", "StreamerFieldRecords3D", "ObnReceiverGathers3D", + "NodalContinuousReceiverGathers3D", "ShotReceiverLineGathers3D", ] +NUM_DEFAULT_TEMPLATES = len(EXPECTED_DEFAULT_TEMPLATE_NAMES) + class MockDatasetTemplate(AbstractDatasetTemplate): """Mock template for testing.""" @@ -245,7 +248,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 assert "Template_One" in templates assert "Template_Two" in templates @@ -255,7 +258,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 +267,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 # Clear all registry.clear() @@ -397,7 +400,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 assert "template1" in templates assert "template2" in templates @@ -440,7 +443,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 # Check all templates are registered for i in range(10):