Refactor revision handling to match new segy format (#494)

* Refactor SEG-Y compatibility with improved revision handling

* Add InvalidMDIOError exception class

* Rename function argument to match signature.

* Add error handling to docstring.

* Update pre-commit hook stages due to depracation

* string compat for <py3.12

* Add unit tests for MDIO compatibility with older versions

* Refactor MDIO version 0.7.4 metadata update logic.

Author: tasansal

.pre-commit-config.yaml

@@ -33,7 +33,7 @@ repos:
         entry: end-of-file-fixer
         language: system
         types: [text]
-        stages: [commit, push, manual]
+        stages: [pre-commit, pre-push, manual]
       - id: flake8
         name: flake8
         entry: flake8
@@ -60,7 +60,7 @@ repos:
         entry: trailing-whitespace-fixer
         language: system
         types: [text]
-        stages: [commit, push, manual]
+        stages: [pre-commit, pre-push, manual]
         args: [--markdown-linebreak-ext=md]
   - repo: https://github.com/pre-commit/mirrors-prettier
     rev: v2.6.0

src/mdio/exceptions.py

@@ -50,3 +50,7 @@ def __init__(self, message, name=None, expected=None):
             message = " - ".join([message, extras])
 
         super().__init__(message)
+
+
+class InvalidMDIOError(MDIOError):
+    """Raised when an invalid MDIO file is encountered."""

src/mdio/segy/compat.py

@@ -8,9 +8,11 @@
 
 from __future__ import annotations
 
+import logging
 import os
 from importlib import metadata
 
+from packaging import version
 from segy.alias.segyio import SEGYIO_BIN_FIELD_MAP
 from segy.alias.segyio import SEGYIO_TRACE_FIELD_MAP
 from segy.schema import HeaderField
@@ -20,17 +22,33 @@
 from segy.schema import TextHeaderSpec
 from segy.schema import TraceDataSpec
 from segy.schema import TraceSpec
+from segy.standards.fields import binary
+
+from mdio.exceptions import InvalidMDIOError
 
 
 MDIO_VERSION = metadata.version("multidimio")
 
 
+logger = logging.getLogger(__name__)
+
+
 def get_binary_fields() -> list[HeaderField]:
     """Generate binary header fields from equinor/segyio fields."""
-    return [field.model for field in SEGYIO_BIN_FIELD_MAP.values()]
+    revision_field = binary.Rev1.SEGY_REVISION.model
+    mdio_v0_bin_fields = []
+
+    # Replace min/max (rev2-ish) with rev1 like parsing.
+    # Ignore minor one, and add the revision as 4-byte.
+    for alias, field in SEGYIO_BIN_FIELD_MAP.items():
+        if alias == "SEGYRevision":
+            mdio_v0_bin_fields.append(revision_field)
+        elif alias != "SEGYRevisionMinor":
+            mdio_v0_bin_fields.append(field.model)
+    return mdio_v0_bin_fields
 
 
-def get_trace_fields(version: str) -> list[HeaderField]:
+def get_trace_fields(version_str: str) -> list[HeaderField]:
     """Generate trace header fields.
 
     This part allows us to configure custom rules for different MDIO versions.
@@ -45,28 +63,29 @@ def get_trace_fields(version: str) -> list[HeaderField]:
     * mdio>=0.8.0 adds an extra field to the end to fill the last 8 bytes
 
     Args:
-        version: MDIO version to generate the trace fields for.
+        version_str: MDIO version to generate the trace fields for.
 
     Returns:
         List of header fields for specified MDIO version trace header encoding.
     """
     trace_fields = [field.model for field in SEGYIO_TRACE_FIELD_MAP.values()]
-    if version > "0.7.4":
+    version_obj = version.parse(version_str)
+    if version_obj > version.parse("0.7.4"):
         trace_fields.append(HeaderField(name="unassigned", byte=233, format="int64"))
     return trace_fields
 
 
-def mdio_segy_spec(version: str | None = None) -> SegySpec:
+def mdio_segy_spec(version_str: str | None = None) -> SegySpec:
     """Get a SEG-Y encoding spec for MDIO based on version."""
     spec_override = os.getenv("MDIO__SEGY__SPEC")
 
     if spec_override is not None:
         return SegySpec.model_validate_json(spec_override)
 
-    version = MDIO_VERSION if version is None else version
+    version_str = MDIO_VERSION if version_str is None else version_str
 
     binary_fields = get_binary_fields()
-    trace_fields = get_trace_fields(version)
+    trace_fields = get_trace_fields(version_str)
 
     return SegySpec(
         segy_standard=None,
@@ -77,3 +96,43 @@ def mdio_segy_spec(version: str | None = None) -> SegySpec:
             data=TraceDataSpec(format=ScalarType.IBM32),  # placeholder
         ),
     )
+
+
+def revision_encode(binary_header: dict, version_str: str) -> 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.
+
+    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"
+
+    try:
+        major = binary_header.pop(major_key)
+        minor = binary_header.pop(minor_key)
+    except KeyError:
+        msg = "Missing revision keys from binary header."
+        logger.error(msg)
+        raise InvalidMDIOError(msg) from KeyError
+
+    code = (major << 8) | minor
+    code_hex = f"0x{code:04x}"
+    binary_header["segy_revision"] = code
+    logger.info(f"Encoded revision {major}.{minor} to {code=} ~ {code_hex}")
+    return binary_header

src/mdio/segy/creation.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import logging
 import os
 from os import path
 from shutil import copyfileobj
@@ -16,12 +17,16 @@
 
 from mdio.api.accessor import MDIOReader
 from mdio.segy.compat import mdio_segy_spec
+from mdio.segy.compat import revision_encode
 
 
 if TYPE_CHECKING:
     from numpy.typing import NDArray
 
 
+logger = logging.getLogger(__name__)
+
+
 def make_segy_factory(
     mdio: MDIOReader,
     spec: SegySpec,
@@ -91,7 +96,9 @@ def mdio_spec_to_segy(
 
     text_str = "\n".join(mdio.text_header)
     text_bytes = factory.create_textual_header(text_str)
-    bin_hdr_bytes = factory.create_binary_header(mdio.binary_header)
+
+    binary_header = revision_encode(mdio.binary_header, mdio_file_version)
+    bin_hdr_bytes = factory.create_binary_header(binary_header)
 
     with open(output_segy_path, mode="wb") as fp:
         fp.write(text_bytes)

tests/unit/test_compat.py

@@ -0,0 +1,108 @@
+"""Test MDIO compatibility with older versions."""
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+import zarr
+from segy import SegyFile
+from segy.factory import SegyFactory
+from segy.standards import get_segy_standard
+
+from mdio import mdio_to_segy
+from mdio import segy_to_mdio
+
+
+# Constants
+MDIO_VERSIONS = ["0.7.4", "0.8.3"]
+SEGY_REVISIONS = [0.0, 0.1, 1.0, 1.1]
+INLINES = (10, 10, 11, 11)
+CROSSLINES = (100, 101, 100, 101)
+INDEX_BYTES = (189, 193)
+API_VERSION_KEY = "api_version"
+BINARY_HEADER_KEY = "binary_header"
+CHUNKED_TRACE_HEADERS_KEY = "chunked_012_trace_headers"
+
+
+def update_mdio_for_version_0_7_4(root):
+    """Update MDIO metadata to mimic version 0.7.4."""
+    # Update binary header revision keys
+    bin_hdr = root.metadata.attrs[BINARY_HEADER_KEY]
+    bin_hdr["SEGYRevision"] = bin_hdr.pop("segy_revision_major")
+    bin_hdr["SEGYRevisionMinor"] = bin_hdr.pop("segy_revision_minor")
+    root.metadata.attrs[BINARY_HEADER_KEY] = bin_hdr
+
+    # Remove trace headers past field 232 (pre-0.8 schema)
+    orig_hdr = root.metadata[CHUNKED_TRACE_HEADERS_KEY]
+    new_dtype = np.dtype(orig_hdr.dtype.descr[:-1])
+    new_hdr = zarr.zeros_like(orig_hdr, dtype=new_dtype)
+    root.metadata.create_dataset(
+        CHUNKED_TRACE_HEADERS_KEY,
+        data=new_hdr,
+        overwrite=True,
+    )
+    zarr.consolidate_metadata(root.store)
+
+
+@pytest.mark.parametrize("mdio_version", MDIO_VERSIONS)
+@pytest.mark.parametrize("segy_revision", SEGY_REVISIONS)
+def test_revision_encode_decode(
+    mdio_version: str, segy_revision: float, tmp_path: Path
+) -> None:
+    """Test binary header major/minor revision roundtrip.
+
+    After introducting TGSAI/segy, we changed the header names. Now we use
+    aliasing and MDIO has a dummy schema. The handling is slightly different
+    for SEG-Y revision major/minor numbers. Testing to ensure they're
+    (de)serialized correctly.
+    """
+    rev1_spec = get_segy_standard(1.0)
+    segy_filename = tmp_path / "segy_input.sgy"
+    mdio_output_filename = tmp_path / "output.mdio"
+    roundtrip_sgy_filename = tmp_path / "roundtrip_output.sgy"
+
+    # Make a rev1 segy
+    factory = SegyFactory(rev1_spec, sample_interval=1000, samples_per_trace=5)
+
+    # We will replace the values in revision fields with these
+    minor, major = np.modf(segy_revision)
+    major, minor = int(major), int(minor * 10)
+    revision_code = (major << 8) | minor
+
+    # Make fake tiny 3D dataset
+    txt_buffer = factory.create_textual_header()
+
+    header = factory.create_trace_header_template(len(INLINES))
+    data = factory.create_trace_sample_template(len(INLINES))
+    header["inline"] = INLINES
+    header["crossline"] = CROSSLINES
+    data[:] = np.arange(len(INLINES))[:, None]
+    trace_buffer = factory.create_traces(header, data)
+
+    # Update revision during bin hdr creation
+    bin_hdr_buffer = factory.create_binary_header(
+        update={"segy_revision": revision_code}
+    )
+    with open(segy_filename, mode="wb") as fp:
+        fp.write(txt_buffer)
+        fp.write(bin_hdr_buffer)
+        fp.write(trace_buffer)
+
+    # Convert SEG-Y to MDIO
+    segy_to_mdio(str(segy_filename), str(mdio_output_filename), index_bytes=INDEX_BYTES)
+
+    # Modify MDIO for specific versions
+    root = zarr.open_group(mdio_output_filename, mode="r+")
+    root.attrs[API_VERSION_KEY] = mdio_version
+    if mdio_version == "0.7.4":
+        update_mdio_for_version_0_7_4(root)
+
+    # Convert MDIO back to SEG-Y
+    mdio_to_segy(str(mdio_output_filename), str(roundtrip_sgy_filename))
+
+    # Assert binary headers and revisions match
+    orig = SegyFile(segy_filename, spec=rev1_spec)
+    rt = SegyFile(roundtrip_sgy_filename, spec=rev1_spec)
+    assert orig.binary_header["segy_revision_major"] == major
+    assert orig.binary_header["segy_revision_minor"] == minor
+    assert orig.binary_header == rt.binary_header