Provide clean disaster recovery interface

Author: BrianMichell

src/mdio/segy/_disaster_recovery_wrapper.py

@@ -4,73 +4,32 @@
 
 from typing import TYPE_CHECKING
 
-from segy.schema import Endianness
-from segy.transforms import ByteSwapTransform
-from segy.transforms import IbmFloatTransform
+from copy import deepcopy
+import numpy as np
 
 if TYPE_CHECKING:
     from numpy.typing import NDArray
     from segy import SegyFile
-    from segy.transforms import Transform
-    from segy.transforms import TransformPipeline
 
+class SegyFileTraceDataWrapper:
 
-def _reverse_single_transform(data: NDArray, transform: Transform, endianness: Endianness) -> NDArray:
-    """Reverse a single transform operation."""
-    if isinstance(transform, ByteSwapTransform):
-        # Reverse the endianness conversion
-        if endianness == Endianness.LITTLE:
-            return data
+    def __init__(self, segy_file: SegyFile, indices: int | list[int] | NDArray | slice):
+        self.segy_file = segy_file
+        self.indices = indices
+        self._header_pipeline = deepcopy(segy_file.accessors.header_decode_pipeline)
+        segy_file.accessors.header_decode_pipeline.transforms = []
+        self.traces = segy_file.trace[indices]
 
-        reverse_transform = ByteSwapTransform(Endianness.BIG)
-        return reverse_transform.apply(data)
+    @property
+    def header(self):
+        # The copy is necessary to avoid applying the pipeline to the original header.
+        return self._header_pipeline.apply(self.traces.header.copy())
 
-    # TODO(BrianMichell): #0000 Do we actually need to worry about IBM/IEEE transforms here?
-    if isinstance(transform, IbmFloatTransform):
-        # Reverse IBM float conversion
-        reverse_direction = "to_ibm" if transform.direction == "to_ieee" else "to_ieee"
-        reverse_transform = IbmFloatTransform(reverse_direction, transform.keys)
-        return reverse_transform.apply(data)
 
-    # For unknown transforms, return data unchanged
-    return data
+    @property
+    def raw_header(self):
+        return np.ascontiguousarray(self.traces.header).view("|V240")
 
-
-def get_header_raw_and_transformed(
-    segy_file: SegyFile, indices: int | list[int] | NDArray | slice, do_reverse_transforms: bool = True
-) -> tuple[NDArray | None, NDArray, NDArray]:
-    """Get both raw and transformed header data.
-
-    Args:
-        segy_file: The SegyFile instance
-        indices: Which headers to retrieve
-        do_reverse_transforms: Whether to apply the reverse transform to get raw data
-
-    Returns:
-        Tuple of (raw_headers, transformed_headers, traces)
-    """
-    traces = segy_file.trace[indices]
-    transformed_headers = traces.header
-
-    # Reverse transforms to get raw data
-    if do_reverse_transforms:
-        raw_headers = _reverse_transforms(
-            transformed_headers, segy_file.header.transform_pipeline, segy_file.spec.endianness
-        )
-    else:
-        raw_headers = None
-
-    return raw_headers, transformed_headers, traces
-
-
-def _reverse_transforms(
-    transformed_data: NDArray, transform_pipeline: TransformPipeline, endianness: Endianness
-) -> NDArray:
-    """Reverse the transform pipeline to get raw data."""
-    raw_data = transformed_data.copy() if hasattr(transformed_data, "copy") else transformed_data
-
-    # Apply transforms in reverse order
-    for transform in reversed(transform_pipeline.transforms):
-        raw_data = _reverse_single_transform(raw_data, transform, endianness)
-
-    return raw_data
+    @property
+    def sample(self):
+        return self.traces.sample

src/mdio/segy/_workers.py

@@ -12,7 +12,7 @@
 
 from mdio.api.io import to_mdio
 from mdio.builder.schemas.dtype import ScalarType
-from mdio.segy._disaster_recovery_wrapper import get_header_raw_and_transformed
+from mdio.segy._disaster_recovery_wrapper import SegyFileTraceDataWrapper
 
 if TYPE_CHECKING:
     from segy.arrays import HeaderArray
@@ -134,17 +134,18 @@ def trace_worker(  # noqa: PLR0913
     if raw_header_key in dataset.data_vars:
         worker_variables.append(raw_header_key)
 
-    from copy import deepcopy  # TODO: Move to head if we need to copy
-    header_pipeline = deepcopy(segy_file.accessors.header_decode_pipeline)
-    segy_file.accessors.header_decode_pipeline.transforms = []
-    traces = segy_file.trace[live_trace_indexes]
+    # traces = segy_file.trace[live_trace_indexes]
+    # Raw headers are not intended to remain as a feature of the SEGY ingestion.
+    # For that reason, we have wrapped the accessors to provide an interface that can be removed
+    # and not require additional changes to the below code.
+    # NOTE: The `raw_header_key` code block should be removed in full as it will become dead code.
+    traces = SegyFileTraceDataWrapper(segy_file, live_trace_indexes)
     ds_to_write = dataset[worker_variables]
 
     if header_key in worker_variables:
         # Create temporary array for headers with the correct shape
         tmp_headers = np.zeros_like(dataset[header_key])
-        # tmp_headers[not_null] = transformed_headers
-        tmp_headers[not_null] = header_pipeline.apply(traces.header.copy())
+        tmp_headers[not_null] = traces.header
         # Create a new Variable object to avoid copying the temporary array
         # The ideal solution is to use `ds_to_write[header_key][:] = tmp_headers`
         # but Xarray appears to be copying memory instead of doing direct assignment.
@@ -155,10 +156,9 @@ def trace_worker(  # noqa: PLR0913
             attrs=ds_to_write[header_key].attrs,
             encoding=ds_to_write[header_key].encoding,  # Not strictly necessary, but safer than not doing it.
         )
-    # del transformed_headers  # Manage memory
     if raw_header_key in worker_variables:
         tmp_raw_headers = np.zeros_like(dataset[raw_header_key])
-        tmp_raw_headers[not_null] = np.ascontiguousarray(traces.header).view("|V240")
+        tmp_raw_headers[not_null] = traces.raw_header
 
         ds_to_write[raw_header_key] = Variable(
             ds_to_write[raw_header_key].dims,