Move SEG-Y metadata away from Dataset root (#667)

* remove old segy revision encoder and rename function

* use new rev encoder and refactor segy creation during export

* add segy related metadata to a separate variable for cleaner dataset metadata

* update tests

---------

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

Author: tasansal

src/mdio/converters/segy.py

@@ -293,41 +293,40 @@ def _populate_coordinates(
     return dataset, drop_vars_delayed
 
 
-def _add_segy_ingest_attributes(dataset: Dataset, segy_file: SegyFile, grid_overrides: dict[str, Any] | None) -> None:
-    text_header = segy_file.text_header.splitlines()
-    # Validate:
-    # text_header this should be a 40-items array of strings with width of 80 characters.
-    item_count = 40
-    if len(text_header) != item_count:
-        err = f"Invalid text header count: expected {item_count}, got {len(text_header)}"
+def _add_segy_file_headers(xr_dataset: xr_Dataset, segy_file: SegyFile) -> xr_Dataset:
+    expected_rows = 40
+    expected_cols = 80
+
+    text_header = segy_file.text_header
+    text_header_rows = 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)}"
         raise ValueError(err)
-    char_count = 80
-    for i, line in enumerate(text_header):
-        if len(line) != char_count:
-            err = f"Invalid text header {i} line length: expected {char_count}, got {len(line)}"
-            raise ValueError(err)
-    ext_text_header = segy_file.ext_text_header
 
-    # If using SegyFile.ext_text_header this should be a minimum of 40 elements and must
-    # capture all textual information (ensure text_header is a subset of ext_text_header).
-    if ext_text_header is not None:
-        for ext_hdr in ext_text_header:
-            text_header.append(ext_hdr.splitlines())
+    if any(text_header_cols_bad):
+        err = f"Invalid text header columns: expected {expected_cols} per line."
+        raise ValueError(err)
+
+    xr_dataset["segy_file_header"] = ((), "")
+    xr_dataset["segy_file_header"].attrs.update(
+        {
+            "textHeader": text_header,
+            "binaryHeader": segy_file.binary_header.to_dict(),
+        }
+    )
+
+    return xr_dataset
+
 
-    # Handle case where it may not have any metadata yet
+def _add_grid_override_to_metadata(dataset: Dataset, grid_overrides: dict[str, Any] | None) -> None:
+    """Add grid override to Dataset metadata if needed."""
     if dataset.metadata.attributes is None:
         dataset.metadata.attributes = {}
 
-    segy_attributes = {
-        "textHeader": text_header,
-        "binaryHeader": segy_file.binary_header.to_dict(),
-    }
-
     if grid_overrides is not None:
-        segy_attributes["gridOverrides"] = grid_overrides
-
-    # Update the attributes with the text and binary headers.
-    dataset.metadata.attributes.update(segy_attributes)
+        dataset.metadata.attributes["gridOverrides"] = grid_overrides
 
 
 def segy_to_mdio(  # noqa PLR0913
@@ -377,7 +376,7 @@ def segy_to_mdio(  # noqa PLR0913
         header_dtype=header_dtype,
     )
 
-    _add_segy_ingest_attributes(dataset=mdio_ds, segy_file=segy_file, grid_overrides=grid_overrides)
+    _add_grid_override_to_metadata(dataset=mdio_ds, grid_overrides=grid_overrides)
 
     xr_dataset: xr_Dataset = to_xarray_dataset(mdio_ds=mdio_ds)
 
@@ -387,8 +386,9 @@ def segy_to_mdio(  # noqa PLR0913
         coords=non_dim_coords,
     )
 
-    xr_dataset.trace_mask.data[:] = grid.live_mask
+    xr_dataset = _add_segy_file_headers(xr_dataset, segy_file)
 
+    xr_dataset.trace_mask.data[:] = grid.live_mask
     # IMPORTANT: Do not drop the "trace_mask" here, as it will be used later in
     # blocked_io.to_zarr() -> _workers.trace_worker()
 

src/mdio/segy/compat.py

@@ -95,29 +95,22 @@ def mdio_segy_spec(version_str: str | None = None) -> SegySpec:
     )
 
 
-def revision_encode(binary_header: dict, version_str: str) -> dict:
+def encode_segy_revision(binary_header: dict) -> dict:
     """Encode revision code to binary header.
 
-    We have two cases where legacy MDIO uses keys "SEGYRevision" and "SEGYRevisionMinor" whereas
-    the new one uses "segy_revision_major" and "segy_revision_minor". Given either case we return
-    the correctly Rev1 like encoded revision code, ready to write to SEG-Y.
+    Return the correctly Rev1-like encoded revision code, ready to write to SEG-Y.
 
     Args:
         binary_header: Dictionary representing the SEG-Y binary header. Contains keys for major
             and minor revision numbers.
-        version_str: MDIO version string to determine the encoding format.
 
     Returns:
         The updated binary header with the encoded revision.
 
     Raises:
         InvalidMDIOError: Raised when binary header in MDIO is broken.
     """
-    version_obj = version.parse(version_str)
-    if version_obj > version.parse("0.7.4"):
-        major_key, minor_key = "segy_revision_major", "segy_revision_minor"
-    else:  # MDIO <0.8
-        major_key, minor_key = "SEGYRevision", "SEGYRevisionMinor"
+    major_key, minor_key = "segy_revision_major", "segy_revision_minor"
 
     try:
         major = binary_header.pop(major_key)

src/mdio/segy/creation.py

@@ -13,7 +13,7 @@
 from tqdm.auto import tqdm
 
 from mdio.api.io import open_mdio
-from mdio.segy.compat import revision_encode
+from mdio.segy.compat import encode_segy_revision
 
 if TYPE_CHECKING:
     import xarray as xr
@@ -25,9 +25,8 @@
 logger = logging.getLogger(__name__)
 
 
-def make_segy_factory(dataset: xr.Dataset, spec: SegySpec) -> SegyFactory:
+def make_segy_factory(spec: SegySpec, binary_header: dict[str, int]) -> SegyFactory:
     """Generate SEG-Y factory from MDIO metadata."""
-    binary_header = dataset.attrs["attributes"]["binaryHeader"]
     sample_interval = binary_header["sample_interval"]
     samples_per_trace = binary_header["samples_per_trace"]
     return SegyFactory(
@@ -63,22 +62,20 @@ def mdio_spec_to_segy(
         Opened Xarray Dataset for MDIO file and SegyFactory
     """
     dataset = open_mdio(input_path, chunks=new_chunks)
-    factory = make_segy_factory(dataset, spec=segy_spec)
 
-    attr = dataset.attrs["attributes"]
+    file_header = dataset["segy_file_header"]
+    text_header = file_header.attrs["textHeader"]
+    binary_header = file_header.attrs["binaryHeader"]
+    binary_header = encode_segy_revision(binary_header)
 
-    txt_header = attr["textHeader"]
-    text_str = "\n".join(txt_header)
-    text_bytes = factory.create_textual_header(text_str)
+    factory = make_segy_factory(spec=segy_spec, binary_header=binary_header)
 
-    bin_header = attr["binaryHeader"]
-    mdio_file_version = dataset.attrs["apiVersion"]
-    binary_header = revision_encode(bin_header, mdio_file_version)
-    bin_hdr_bytes = factory.create_binary_header(binary_header)
+    text_header_bytes = factory.create_textual_header(text_header)
+    binary_header_bytes = factory.create_binary_header(binary_header)
 
     with output_path.open(mode="wb") as fp:
-        fp.write(text_bytes)
-        fp.write(bin_hdr_bytes)
+        fp.write(text_header_bytes)
+        fp.write(binary_header_bytes)
 
     return dataset, factory
 

tests/integration/test_segy_import_export.py

@@ -73,7 +73,7 @@ def test_import_4d_segy(  # noqa: PLR0913
 
         ds = open_mdio(zarr_tmp)
 
-        assert ds.attrs["attributes"]["binaryHeader"]["samples_per_trace"] == num_samples
+        assert ds["segy_file_header"].attrs["binaryHeader"]["samples_per_trace"] == num_samples
         assert ds.attrs["attributes"]["gridOverrides"] == grid_override
 
         assert npt.assert_array_equal(ds["shot_point"], shots)
@@ -120,7 +120,7 @@ def test_import_4d_segy(  # noqa: PLR0913
 
         ds = open_mdio(zarr_tmp)
 
-        assert ds.attrs["attributes"]["binaryHeader"]["samples_per_trace"] == num_samples
+        assert ds["segy_file_header"].attrs["binaryHeader"]["samples_per_trace"] == num_samples
         assert ds.attrs["attributes"].get("gridOverrides", None) == grid_override  # may not exist, so default=None
 
         xrt.assert_duckarray_equal(ds["shot_point"], shots)
@@ -261,13 +261,15 @@ def test_dataset_metadata(self, zarr_tmp: Path) -> None:
 
         attributes = ds.attrs["attributes"]
         assert attributes is not None
-        assert len(attributes) == 5
+        assert len(attributes) == 3
         # Validate all attributes provided by the abstract template
         assert attributes["defaultVariableName"] == "amplitude"
         assert attributes["surveyType"] == "3D"
         assert attributes["gatherType"] == "stacked"
-        assert attributes["textHeader"] == text_header_teapot_dome()
-        assert attributes["binaryHeader"] == binary_header_teapot_dome()
+
+        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()
 
     def test_variable_metadata(self, zarr_tmp: Path) -> None:
         """Metadata reading tests."""

tests/integration/testing_data.py

@@ -34,9 +34,9 @@ def custom_teapot_dome_segy_spec(keep_unaltered: bool) -> SegySpec:
     )
 
 
-def text_header_teapot_dome() -> list[str]:
+def text_header_teapot_dome() -> str:
     """Return the teapot dome expected text header."""
-    return [
+    header_rows = [
         "C 1 CLIENT: ROCKY MOUNTAIN OILFIELD TESTING CENTER                              ",
         "C 2 PROJECT: NAVAL PETROLEUM RESERVE #3 (TEAPOT DOME); NATRONA COUNTY, WYOMING  ",
         "C 3 LINE: 3D                                                                    ",
@@ -78,6 +78,7 @@ def text_header_teapot_dome() -> list[str]:
         "C39               (voice) 303.694.9629 (fax) 303.771.1646                       ",
         "C40 END EBCDIC                                                                  ",
     ]
+    return "\n".join(header_rows)
 
 
 def binary_header_teapot_dome() -> dict[str, int]: