create_empty_like

Author: dmitriyrepin

pyproject.toml

@@ -55,6 +55,7 @@ dev = [
     "pre-commit-hooks>=6.0.0",
     "pytest>=8.4.2",
     "pytest-dependency>=0.6.0",
+    "pytest-order>=1.3.0",
     "typeguard>=4.4.4",
     "xdoctest[colors]>=1.3.0",
     "Pygments>=2.19.2"

src/mdio/creators/__init__.py

@@ -0,0 +1,5 @@
+"""MDIO Data creation API."""
+
+from mdio.creators.mdio import create_empty_like
+
+__all__ = ["create_empty_like"]

src/mdio/creators/mdio.py

@@ -0,0 +1,105 @@
+"""Creating MDIO v1 datasets."""
+
+from __future__ import annotations
+
+from datetime import UTC
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+from mdio.api.io import _normalize_path
+from mdio.api.io import open_mdio
+from mdio.api.io import to_mdio
+from mdio.builder.template_registry import TemplateRegistry
+from mdio.builder.xarray_builder import to_xarray_dataset
+from mdio.converters.segy import populate_dim_coordinates
+from mdio.converters.type_converter import to_structured_type
+from mdio.core.grid import Grid
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from segy.schema import HeaderSpec
+    from upath import UPath
+    from xarray import Dataset as xr_Dataset
+
+    from mdio.builder.schemas import Dataset
+    from mdio.core.dimension import Dimension
+
+
+def create_empty_like(  # noqa PLR0913
+    input_path: UPath | Path | str,
+    output_path: UPath | Path | str,
+    keep_coordinates: bool = False,
+    overwrite: bool = False,
+) -> xr_Dataset:
+    """A function that creates an empty MDIO v1 file with the same structure as an existing one.
+
+    Args:
+        input_path: The path of the input MDIO file.
+        output_path: The path of the output MDIO file.
+                     If None, the output will not be written to disk.
+        keep_coordinates: Whether to keep the coordinates in the output file.
+        overwrite: Whether to overwrite the output file if it exists.
+
+    Returns:
+        The output MDIO dataset.
+
+    Raises:
+        FileExistsError: If the output location already exists and overwrite is False.
+    """
+    input_path = _normalize_path(input_path)
+    output_path = _normalize_path(output_path) if output_path is not None else None
+
+    if not overwrite and output_path is not None and output_path.exists():
+        err = f"Output location '{output_path.as_posix()}' exists. Set `overwrite=True` if intended."
+        raise FileExistsError(err)
+
+    ds = open_mdio(input_path)
+
+    # Create a copy with the same structure but no data or,
+    # optionally, coordinates
+    ds_output = ds.copy(data=None).reset_coords(drop=not keep_coordinates)
+
+    # Dataset
+    # Keep the name (which is the same as the used template name) and the original API version
+    # ds_output.attrs["name"]
+    # ds_output.attrs["apiVersion"]
+    ds_output.attrs["createdOn"] = datetime.now(UTC)
+
+    # Coordinates
+    if not keep_coordinates:
+        for coord_name in ds_output.coords:
+            ds_output[coord_name].attrs["unitsV1"] = None
+
+    # MDIO attributes
+    attr = ds_output.attrs["attributes"]
+    if attr is not None:
+        attr.pop("gridOverrides", None)  # Empty dataset should not have gridOverrides
+        # Keep the original values for the following attributes
+        # attr["defaultVariableName"]
+        # attr["surveyType"]
+        # attr["gatherType"]
+
+    # "All traces should be marked as dead in empty dataset"
+    if "trace_mask" in ds_output.variables:
+        ds_output["trace_mask"][:] = False
+
+    # Data variable
+    var_name = attr["defaultVariableName"]
+    var = ds_output[var_name]
+    var.attrs["statsV1"] = None
+    if not keep_coordinates:
+        var.attrs["unitsV1"] = None
+
+    # SEG-Y file header
+    if "segy_file_header" in ds_output.variables:
+        segy_file_header = ds_output["segy_file_header"]
+        if segy_file_header is not None:
+            segy_file_header.attrs["textHeader"] = None
+            segy_file_header.attrs["binaryHeader"] = None
+            segy_file_header.attrs["rawBinaryHeader"] = None
+
+    if output_path is not None:
+        to_mdio(ds_output, output_path=output_path, mode="w", compute=True)
+
+    return ds_output

tests/conftest.py

@@ -46,6 +46,17 @@ def zarr_tmp(tmp_path_factory: pytest.TempPathFactory) -> Path:
     """Make a temp file for the output MDIO."""
     return tmp_path_factory.mktemp(r"mdio")
 
+@pytest.fixture(scope="session")
+def teapot_mdio_tmp(tmp_path_factory: pytest.TempPathFactory) -> Path:
+    """Make a temp file for the output MDIO."""
+    return tmp_path_factory.mktemp(r"teapot.mdio")
+
+
+@pytest.fixture(scope="module")
+def mdio_4d_tmp(tmp_path_factory: pytest.TempPathFactory) -> Path:
+    """Make a temp file for the output MDIO."""
+    return tmp_path_factory.mktemp(r"tmp_4d.mdio")
+
 
 @pytest.fixture(scope="module")
 def zarr_tmp2(tmp_path_factory: pytest.TempPathFactory) -> Path:  # pragma: no cover - used by disabled test
@@ -58,3 +69,31 @@ def segy_export_tmp(tmp_path_factory: pytest.TempPathFactory) -> Path:
     """Make a temp file for the round-trip IBM SEG-Y."""
     tmp_dir = tmp_path_factory.mktemp("segy")
     return tmp_dir / "teapot_roundtrip.segy"
+
+
+@pytest.fixture(scope="class")
+def empty_mdio_with_headers(tmp_path_factory: pytest.TempPathFactory) -> Path:
+    """Make a temp file for empty MDIO testing."""
+    path = tmp_path_factory.mktemp(r"empty_with_headers.mdio")
+    return path
+
+
+# @pytest.fixture(scope="session")
+# def tmp_path_factory() -> pytest.TempPathFactory:
+#     """Custom tmp_path_factory implementation for local debugging."""
+#     from pathlib import Path  # noqa: PLC0415
+
+#     class DebugTempPathFactory:
+#         def __init__(self) -> None:
+#             pass
+
+#         def mktemp(self, basename: str, numbered: bool = True) -> Path:
+#             _ = numbered
+#             path = self.getbasetemp() / basename
+#             path.mkdir(parents=True, exist_ok=True)
+#             return path
+
+#         def getbasetemp(self) -> Path:
+#             return Path("tmp")
+
+#     return DebugTempPathFactory()

tests/integration/test_create_empty.py

@@ -0,0 +1,151 @@
+"""Test for create_empty_mdio function."""
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING
+
+import numpy as np
+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
+from mdio.builder.schemas.v1.units import SpeedUnitEnum
+from mdio.builder.schemas.v1.units import SpeedUnitModel
+from mdio.builder.schemas.v1.units import TimeUnitEnum
+from mdio.builder.schemas.v1.units import TimeUnitModel
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from xarray import Dataset as xr_Dataset
+
+
+from tests.integration.test_segy_roundtrip_teapot import get_teapot_segy_spec
+from tests.integration.testing_helpers import get_values
+from tests.integration.testing_helpers import validate_variable
+
+from mdio import __version__
+from mdio.api.io import open_mdio
+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_like
+
+
+@pytest.mark.order(1000)
+class TestCreateEmptyPostStack3DTimeMdio:
+    """Tests for create_empty_mdio function."""
+
+    @classmethod
+    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=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),
+        ]
+        hs: HeaderSpec = get_segy_standard(1.0).trace.header
+        hs.customize(fields=trace_header_fields)
+        return hs
+
+    @classmethod
+    def _validate_dataset_metadata(cls, ds: xr_Dataset) -> None:
+        """Validate the dataset metadata."""
+        # Check basic metadata attributes
+        expected_attrs = {
+            "apiVersion": __version__,
+            "name": "PostStack3DTime",
+        }
+        actual_attrs_json = ds.attrs
+
+        # Compare one by one due to ever changing createdOn
+        for key, value in expected_attrs.items():
+            assert key in actual_attrs_json
+            if key == "createdOn":
+                assert actual_attrs_json[key] is not None
+            else:
+                assert actual_attrs_json[key] == value
+
+        # Check that createdOn exists
+        assert "createdOn" in actual_attrs_json
+        assert actual_attrs_json["createdOn"] is not None
+
+        # Validate template attributes
+        attributes = ds.attrs["attributes"]
+        assert attributes is not None
+        assert len(attributes) == 3
+        # Validate all attributes provided by the abstract template
+        assert attributes["defaultVariableName"] == "amplitude"
+        assert attributes["surveyType"] == "3D"
+        assert attributes["gatherType"] == "stacked"
+        assert "gridOverrides" not in attributes, "Empty dataset should not have gridOverrides"
+
+    @classmethod
+    def _validate_empty_mdio_dataset(cls, ds: xr_Dataset, has_headers: bool) -> None:
+        """Validate an empty MDIO dataset structure and content."""
+        # Check that the dataset has the expected shape
+        assert ds.sizes == {"inline": 345, "crossline": 188, "time": 1501}
+
+        # Validate the dimension coordinate variables
+        validate_variable(ds, "inline", (345,), ("inline",), np.int32, range(1, 346), get_values)
+        validate_variable(ds, "crossline", (188,), ("crossline",), np.int32, range(1, 189), get_values)
+        validate_variable(ds, "time", (1501,), ("time",), np.int32, range(0, 3002, 2), get_values)
+
+        # Validate the non-dimensional coordinate variables (should be empty for empty dataset)
+        validate_variable(ds, "cdp_x", (345, 188), ("inline", "crossline"), np.float64, None, None)
+        validate_variable(ds, "cdp_y", (345, 188), ("inline", "crossline"), np.float64, None, None)
+
+        if has_headers:
+            segy_spec = get_teapot_segy_spec()
+            # Validate the headers (should be empty for empty dataset)
+            # Infer the dtype from segy_spec and ignore endianness
+            header_dtype = segy_spec.trace.header.dtype.newbyteorder("native")
+            validate_variable(ds, "headers", (345, 188), ("inline", "crossline"), header_dtype, None, None)
+            validate_variable(ds, "segy_file_header", (), (), np.dtype("U1"), None, None)
+
+            assert "segy_file_header" in ds.variables
+            assert ds["segy_file_header"].attrs.get("textHeader", None) is None, (
+                "TextHeader should be empty for empty dataset"
+            )
+            assert ds["segy_file_header"].attrs.get("binaryHeader", None) is None, (
+                "BinaryHeader should be empty for empty dataset"
+            )
+            assert ds["segy_file_header"].attrs.get("rawBinaryHeader", None) is None, (
+                "RawBinaryHeader should be empty for empty dataset"
+            )
+        else:
+            assert "headers" not in ds.variables
+            assert "segy_file_header" not in ds.variables
+
+        # Validate the trace mask
+        validate_variable(ds, "trace_mask", (345, 188), ("inline", "crossline"), np.bool_, None, None)
+        trace_mask = ds["trace_mask"].values
+        assert not np.any(trace_mask), "All traces should be marked as dead in empty dataset"
+
+        # Validate the amplitude data (should be empty)
+        validate_variable(ds, "amplitude", (345, 188, 1501), ("inline", "crossline", "time"), np.float32, None, None)
+        assert ds["amplitude"].attrs.get("statsV1", None) is None, "StatsV1 should be empty for empty dataset"
+        assert ds["amplitude"].attrs.get("unitsV1", None) is None, "UnitsV1 should be empty for empty dataset"
+
+    @pytest.mark.order(1001)
+    @pytest.mark.dependency
+    def test_create_empty_like(self, teapot_mdio_tmp: Path, empty_mdio_with_headers: Path) -> None:
+        """Create an empty MDIO file like the input file."""
+        _ = empty_mdio_with_headers
+        ds = create_empty_like(
+            input_path=teapot_mdio_tmp,
+            output_path=None,  # We don't want to write to disk for now
+            keep_coordinates=True,
+            overwrite=True,
+        )
+        self._validate_dataset_metadata(ds)
+        self._validate_empty_mdio_dataset(ds, has_headers=True)