Begin implementation of header disaster recovery array

Author: BrianMichell

src/mdio/constants.py

@@ -56,4 +56,5 @@
     ScalarType.COMPLEX64: complex(np_nan, np_nan),
     ScalarType.COMPLEX128: complex(np_nan, np_nan),
     ScalarType.COMPLEX256: complex(np_nan, np_nan),
+    ScalarType.HEADERS: b"\x00" * 240,
 }

src/mdio/converters/segy.py

@@ -22,6 +22,11 @@
 from mdio.schemas.v1.units import AllUnits
 from mdio.schemas.v1.units import LengthUnitEnum
 from mdio.schemas.v1.units import LengthUnitModel
+from mdio.schemas.dtype import ScalarType
+from mdio.schemas.v1.variable import Variable
+from mdio.schemas.compressors import Blosc, BloscCname
+from mdio.schemas.metadata import ChunkGridMetadata
+from mdio.schemas.chunk_grid import RegularChunkGrid, RegularChunkShape
 from mdio.segy import blocked_io
 from mdio.segy.utilities import get_grid_plan
 
@@ -312,6 +317,56 @@ def _add_text_binary_headers(dataset: Dataset, segy_file: SegyFile) -> None:
     )
 
 
+def _add_raw_headers_to_template(mdio_template: AbstractDatasetTemplate) -> None:
+    """Add raw headers capability to the MDIO template by monkey-patching its _add_variables method.
+    
+    This function modifies the template's _add_variables method to also add a raw headers variable
+    with the following characteristics:
+    - Same rank as the Headers variable (all dimensions except vertical)
+    - Name: "RawHeaders"
+    - Type: ScalarType.HEADERS
+    - No coordinates
+    - zstd compressor
+    - No additional metadata
+    - Chunked the same as the Headers variable
+    
+    Args:
+        mdio_template: The MDIO template to mutate
+    """
+    # Store the original _add_variables method
+    original_add_variables = mdio_template._add_variables
+    
+    def enhanced_add_variables():
+        # Call the original method first
+        original_add_variables()
+        
+        # Now add the raw headers variable
+        chunk_shape = mdio_template._var_chunk_shape[:-1]
+        
+        # Create chunk grid metadata
+        chunk_metadata = ChunkGridMetadata(
+            chunk_grid=RegularChunkGrid(
+                configuration=RegularChunkShape(chunk_shape=chunk_shape)
+            )
+        )
+        
+        # Add the raw headers variable using the builder's add_variable method
+        mdio_template._builder.add_variable(
+            name="raw_headers",
+            long_name="Raw Headers",
+            dimensions=mdio_template._dim_names[:-1],  # All dimensions except vertical
+            data_type=ScalarType.HEADERS,
+            compressor=Blosc(cname=BloscCname.zstd),
+            coordinates=None,  # No coordinates as specified
+            metadata_info=[chunk_metadata],
+        )
+    
+    # Replace the template's _add_variables method
+    mdio_template._add_variables = enhanced_add_variables
+
+    return mdio_template
+
+
 def segy_to_mdio(
     segy_spec: SegySpec,
     mdio_template: AbstractDatasetTemplate,
@@ -350,6 +405,9 @@ def segy_to_mdio(
     # https://github.com/TGSAI/mdio-python/issues/601
     headers = to_structured_type(segy_spec.trace.header.dtype)
 
+    if os.getenv("MDIO__DO_RAW_HEADERS") == "1":
+        mdio_template = _add_raw_headers_to_template(mdio_template)
+
     horizontal_unit = _get_horizontal_coordinate_unit(segy_dimensions)
     mdio_ds: Dataset = mdio_template.build_dataset(
         name=mdio_template.name,

src/mdio/converters/type_converter.py

@@ -78,6 +78,8 @@ def to_structured_type(data_type: np_dtype) -> StructuredType:
 def to_numpy_dtype(data_type: ScalarType | StructuredType) -> np_dtype:
     """Get the numpy dtype for a variable."""
     if isinstance(data_type, ScalarType):
+        if data_type == ScalarType.HEADERS:
+            return np_dtype("|V240")
         return np_dtype(data_type.value)
     if isinstance(data_type, StructuredType):
         return np_dtype([(f.name, f.format.value) for f in data_type.fields])

src/mdio/schemas/dtype.py

@@ -32,6 +32,7 @@ class ScalarType(StrEnum):
     COMPLEX64 = "complex64"
     COMPLEX128 = "complex128"
     COMPLEX256 = "complex256"
+    HEADERS = "r1920"  # Raw number of BITS, must be a multiple of 8
 
 
 class StructuredField(CamelCaseStrictModel):

src/mdio/segy/_workers.py

@@ -122,12 +122,15 @@ def trace_worker(  # noqa: PLR0913
     traces = segy_file.trace[live_trace_indexes]
 
     header_key = "headers"
+    raw_header_key = "raw_headers"
 
     # 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:
+        worker_variables.append(raw_header_key)
 
     ds_to_write = dataset[worker_variables]
 
@@ -146,6 +149,15 @@ 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.
         )
+    if raw_header_key in worker_variables:
+        tmp_raw_headers = np.zeros_like(dataset[raw_header_key])
+        tmp_raw_headers[not_null] = traces.header.view("|V240")
+        ds_to_write[raw_header_key] = Variable(
+            ds_to_write[raw_header_key].dims,
+            tmp_raw_headers,
+            attrs=ds_to_write[raw_header_key].attrs,
+            encoding=ds_to_write[raw_header_key].encoding,  # Not strictly necessary, but safer than not doing it.
+        )
 
     data_variable = ds_to_write[data_variable_name]
     fill_value = _get_fill_value(ScalarType(data_variable.dtype.name))