Further cleanup and optimization

Author: BrianMichell

src/mdio/segy/_disaster_recovery_wrapper.py

@@ -37,8 +37,9 @@ def _reverse_single_transform(data: NDArray, transform: Transform) -> NDArray:
 
 def get_header_raw_and_transformed(
     segy_file: SegyFile,
-    indices: int | list[int] | NDArray | slice
-) -> tuple[NDArray, NDArray, NDArray]:
+    indices: int | list[int] | NDArray | slice,
+    do_reverse_transforms: bool = True
+) -> tuple[NDArray | None, NDArray, NDArray]:
     """Get both raw and transformed header data.
 
     Args:
@@ -53,7 +54,10 @@ def get_header_raw_and_transformed(
     transformed_headers = traces.header
 
     # Reverse transforms to get raw data
-    raw_headers = _reverse_transforms(transformed_headers, segy_file.header.transform_pipeline)
+    if do_reverse_transforms:
+        raw_headers = _reverse_transforms(transformed_headers, segy_file.header.transform_pipeline)
+    else:
+        raw_headers = None
 
     return raw_headers, transformed_headers, traces
 

src/mdio/segy/_workers.py

@@ -9,7 +9,6 @@
 
 import numpy as np
 from segy import SegyFile
-from segy.indexing import merge_cat_file
 
 from mdio.api.io import to_mdio
 from mdio.builder.schemas.dtype import ScalarType
@@ -121,19 +120,23 @@ def trace_worker(  # noqa: PLR0913
     zarr_config.set({"threading.max_workers": 1})
 
     live_trace_indexes = local_grid_map[not_null].tolist()
-    raw_headers, transformed_headers, traces = get_header_raw_and_transformed(segy_file, live_trace_indexes)
 
     header_key = "headers"
     raw_header_key = "raw_headers"
 
+    # Used to disable the reverse transforms if we aren't going to write the raw headers
+    do_reverse_transforms = False
+
     # Get subset of the dataset that has not yet been saved
     # The headers might not be present in the dataset
     worker_variables = [data_variable_name]
     if header_key in dataset.data_vars:  # Keeping the `if` here to allow for more worker configurations
         worker_variables.append(header_key)
     if raw_header_key in dataset.data_vars:
+        do_reverse_transforms = True
         worker_variables.append(raw_header_key)
 
+    raw_headers, transformed_headers, traces = get_header_raw_and_transformed(segy_file, live_trace_indexes, do_reverse_transforms=do_reverse_transforms)
     ds_to_write = dataset[worker_variables]
 
     if header_key in worker_variables: