Environment variable flag to preserve unmodified, raw binary header from SEG-Y in segy_file_header variable. (#683)

* binary_header_disaster_recovery 2

* Update segy.py

Make raw binary header serialization conditional to match with pending raw trace headers PR

* Update test_segy_roundtrip_teapot.py

Set the raw headers environment variable

* refactor a little

* monkey patch env vars for teapot

---------

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

Author: 3 people

src/mdio/converters/segy.py

@@ -2,8 +2,10 @@
 
 from __future__ import annotations
 
+import base64
 import logging
 import os
+from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -123,6 +125,35 @@ def grid_density_qc(grid: Grid, num_traces: int) -> None:
         raise GridTraceSparsityError(grid.shape, num_traces, msg)
 
 
+@dataclass
+class SegyFileHeaderDump:
+    """Segy metadata information."""
+
+    text_header: str
+    binary_header_dict: dict
+    raw_binary_headers: bytes
+
+
+def _get_segy_file_header_dump(segy_file: SegyFile) -> SegyFileHeaderDump:
+    """Reads information from a SEG-Y file."""
+    text_header = segy_file.text_header
+
+    raw_binary_headers: bytes = segy_file.fs.read_block(
+        fn=segy_file.url,
+        offset=segy_file.spec.binary_header.offset,
+        length=segy_file.spec.binary_header.itemsize,
+    )
+
+    # We read here twice, but it's ok for now. Only 400-bytes.
+    binary_header_dict = segy_file.binary_header.to_dict()
+
+    return SegyFileHeaderDump(
+        text_header=text_header,
+        binary_header_dict=binary_header_dict,
+        raw_binary_headers=raw_binary_headers,
+    )
+
+
 def _scan_for_headers(
     segy_file: SegyFile,
     template: AbstractDatasetTemplate,
@@ -151,12 +182,12 @@ def _scan_for_headers(
     return segy_dimensions, segy_headers
 
 
-def _build_and_check_grid(segy_dimensions: list[Dimension], segy_file: SegyFile, segy_headers: SegyHeaderArray) -> Grid:
+def _build_and_check_grid(segy_dimensions: list[Dimension], num_traces: int, segy_headers: SegyHeaderArray) -> Grid:
     """Build and check the grid from the SEG-Y headers and dimensions.
 
     Args:
         segy_dimensions: List of of all SEG-Y dimensions to build grid from.
-        segy_file: Instance of SegyFile to check for trace count.
+        num_traces: Number of traces in the SEG-Y file.
         segy_headers: Headers read in from SEG-Y file for building the trace map.
 
     Returns:
@@ -166,15 +197,15 @@ def _build_and_check_grid(segy_dimensions: list[Dimension], segy_file: SegyFile,
         GridTraceCountError: If number of traces in SEG-Y file does not match the parsed grid
     """
     grid = Grid(dims=segy_dimensions)
-    grid_density_qc(grid, segy_file.num_traces)
+    grid_density_qc(grid, num_traces)
     grid.build_map(segy_headers)
     # Check grid validity by comparing trace numbers
-    if np.sum(grid.live_mask) != segy_file.num_traces:
+    if np.sum(grid.live_mask) != num_traces:
         for dim_name in grid.dim_names:
             dim_min, dim_max = grid.get_min(dim_name), grid.get_max(dim_name)
             logger.warning("%s min: %s max: %s", dim_name, dim_min, dim_max)
         logger.warning("Ingestion grid shape: %s.", grid.shape)
-        raise GridTraceCountError(np.sum(grid.live_mask), segy_file.num_traces)
+        raise GridTraceCountError(np.sum(grid.live_mask), num_traces)
     return grid
 
 
@@ -301,20 +332,19 @@ def _populate_coordinates(
     return dataset, drop_vars_delayed
 
 
-def _add_segy_file_headers(xr_dataset: xr_Dataset, segy_file: SegyFile) -> xr_Dataset:
+def _add_segy_file_headers(xr_dataset: xr_Dataset, segy_file_header_dump: SegyFileHeaderDump) -> xr_Dataset:
     save_file_header = os.getenv("MDIO__IMPORT__SAVE_SEGY_FILE_HEADER", "") in ("1", "true", "yes", "on")
     if not save_file_header:
         return xr_dataset
 
     expected_rows = 40
     expected_cols = 80
 
-    text_header = segy_file.text_header
-    text_header_rows = text_header.splitlines()
+    text_header_rows = segy_file_header_dump.text_header.splitlines()
     text_header_cols_bad = [len(row) != expected_cols for row in text_header_rows]
 
     if len(text_header_rows) != expected_rows:
-        err = f"Invalid text header count: expected {expected_rows}, got {len(text_header)}"
+        err = f"Invalid text header count: expected {expected_rows}, got {len(segy_file_header_dump.text_header)}"
         raise ValueError(err)
 
     if any(text_header_cols_bad):
@@ -324,10 +354,13 @@ def _add_segy_file_headers(xr_dataset: xr_Dataset, segy_file: SegyFile) -> xr_Da
     xr_dataset["segy_file_header"] = ((), "")
     xr_dataset["segy_file_header"].attrs.update(
         {
-            "textHeader": text_header,
-            "binaryHeader": segy_file.binary_header.to_dict(),
+            "textHeader": segy_file_header_dump.text_header,
+            "binaryHeader": segy_file_header_dump.binary_header_dict,
         }
     )
+    if os.getenv("MDIO__IMPORT__RAW_HEADERS") in ("1", "true", "yes", "on"):
+        raw_binary_base64 = base64.b64encode(segy_file_header_dump.raw_binary_headers).decode("ascii")
+        xr_dataset["segy_file_header"].attrs.update({"rawBinaryHeader": raw_binary_base64})
 
     return xr_dataset
 
@@ -428,10 +461,11 @@ def segy_to_mdio(  # noqa PLR0913
 
     segy_settings = SegySettings(storage_options=input_path.storage_options)
     segy_file = SegyFile(url=input_path.as_posix(), spec=segy_spec, settings=segy_settings)
+    segy_info: SegyFileHeaderDump = _get_segy_file_header_dump(segy_file)
 
     segy_dimensions, segy_headers = _scan_for_headers(segy_file, mdio_template, grid_overrides)
 
-    grid = _build_and_check_grid(segy_dimensions, segy_file, segy_headers)
+    grid = _build_and_check_grid(segy_dimensions, segy_file.num_traces, segy_headers)
 
     _, non_dim_coords = _get_coordinates(grid, segy_headers, mdio_template)
     header_dtype = to_structured_type(segy_spec.trace.header.dtype)
@@ -461,7 +495,7 @@ def segy_to_mdio(  # noqa PLR0913
         coords=non_dim_coords,
     )
 
-    xr_dataset = _add_segy_file_headers(xr_dataset, segy_file)
+    xr_dataset = _add_segy_file_headers(xr_dataset, segy_info)
 
     xr_dataset.trace_mask.data[:] = grid.live_mask
     # IMPORTANT: Do not drop the "trace_mask" here, as it will be used later in

tests/conftest.py

@@ -23,7 +23,9 @@
 def fake_segy_tmp(tmp_path_factory: pytest.TempPathFactory) -> Path:
     """Make a temp file for the fake SEG-Y files we are going to create."""
     if DEBUG_MODE:
-        return Path("TMP/fake_segy")
+        tmp_dir = Path("tmp/fake_segy")
+        tmp_dir.mkdir(parents=True, exist_ok=True)
+        return tmp_dir
     return tmp_path_factory.mktemp(r"fake_segy")
 
 
@@ -37,7 +39,7 @@ def segy_input_uri() -> str:
 def segy_input(segy_input_uri: str, tmp_path_factory: pytest.TempPathFactory) -> Path:
     """Download teapot dome dataset for testing."""
     if DEBUG_MODE:
-        tmp_dir = Path("TMP/segy")
+        tmp_dir = Path("tmp/segy")
         tmp_dir.mkdir(parents=True, exist_ok=True)
     else:
         tmp_dir = tmp_path_factory.mktemp("segy")
@@ -50,23 +52,23 @@ def segy_input(segy_input_uri: str, tmp_path_factory: pytest.TempPathFactory) ->
 def zarr_tmp(tmp_path_factory: pytest.TempPathFactory) -> Path:
     """Make a temp file for the output MDIO."""
     if DEBUG_MODE:
-        return Path("TMP/mdio")
+        return Path("tmp/mdio")
     return tmp_path_factory.mktemp(r"mdio")
 
 
 @pytest.fixture(scope="module")
 def zarr_tmp2(tmp_path_factory: pytest.TempPathFactory) -> Path:
     """Make a temp file for the output MDIO."""
     if DEBUG_MODE:
-        return Path("TMP/mdio2")
+        return Path("tmp/mdio2")
     return tmp_path_factory.mktemp(r"mdio2")
 
 
 @pytest.fixture(scope="session")
 def segy_export_tmp(tmp_path_factory: pytest.TempPathFactory) -> Path:
     """Make a temp file for the round-trip IBM SEG-Y."""
     if DEBUG_MODE:
-        tmp_dir = Path("TMP/segy")
+        tmp_dir = Path("tmp/segy")
         tmp_dir.mkdir(parents=True, exist_ok=True)
     else:
         tmp_dir = tmp_path_factory.mktemp("segy")

tests/integration/test_segy_import_export_masked.py

@@ -290,7 +290,7 @@ def export_masked_path(tmp_path_factory: pytest.TempPathFactory, raw_headers_env
     path_suffix = "with_raw_headers" if raw_headers_enabled else "without_raw_headers"
 
     if DEBUG_MODE:
-        return Path(f"TMP/export_masked_{path_suffix}")
+        return Path(f"tmp/export_masked_{path_suffix}")
     return tmp_path_factory.getbasetemp() / f"export_masked_{path_suffix}"
 
 

tests/integration/test_segy_roundtrip_teapot.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import json
-import os
 from typing import TYPE_CHECKING
 
 import dask
@@ -25,13 +24,20 @@
 from mdio.converters.segy import segy_to_mdio
 
 if TYPE_CHECKING:
+    from collections.abc import Generator
     from pathlib import Path
 
     from segy.schema import SegySpec
 
 
 dask.config.set(scheduler="synchronous")
-os.environ["MDIO__IMPORT__SAVE_SEGY_FILE_HEADER"] = "true"
+
+
+@pytest.fixture
+def set_env_vars(monkeypatch: Generator[pytest.MonkeyPatch]) -> None:
+    """Set environment variables for the Teapot dome tests."""
+    monkeypatch.setenv("MDIO__IMPORT__SAVE_SEGY_FILE_HEADER", "true")
+    monkeypatch.setenv("MDIO__IMPORT__RAW_HEADERS", "true")
 
 
 @pytest.fixture
@@ -130,11 +136,29 @@ def binary_header_teapot_dome() -> dict[str, int]:
     }
 
 
+def raw_binary_header_teapot_dome() -> str:
+    """Return the teapot dome expected raw binary header, base64 encoded."""
+    return (
+        "AAAnDwAAJw8AAAABALwAAAfQAAAF3QXdAAEAOQAEAAEAAAAAAAAAAAAAAAAAAAAAAAIAAQAEAAIAAQAAAAAAAAAAAAAAAAAAAAAA"
+        "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
+        "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
+        "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
+        "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
+        "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=="
+    )
+
+
 class TestTeapotRoundtrip:
     """Tests for Teapot Dome data ingestion and export."""
 
     @pytest.mark.dependency
-    def test_teapot_import(self, segy_input: Path, zarr_tmp: Path, teapot_segy_spec: SegySpec) -> None:
+    @pytest.mark.usefixtures("set_env_vars")
+    def test_teapot_import(
+        self,
+        segy_input: Path,
+        zarr_tmp: Path,
+        teapot_segy_spec: SegySpec,
+    ) -> None:
         """Test importing a SEG-Y file to MDIO.
 
         NOTE: This test must be executed before the 'TestReader' and 'TestExport' tests.
@@ -176,6 +200,7 @@ def test_dataset_metadata(self, zarr_tmp: Path) -> None:
         segy_file_header = ds["segy_file_header"]
         assert segy_file_header.attrs["textHeader"] == text_header_teapot_dome()
         assert segy_file_header.attrs["binaryHeader"] == binary_header_teapot_dome()
+        assert segy_file_header.attrs["rawBinaryHeader"] == raw_binary_header_teapot_dome()
 
     def test_variable_metadata(self, zarr_tmp: Path) -> None:
         """Metadata reading tests."""