Add export to segy to test_populate_empty_dataset

Author: dmitriyrepin

src/mdio/creators/mdio.py

@@ -66,6 +66,11 @@ def create_empty(  # noqa PLR0913
     drop_vars_delayed = []
     dataset, drop_vars_delayed = populate_dim_coordinates(xr_dataset, grid, drop_vars_delayed=drop_vars_delayed)
 
+    if headers:
+        # Since the headers were provided, the user wants to export to SEG-Y
+        # Add a dummy segy_file_header variable used to export to SEG-Y
+        dataset["segy_file_header"] = ((), "")
+
     # Create the Zarr store with the correct structure but with empty arrays
     to_mdio(dataset, output_path=output_path, mode="w", compute=False)
 

tests/integration/test_create_empty.py

@@ -9,6 +9,8 @@
 import pytest
 from segy.schema import HeaderField
 from segy.schema import HeaderSpec
+from segy.schema import ScalarType
+from segy.standards import get_segy_standard
 
 from mdio.builder.schemas.v1.units import LengthUnitEnum
 from mdio.builder.schemas.v1.units import LengthUnitModel
@@ -22,6 +24,7 @@
 
     from xarray import Dataset as xr_Dataset
 
+
 from tests.integration.testing_helpers import get_values
 from tests.integration.testing_helpers import validate_variable
 
@@ -30,6 +33,7 @@
 from mdio.api.io import to_mdio
 from mdio.builder.schemas.v1.stats import CenteredBinHistogram
 from mdio.builder.schemas.v1.stats import SummaryStatistics
+from mdio.converters.mdio import mdio_to_segy
 from mdio.core import Dimension
 from mdio.creators.mdio import create_empty
 
@@ -38,15 +42,18 @@ class TestCreateEmptyPostStack3DTimeMdio:
     """Tests for create_empty_mdio function."""
 
     @classmethod
-    def _get_header_spec(cls) -> HeaderSpec:
+    def _get_customized_v10_trace_header_spec(cls) -> HeaderSpec:
         """Get the header spec for the MDIO dataset."""
         trace_header_fields = [
-            HeaderField(name="inline", byte=17, format="int32"),
-            HeaderField(name="crossline", byte=13, format="int32"),
-            HeaderField(name="cdp_x", byte=181, format="int32"),
-            HeaderField(name="cdp_y", byte=185, format="int32"),
+            HeaderField(name="inline", byte=17, format=ScalarType.INT32),
+            HeaderField(name="crossline", byte=13, format=ScalarType.INT32),
+            HeaderField(name="cdp_x", byte=181, format=ScalarType.INT32),
+            HeaderField(name="cdp_y", byte=185, format=ScalarType.INT32),
+            HeaderField(name="coordinate_scalar", byte=71, format=ScalarType.INT16),
         ]
-        return HeaderSpec(fields=trace_header_fields)
+        hs: HeaderSpec = get_segy_standard(1.0).trace.header
+        hs.customize(fields=trace_header_fields)
+        return hs
 
     @classmethod
     def _validate_empty_mdio_dataset(cls, ds: xr_Dataset, has_headers: bool) -> None:
@@ -66,10 +73,12 @@ def _validate_empty_mdio_dataset(cls, ds: xr_Dataset, has_headers: bool) -> None
         if has_headers:
             # Validate the headers (should be empty for empty dataset)
             # Infer the dtype from segy_spec and ignore endianness
-            header_dtype = cls._get_header_spec().dtype.newbyteorder("native")
+            header_dtype = cls._get_customized_v10_trace_header_spec().dtype.newbyteorder("native")
             validate_variable(ds, "headers", (200, 300), ("inline", "crossline"), header_dtype, None, None)
+            validate_variable(ds, "segy_file_header", (), (), np.dtype("U1"), None, None)
         else:
             assert "headers" not in ds.variables
+            assert "segy_file_header" not in ds.variables
 
         # Validate the trace mask (should be all True for empty dataset)
         validate_variable(ds, "trace_mask", (200, 300), ("inline", "crossline"), np.bool_, None, None)
@@ -89,7 +98,9 @@ def _create_empty_mdio(cls, create_headers: bool, output_path: Path, overwrite:
             Dimension(name="time", coords=range(0, 3000, 4)),  # 0-3 seconds 4ms sample rate
         ]
 
-        headers = cls._get_header_spec() if create_headers else None
+        # If later on, we want to export to SEG-Y, we need to provide the trace header spec.
+        # The HeaderSpec can be either standard or customized.
+        headers = cls._get_customized_v10_trace_header_spec() if create_headers else None
         create_empty(
             mdio_template_name="PostStack3DTime",
             dimensions=dims,
@@ -167,7 +178,7 @@ def test_overwrite_behavior(self, empty_mdio_dir: Path) -> None:
         garbage_file = empty_mdio / "garbage.txt"
         garbage_file.write_text("This is garbage data that should be overwritten")
         garbage_dir = empty_mdio / "garbage_dir"
-        garbage_dir.mkdir()
+        garbage_dir.mkdir(exist_ok=True)
         (garbage_dir / "nested_garbage.txt").write_text("More garbage")
 
         # Verify the directory exists with garbage data
@@ -201,7 +212,9 @@ def test_populate_empty_dataset(self, mdio_with_headers: Path) -> None:
         #   HACK: in this example, we will use this variable to store the velocity data
         # * 'trace_mask' variable was created and pre-populated with 'False' fill values
         #   (all traces are marked as dead)
-        # * 'headers' segy trace headers variable was created (if the dataset was created with headers not None)
+        # * 'headers' and 'segy_file_header' variables were created (if the dataset was created with
+        #   headers not None). The 'headers' variable structured datatype is defined by the HeaderSpec
+        # that was used to create the empty MDIO
         # * dataset attribute called 'attributes' was created
         ds = open_mdio(mdio_with_headers)
 
@@ -223,7 +236,7 @@ def test_populate_empty_dataset(self, mdio_with_headers: Path) -> None:
             sum_squares=(np.ma.power(nonzero_samples, 2).sum(dtype="float64")),
             histogram=CenteredBinHistogram(bin_centers=[], counts=[]),
         )
-        velocity.attrs["statsV1"] = stats.model_dump_json()
+        velocity.attrs["statsV1"] = stats.model_dump(mode="json")
 
         # 3) Populate the non-dimensional coordinate variables 'cdp_x' and 'cdp_y' (optional)
         origin = [270000, 3290000]  # survey x, y origin
@@ -239,24 +252,70 @@ def test_populate_empty_dataset(self, mdio_with_headers: Path) -> None:
         ds.trace_mask[:] = ~np.isnan(velocity[:, :, 0])
 
         # 5) Set coordinate and data variable units (optional)
-        ds.time["unitsV1"] = TimeUnitModel(time=TimeUnitEnum.MILLISECOND).model_dump_json()
+        ds.time.attrs["unitsV1"] = TimeUnitModel(time=TimeUnitEnum.MILLISECOND).model_dump(mode="json")
 
-        ds.cdp_x.attrs["unitsV1"] = LengthUnitModel(length=LengthUnitEnum.FOOT).model_dump_json()
-        ds.cdp_x.attrs["unitsV1"] = LengthUnitModel(length=LengthUnitEnum.FOOT).model_dump_json()
+        ds.cdp_x.attrs["unitsV1"] = LengthUnitModel(length=LengthUnitEnum.FOOT).model_dump(mode="json")
+        ds.cdp_x.attrs["unitsV1"] = LengthUnitModel(length=LengthUnitEnum.FOOT).model_dump(mode="json")
 
-        velocity.attrs["unitsV1"] = SpeedUnitModel(speed=SpeedUnitEnum.FEET_PER_SECOND).model_dump_json()
+        velocity.attrs["unitsV1"] = SpeedUnitModel(speed=SpeedUnitEnum.FEET_PER_SECOND).model_dump(mode="json")
 
-        # 6) Populate dataset's segy trace headers, if those were created (optional)
+        # 6) Populate dataset's segy trace headers, if those were created (required only if we want to export to SEG-Y)
         if "headers" in ds.variables:
-            # numpy broadcasting (200, 1) array to (200, 300) array
-            ds["headers"].values["inline"] = ds.inline.values[:, np.newaxis]
-            # numpy broadcasting (1, 300) array to (200, 300) array
-            ds["headers"].values["crossline"] = ds.crossline.values[np.newaxis, :]
-            ds["headers"]["cdp_x"][:] = ds.cdp_x
-            ds["headers"]["cdp_y"][:] = ds.cdp_y
+            # Both the structured "headers" and the dummy "segy_file_header" variables are
+            # required to enable SEG-Y to MDIO conversion
+
+            # Populate the structured trace "headers" variable
+            ds["headers"].values["inline"] = inline_grid
+            ds["headers"].values["crossline"] = xline_grid
+            # coordinate_scalar:
+            # Scalar to be applied to all coordinates specified in Standard Trace Header bytes
+            # 73–88 and to bytes Trace Header 181–188 to give the real value. Scalar = 1,
+            # ±10, ±100, ±1000, or ±10,000. If positive, scalar is used as a multiplier; if
+            # negative, scalar is used as divisor. A value of zero is assumed to be a scalar
+            # value of 1.
+            ds["headers"].values["coordinate_scalar"][:] = np.int16(-100)
+            ds["headers"].values["cdp_x"][:] = np.int32(ds.cdp_x * 100)
+            ds["headers"].values["cdp_y"][:] = np.int32(ds.cdp_y * 100)
+
+            # Fill its metadata (.attrs) with 'textHeader' and 'binaryHeader'.
+            ds["segy_file_header"].attrs.update(
+                {
+                    "textHeader": "\n".join(
+                        [
+                            "C01 BYTES 13-16: CROSSLINE         " + " " * 47,
+                            "C02 BYTES 17-20: INLINE            " + " " * 47,
+                            "C03 BYTES 71-74: COORDINATE SCALAR " + " " * 47,
+                            "C04 BYTES 181-184: CDP X           " + " " * 47,
+                            "C05 BYTES 185-188: CDP Y           " + " " * 47,
+                            *(f"C{i:02d}" + " " * 77 for i in range(6, 41)),
+                        ]
+                    ),
+                    "binaryHeader": {
+                        "data_sample_format": 1,
+                        "sample_interval": int(ds.time[1] - ds.time[0]),
+                        "samples_per_trace": ds.time.size,
+                        "segy_revision_major": 0,
+                        "segy_revision_minor": 0,
+                    },
+                }
+            )
 
         # 7) Create dataset's custom attributes (optional)
         ds.attrs["attributes"]["createdBy"] = "John Doe"
 
-        output_path = mdio_with_headers.parent / "populated_empty.mdio"
-        to_mdio(ds, output_path=output_path, mode="w", compute=True)
+        # 8) Export to MDIO
+        output_path_mdio = mdio_with_headers.parent / "populated_empty.mdio"
+        to_mdio(ds, output_path=output_path_mdio, mode="w", compute=True)
+
+        # 9) Convert the populated emptyMDIO to SEG-Y
+        if "headers" in ds.variables:
+            # Select the SEG-Y standard to use for the conversion
+            custom_segy_spec = get_segy_standard(1.0)
+            # Customize to use the same HeaderSpec that was used to create the empty MDIO
+            custom_segy_spec.trace.header = self._get_customized_v10_trace_header_spec()
+            # Convert the MDIO file to SEG-Y
+            mdio_to_segy(
+                segy_spec=custom_segy_spec,
+                input_path=output_path_mdio,
+                output_path=mdio_with_headers.parent / "populated_empty.sgy",
+            )