Ensure header overrides functionality of TGSAI/segy is available and applied. (TGSAI#700)

* Support SEG-Y header overrides during file ingestion

* Clean up redundant SEG-Y file initialization and fix import naming inconsistency

* Refactor SEG-Y file initialization and adjust import formatting for clarity

* refine coordinate unit extraction logic

* Refactor coordinate scalar handling: extract scalar logic into `scalar.py`, apply scalars to coordinates, and streamline related imports

* Add unit tests for coordinate scalar getters and apply functions

* Fix typo in coordinate scalar key: `cdp_y_` → `cdp_y`

* remove redundant checks

* more informative warning message

* update schema chunk shape for shots

* ensure result it wrapped in HeaderArray

* add tutorial for dealing with corrupt SEG-Y files.

* nicer syntax

* fix broken tests after schema change

---------

Co-authored-by: Altay Sansal <[email protected]>

Author: tasansal

docs/tutorials/corrupt_files.ipynb

@@ -7,4 +7,5 @@ quickstart
 creation
 compression
 rechunking
+corrupt_files
 ```

docs/tutorials/index.md

@@ -18,7 +18,7 @@ def __init__(self, data_domain: SeismicDataDomain):
         self._coord_dim_names = ("shot_point", "channel")
         self._dim_names = (*self._coord_dim_names, self._data_domain)
         self._coord_names = ("gun", "source_coord_x", "source_coord_y", "group_coord_x", "group_coord_y")
-        self._var_chunk_shape = (16, 64, 1024)
+        self._var_chunk_shape = (16, 32, 2048)
 
     @property
     def _name(self) -> str:

src/mdio/builder/templates/seismic_2d_prestack_shot.py

@@ -18,7 +18,7 @@ def __init__(self, data_domain: SeismicDataDomain):
         self._coord_dim_names = ("shot_point", "cable", "channel")
         self._dim_names = (*self._coord_dim_names, self._data_domain)
         self._coord_names = ("gun", "source_coord_x", "source_coord_y", "group_coord_x", "group_coord_y")
-        self._var_chunk_shape = (8, 2, 128, 1024)
+        self._var_chunk_shape = (8, 1, 128, 2048)
 
     @property
     def _name(self) -> str:

src/mdio/builder/templates/seismic_3d_prestack_shot.py

@@ -12,8 +12,9 @@
 import zarr
 from segy import SegyFile
 from segy.config import SegyFileSettings
-from segy.standards.codes import MeasurementSystem as segy_MeasurementSystem
-from segy.standards.fields.trace import Rev0 as TraceHeaderFieldsRev0
+from segy.config import SegyHeaderOverrides
+from segy.standards.codes import MeasurementSystem as SegyMeasurementSystem
+from segy.standards.fields import binary as binary_header_fields
 
 from mdio.api.io import _normalize_path
 from mdio.api.io import to_mdio
@@ -36,6 +37,9 @@
 from mdio.core.utils_write import MAX_SIZE_LIVE_MASK
 from mdio.core.utils_write import get_constrained_chunksize
 from mdio.segy import blocked_io
+from mdio.segy.scalar import SCALE_COORDINATE_KEYS
+from mdio.segy.scalar import _apply_coordinate_scalar
+from mdio.segy.scalar import _get_coordinate_scalar
 from mdio.segy.utilities import get_grid_plan
 
 if TYPE_CHECKING:
@@ -54,6 +58,9 @@
 logger = logging.getLogger(__name__)
 
 
+MEASUREMENT_SYSTEM_KEY = binary_header_fields.Rev0.MEASUREMENT_SYSTEM_CODE.model.name
+
+
 def grid_density_qc(grid: Grid, num_traces: int) -> None:
     """Quality control for sensible grid density during SEG-Y to MDIO conversion.
 
@@ -269,6 +276,7 @@ def populate_non_dim_coordinates(
     grid: Grid,
     coordinates: dict[str, SegyHeaderArray],
     drop_vars_delayed: list[str],
+    horizontal_coordinate_scalar: int,
 ) -> tuple[xr_Dataset, list[str]]:
     """Populate the xarray dataset with coordinate variables."""
     non_data_domain_dims = grid.dim_names[:-1]  # minus the data domain dimension
@@ -282,6 +290,10 @@ def populate_non_dim_coordinates(
 
         not_null = coord_trace_indices != grid.map.fill_value
         tmp_coord_values[not_null] = coord_values[coord_trace_indices[not_null]]
+
+        if coord_name in SCALE_COORDINATE_KEYS:
+            tmp_coord_values = _apply_coordinate_scalar(tmp_coord_values, horizontal_coordinate_scalar)
+
         dataset[coord_name][:] = tmp_coord_values
         drop_vars_delayed.append(coord_name)
 
@@ -291,16 +303,22 @@ def populate_non_dim_coordinates(
     return dataset, drop_vars_delayed
 
 
-def _get_horizontal_coordinate_unit(segy_headers: list[Dimension]) -> LengthUnitModel | None:
+def _get_horizontal_coordinate_unit(segy_info: SegyFileHeaderDump) -> LengthUnitModel | None:
     """Get the coordinate unit from the SEG-Y headers."""
-    name = TraceHeaderFieldsRev0.COORDINATE_UNIT.name.upper()
-    unit_hdr = next((c for c in segy_headers if c.name.upper() == name), None)
-    if unit_hdr is None or len(unit_hdr.coords) == 0:
+    measurement_system_code = int(segy_info.binary_header_dict[MEASUREMENT_SYSTEM_KEY])
+
+    if measurement_system_code not in (1, 2):
+        logger.warning(
+            "Unexpected value in coordinate unit (%s) header: %s. Can't extract coordinate unit and will "
+            "ingest without coordinate units.",
+            MEASUREMENT_SYSTEM_KEY,
+            measurement_system_code,
+        )
         return None
 
-    if segy_MeasurementSystem(unit_hdr.coords[0]) == segy_MeasurementSystem.METERS:
+    if measurement_system_code == SegyMeasurementSystem.METERS:
         unit = LengthUnitEnum.METER
-    if segy_MeasurementSystem(unit_hdr.coords[0]) == segy_MeasurementSystem.FEET:
+    if measurement_system_code == SegyMeasurementSystem.FEET:
         unit = LengthUnitEnum.FOOT
 
     return LengthUnitModel(length=unit)
@@ -310,6 +328,7 @@ def _populate_coordinates(
     dataset: xr_Dataset,
     grid: Grid,
     coords: dict[str, SegyHeaderArray],
+    horizontal_coordinate_scalar: int,
 ) -> tuple[xr_Dataset, list[str]]:
     """Populate dim and non-dim coordinates in the xarray dataset and write to Zarr.
 
@@ -319,6 +338,7 @@ def _populate_coordinates(
         dataset: The xarray dataset to populate.
         grid: The grid object containing the grid map.
         coords: The non-dim coordinates to populate.
+        horizontal_coordinate_scalar: The X/Y coordinate scalar from the SEG-Y file.
 
     Returns:
         Xarray dataset with filled coordinates and updated variables to drop after writing
@@ -329,7 +349,11 @@ def _populate_coordinates(
 
     # Populate the non-dimension coordinate variables (N-dim arrays)
     dataset, vars_to_drop_later = populate_non_dim_coordinates(
-        dataset, grid, coordinates=coords, drop_vars_delayed=drop_vars_delayed
+        dataset,
+        grid,
+        coordinates=coords,
+        drop_vars_delayed=drop_vars_delayed,
+        horizontal_coordinate_scalar=horizontal_coordinate_scalar,
     )
 
     return dataset, drop_vars_delayed
@@ -465,6 +489,7 @@ def segy_to_mdio(  # noqa PLR0913
     output_path: UPath | Path | str,
     overwrite: bool = False,
     grid_overrides: dict[str, Any] | None = None,
+    segy_header_overrides: SegyHeaderOverrides | None = None,
 ) -> None:
     """A function that converts a SEG-Y file to an MDIO v1 file.
 
@@ -477,6 +502,7 @@ def segy_to_mdio(  # noqa PLR0913
         output_path: The universal path for the output MDIO v1 file.
         overwrite: Whether to overwrite the output file if it already exists. Defaults to False.
         grid_overrides: Option to add grid overrides.
+        segy_header_overrides: Option to override specific SEG-Y headers during ingestion.
 
     Raises:
         FileExistsError: If the output location already exists and overwrite is False.
@@ -489,7 +515,12 @@ def segy_to_mdio(  # noqa PLR0913
         raise FileExistsError(err)
 
     segy_settings = SegyFileSettings(storage_options=input_path.storage_options)
-    segy_file = SegyFile(url=input_path.as_posix(), spec=segy_spec, settings=segy_settings)
+    segy_file = SegyFile(
+        url=input_path.as_posix(),
+        spec=segy_spec,
+        settings=segy_settings,
+        header_overrides=segy_header_overrides,
+    )
     segy_info: SegyFileHeaderDump = _get_segy_file_header_dump(segy_file)
 
     segy_dimensions, segy_headers = _scan_for_headers(segy_file, mdio_template, grid_overrides)
@@ -506,7 +537,7 @@ def segy_to_mdio(  # noqa PLR0913
             logger.warning("MDIO__IMPORT__RAW_HEADERS is experimental and expected to change or be removed.")
             mdio_template = _add_raw_headers_to_template(mdio_template)
 
-    horizontal_unit = _get_horizontal_coordinate_unit(segy_dimensions)
+    horizontal_unit = _get_horizontal_coordinate_unit(segy_info)
     mdio_ds: Dataset = mdio_template.build_dataset(
         name=mdio_template.name,
         sizes=grid.shape,
@@ -523,10 +554,12 @@ def segy_to_mdio(  # noqa PLR0913
 
     xr_dataset: xr_Dataset = to_xarray_dataset(mdio_ds=mdio_ds)
 
+    coordinate_scalar = _get_coordinate_scalar(segy_file)
     xr_dataset, drop_vars_delayed = _populate_coordinates(
         dataset=xr_dataset,
         grid=grid,
         coords=non_dim_coords,
+        horizontal_coordinate_scalar=coordinate_scalar,
     )
 
     xr_dataset = _add_segy_file_headers(xr_dataset, segy_info)

src/mdio/converters/segy.py

@@ -5,17 +5,16 @@
 import os
 from typing import TYPE_CHECKING
 from typing import TypedDict
-from typing import cast
 
 import numpy as np
 from segy import SegyFile
+from segy.arrays import HeaderArray
 
 from mdio.api.io import to_mdio
 from mdio.builder.schemas.dtype import ScalarType
 from mdio.segy._raw_trace_wrapper import SegyFileRawTraceWrapper
 
 if TYPE_CHECKING:
-    from segy.arrays import HeaderArray
     from segy.config import SegyFileSettings
     from segy.schema import SegySpec
     from upath import UPath
@@ -79,7 +78,7 @@ def header_scan_worker(
     # (singleton) so we can concat and assign stuff later.
     trace_header = np.array(trace_header, dtype=new_dtype, ndmin=1)
 
-    return cast("HeaderArray", trace_header)
+    return HeaderArray(trace_header)  # wrap back so we can use aliases
 
 
 def trace_worker(  # noqa: PLR0913

src/mdio/segy/_workers.py

@@ -91,6 +91,7 @@ def to_zarr(  # noqa: PLR0913, PLR0915
         "url": segy_file.fs.unstrip_protocol(segy_file.url),
         "spec": segy_file.spec,
         "settings": segy_file.settings,
+        "header_overrides": segy_file.header_overrides,
     }
     with executor:
         futures = []

src/mdio/segy/blocked_io.py

@@ -24,7 +24,7 @@
 
 def parse_headers(
     segy_file: SegyFile,
-    subset: list[str] | None = None,
+    subset: tuple[str, ...] | None = None,
     block_size: int = 10000,
     progress_bar: bool = True,
 ) -> HeaderArray:

src/mdio/segy/parsers.py

@@ -0,0 +1,59 @@
+"""Utilities to read, parse, and apply coordinate scalars."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from segy.standards import SegyStandard
+from segy.standards.fields import trace as trace_header_fields
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+    from segy import SegyFile
+
+
+logger = logging.getLogger(__name__)
+
+
+COORD_SCALAR_KEY = trace_header_fields.Rev0.COORDINATE_SCALAR.model.name
+VALID_COORD_SCALAR = {1, 10, 100, 1000, 10000}
+SCALE_COORDINATE_KEYS = [
+    "cdp_x",
+    "cdp_y",
+    "source_coord_x",
+    "source_coord_y",
+    "group_coord_x",
+    "group_coord_y",
+]
+
+
+def _get_coordinate_scalar(segy_file: SegyFile) -> int:
+    """Get and parse the coordinate scalar from the first SEG-Y trace header."""
+    file_revision = segy_file.spec.segy_standard
+    first_header = segy_file.header[0]
+    coord_scalar = int(first_header[COORD_SCALAR_KEY])
+
+    # Per Rev2, standardize 0 to 1 if a file is 2+.
+    if coord_scalar == 0 and file_revision >= SegyStandard.REV2:
+        logger.warning("Coordinate scalar is 0 and file is %s. Setting to 1.", file_revision)
+        return 1
+
+    def validate_segy_scalar(scalar: int) -> bool:
+        """Validate if coord scalar matches the seg-y standard."""
+        logger.debug("Coordinate scalar is %s", scalar)
+        return abs(scalar) in VALID_COORD_SCALAR  # valid values
+
+    is_valid = validate_segy_scalar(coord_scalar)
+    if not is_valid:
+        msg = f"Invalid coordinate scalar: {coord_scalar} for file revision {file_revision}."
+        raise ValueError(msg)
+
+    logger.info("Coordinate scalar is parsed as %s", coord_scalar)
+    return coord_scalar
+
+
+def _apply_coordinate_scalar(data: NDArray, scalar: int) -> NDArray:
+    if scalar < 0:
+        scalar = 1 / scalar
+    return data * abs(scalar)

src/mdio/segy/scalar.py

@@ -0,0 +1,92 @@
+"""Tests for coordinate scalar getters and apply functions."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+from segy import SegyFile
+from segy.standards import SegyStandard
+from segy.standards.fields import trace as trace_header_fields
+
+from mdio.segy.scalar import _apply_coordinate_scalar
+from mdio.segy.scalar import _get_coordinate_scalar
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+COORD_SCALAR_KEY = trace_header_fields.Rev0.COORDINATE_SCALAR.model.name
+
+
+@pytest.fixture
+def mock_segy_file() -> SegyFile:
+    """Mock SegyFile object."""
+    segy_file = MagicMock(spec=SegyFile)
+    segy_file.spec = MagicMock()
+    segy_file.header = [MagicMock()]
+    return segy_file
+
+
+@pytest.mark.parametrize("scalar", [1, 100, 10000, -10, -1000])
+def test_get_coordinate_scalar_valid(mock_segy_file: SegyFile, scalar: int) -> None:
+    """Test valid options when getting coordinate scalar."""
+    mock_segy_file.spec.segy_standard = SegyStandard.REV1
+    mock_segy_file.header[0].__getitem__.return_value = scalar
+
+    result = _get_coordinate_scalar(mock_segy_file)
+
+    assert result == scalar
+
+
+@pytest.mark.parametrize(
+    "revision",
+    [SegyStandard.REV2, SegyStandard.REV21],
+)
+def test_get_coordinate_scalar_zero_rev2_plus(mock_segy_file: SegyFile, revision: SegyStandard) -> None:
+    """Test when scalar is normalized to 1 (from 0) in Rev2+."""
+    mock_segy_file.spec.segy_standard = revision
+    mock_segy_file.header[0].__getitem__.return_value = 0
+
+    result = _get_coordinate_scalar(mock_segy_file)
+
+    assert result == 1
+
+
+@pytest.mark.parametrize(
+    ("scalar", "revision", "error_msg"),
+    [
+        (0, SegyStandard.REV0, "Invalid coordinate scalar: 0 for file revision SegyStandard.REV0."),
+        (110, SegyStandard.REV1, "Invalid coordinate scalar: 110 for file revision SegyStandard.REV1."),
+        (32768, SegyStandard.REV1, "Invalid coordinate scalar: 32768 for file revision SegyStandard.REV1."),
+    ],
+)
+def test_get_coordinate_scalar_invalid(
+    mock_segy_file: SegyFile, scalar: int, revision: SegyStandard, error_msg: str
+) -> None:
+    """Test invalid options when getting coordinate scalar."""
+    mock_segy_file.spec.segy_standard = revision
+    mock_segy_file.header[0].__getitem__.return_value = scalar
+
+    with pytest.raises(ValueError, match=error_msg):
+        _get_coordinate_scalar(mock_segy_file)
+
+
+@pytest.mark.parametrize(
+    ("data", "scalar", "expected"),
+    [
+        # POSITIVE
+        (np.array([1, 2, 3]), 1, np.array([1, 2, 3])),
+        (np.array([1, 2, 3]), 10, np.array([10, 20, 30])),
+        (np.array([[1, 2], [3, 4]]), 1000, np.array([[1000, 2000], [3000, 4000]])),
+        # NEGATIVE
+        (np.array([1, 2, 3]), -1, np.array([1, 2, 3])),
+        (np.array([10, 20, 30]), -10, np.array([1, 2, 3])),
+        (np.array([[1000, 2000], [3000, 4000]]), -1000, np.array([[1, 2], [3, 4]])),
+    ],
+)
+def test_apply_coordinate_scalar(data: NDArray, scalar: int, expected: NDArray) -> None:
+    """Test applying coordinate scalar with negative and positive code."""
+    result = _apply_coordinate_scalar(data, scalar)
+    assert np.allclose(result, expected)