Working end-to-end examples

Author: BrianMichell

src/mdio/converters/segy.py

@@ -12,6 +12,8 @@
 from segy.config import SegySettings
 from segy.standards.codes import MeasurementSystem as segy_MeasurementSystem
 from segy.standards.fields.trace import Rev0 as TraceHeaderFieldsRev0
+from segy.schema import HeaderField
+from segy.schema import ScalarType as ScalarType2
 
 from mdio.api.io import _normalize_path
 from mdio.api.io import to_mdio
@@ -340,6 +342,45 @@ def _add_grid_override_to_metadata(dataset: Dataset, grid_overrides: dict[str, A
     if grid_overrides is not None:
         dataset.metadata.attributes["gridOverrides"] = grid_overrides
 
+def _scalar_to_size(scalar: ScalarType2) -> int:
+    if scalar == ScalarType2.UINT8:
+        return 1
+    elif scalar == ScalarType2.UINT16:
+        return 2
+    elif scalar == ScalarType2.UINT32:
+        return 4
+    elif scalar == ScalarType2.UINT64:
+        return 8
+    elif scalar == ScalarType2.INT8:
+        return 1
+    elif scalar == ScalarType2.INT16:
+        return 2
+    elif scalar == ScalarType2.INT32:
+        return 4
+    elif scalar == ScalarType2.INT64:
+        return 8
+    elif scalar == ScalarType2.FLOAT32:
+        return 4
+    elif scalar == ScalarType2.FLOAT64:
+        return 8
+    elif scalar == ScalarType2.FLOAT16:
+        return 2
+    elif scalar == ScalarType2.STRING8:
+        return 8
+    else:
+        raise ValueError(f"Invalid scalar type: {scalar}")
+
+def _customize_segy_spec(segy_spec: SegySpec) -> SegySpec:
+    assigned_bytes = []
+    for field in segy_spec.trace.header.fields:
+        byte = field.byte-1
+        for i in range(byte, byte + _scalar_to_size(field.format)):
+            assigned_bytes.append(i)
+    unassigned_bytes = [i for i in range(240) if i not in assigned_bytes]
+    field_to_customize = [HeaderField(name=f"Field_{i}", format=ScalarType.UINT8, byte=i+1) for i in unassigned_bytes]
+    segy_spec = segy_spec.customize(trace_header_fields=field_to_customize)
+    return segy_spec
+
 
 def _add_raw_headers_to_template(mdio_template: AbstractDatasetTemplate) -> AbstractDatasetTemplate:
     """Add raw headers capability to the MDIO template by monkey-patching its _add_variables method.
@@ -421,6 +462,9 @@ def segy_to_mdio(  # noqa PLR0913
     input_path = _normalize_path(input_path)
     output_path = _normalize_path(output_path)
 
+    if os.getenv("MDIO__DO_RAW_HEADERS") == "1":
+        segy_spec = _customize_segy_spec(segy_spec)
+
     if not overwrite and output_path.exists():
         err = f"Output location '{output_path.as_posix()}' exists. Set `overwrite=True` if intended."
         raise FileExistsError(err)

src/mdio/segy/_workers.py

@@ -148,25 +148,8 @@ def trace_worker(  # noqa: PLR0913
     ds_to_write = dataset[worker_variables]
 
     if raw_header_key in worker_variables:
-        # from segy.schema import HeaderField, ScalarType
-        from segy.schema import HeaderField
-        from segy.schema import ScalarType as ScalarType2
-
-        new_spec = deepcopy(segy_file.spec)
-        new_spec.trace.header.fields = []
-        new_fields = []
-
-        for i in range(240):
-            new_fields.append(HeaderField(name=f"Field_{i}", format=ScalarType.UINT8, byte=i+1))
-
-        # # new_spec.trace.header.fields = new_fields
-        new_spec = new_spec.customize(trace_header_fields=new_fields)
-        updated_segy_file = SegyFile(segy_file.url, spec=new_spec)
-        updated_traces = updated_segy_file.trace[live_trace_indexes]
-
         tmp_raw_headers = np.zeros_like(dataset[raw_header_key])
 
-
         # _foo = traces.tobytes()[:240]
 
         # _asContiguousCopy = np.ascontiguousarray(traces.header.copy()).view("|V240")
@@ -185,8 +168,6 @@ def trace_worker(  # noqa: PLR0913
 
         # _aahhhhhhh = segy_file.header[0]
 
-    
-
         # _asBytes = traces.header.tobytes()
         # _are_equal = _foo == _asBytes[:240]
         # _asBytesLen = len(_asBytes)
@@ -195,18 +176,18 @@ def trace_worker(  # noqa: PLR0913
         # _asContiguousBytesView = np.ascontiguousarray(traces.header[0]).tobytes().view("|V240")
         # _asBuffer = memoryview(traces.header)
 
-
         # _type = traces.header.dtype
 
         # tmp_raw_headers[not_null] = traces.raw_header
-        # tmp_raw_headers[not_null] = np.ascontiguousarray(traces.header.copy()).view("|V240")
+        # tmp_raw_headers[not_null] = np.ascontiguousarray(traces.header.copy(order="C")).view("|V240")
+        tmp_raw_headers[not_null] = np.ascontiguousarray(traces.header.copy()).view("|V240")
         # tmp_raw_headers[not_null] = _asContiguousBytes
         # tmp_raw_headers[not_null] = _asBytes
         # tmp_raw_headers[not_null] = traces.header.view("|V240")
         # tmp_raw_headers[not_null] = np.ascontiguousarray(traces.header.copy()).view("|V240")  # Leaks numpy metadata
 
 
-        tmp_raw_headers[not_null] = updated_traces.header.view("|V240")
+        # tmp_raw_headers[not_null] = updated_traces.header.view("|V240")
 
         ds_to_write[raw_header_key] = Variable(
             ds_to_write[raw_header_key].dims,

tests/integration/test_segy_import_export_masked.py

@@ -547,7 +547,6 @@ def read_segy_trace_header(trace_index: int) -> bytes:
         segy_trace_idx = 0
         flat_mask = trace_mask.ravel()
         flat_raw_headers = raw_headers_data.ravel()  # Flatten to 1D array of 240-byte header records
-        print(f"Flat mask shape: {flat_mask.shape}")
 
         operation = 'w'