Merge upstream/main' into create_empty

dmitriyrepin · dmitriyrepin · commit 251e2f68cc45 · 2025-10-27T19:29:43.000Z
diff --git a/src/mdio/creators/mdio.py b/src/mdio/creators/mdio.py
@@ -20,11 +20,12 @@
     from xarray import Dataset as xr_Dataset
 
     from mdio.builder.schemas import Dataset
+    from mdio.builder.templates.base import AbstractDatasetTemplate
     from mdio.core.dimension import Dimension
 
 
 def create_empty(  # noqa PLR0913
-    mdio_template_name: str,
+    mdio_template: AbstractDatasetTemplate | str,
     dimensions: list[Dimension],
     output_path: UPath | Path | str,
     headers: HeaderSpec | None = None,
@@ -33,7 +34,14 @@ def create_empty(  # noqa PLR0913
     """A function that creates an empty MDIO v1 file with known dimensions.
 
     Args:
-        mdio_template_name: The MDIO template to use to define the dataset structure.
+        mdio_template: The MDIO template or template name to use to define the dataset structure.
+            NOTE: If you want to have a unit-aware MDIO model, you need to add the units
+            to the template before calling this function. For example:
+            'unit_aware_template = TemplateRegistry().get("PostStack3DTime")'
+            'unit_aware_template.add_units({"time": UNITS_SECOND})'
+            'unit_aware_template.add_units({"cdp_x": UNITS_METER})'
+            'unit_aware_template.add_units({"cdp_y": UNITS_METER})'
+            'create_empty(unit_aware_template, dimensions, output_path, headers, overwrite)'
         dimensions: The dimensions of the MDIO file.
         output_path: The universal path for the output MDIO v1 file.
         headers: SEG-Y v1.0 trace headers. Defaults to None.
@@ -50,13 +58,11 @@ def create_empty(  # noqa PLR0913
 
     header_dtype = to_structured_type(headers.dtype) if headers else None
     grid = Grid(dims=dimensions)
-    mdio_template = TemplateRegistry().get(mdio_template_name)
-    mdio_ds: Dataset = mdio_template.build_dataset(
-        name=mdio_template_name,
-        sizes=grid.shape,
-        horizontal_coord_unit=None,
-        header_dtype=header_dtype,
-    )
+    if isinstance(mdio_template, str):
+        # A template name is passed in. Get a unit-unaware template from registry
+        mdio_template = TemplateRegistry().get(mdio_template)
+    # Build the dataset using the template
+    mdio_ds: Dataset = mdio_template.build_dataset(name=mdio_template.name, sizes=grid.shape, header_dtype=header_dtype)
 
     # Convert to xarray dataset
     xr_dataset: xr_Dataset = to_xarray_dataset(mdio_ds=mdio_ds)
diff --git a/tests/integration/test_create_empty.py b/tests/integration/test_create_empty.py
@@ -26,17 +26,61 @@
 
 
 from tests.integration.testing_helpers import get_values
-from tests.integration.testing_helpers import validate_variable
+from tests.integration.testing_helpers import validate_xr_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.builder.templates.seismic_3d_poststack import Seismic3DPostStackTemplate
 from mdio.converters.mdio import mdio_to_segy
 from mdio.core import Dimension
 from mdio.creators.mdio import create_empty
 
+UNITS_NONE = None
+UNITS_METER = LengthUnitModel(length=LengthUnitEnum.METER)
+UNITS_SECOND = TimeUnitModel(time=TimeUnitEnum.SECOND)
+UNITS_METER_PER_SECOND = SpeedUnitModel(speed=SpeedUnitEnum.METER_PER_SECOND)
+UNITS_FOOT = LengthUnitModel(length=LengthUnitEnum.FOOT)
+UNITS_FEET_PER_SECOND = SpeedUnitModel(speed=SpeedUnitEnum.FEET_PER_SECOND)
+
+
+class PostStack3DVelocityTemplate(Seismic3DPostStackTemplate):
+    """Custom template that uses 'velocity' as the default variable name instead of 'amplitude'."""
+
+    @property
+    def _default_variable_name(self) -> str:
+        """Override the default variable name."""
+        return "velocity"
+
+    def __init__(self, data_domain: str, is_metric: bool) -> None:
+        super().__init__(data_domain)
+        if is_metric:
+            self._units.update(
+                {
+                    "time": UNITS_SECOND,
+                    "cdp_x": UNITS_METER,
+                    "cdp_y": UNITS_METER,
+                    "velocity": UNITS_METER_PER_SECOND,
+                }
+            )
+        else:
+            self._units.update(
+                {
+                    "time": UNITS_SECOND,
+                    "cdp_x": UNITS_FOOT,
+                    "cdp_y": UNITS_FOOT,
+                    "velocity": UNITS_FEET_PER_SECOND,
+                }
+            )
+
+    @property
+    def _name(self) -> str:
+        """Override the name of the template."""
+        domain_suffix = self._data_domain.capitalize()
+        return f"PostStack3DVelocity{domain_suffix}"
+
 
 class TestCreateEmptyPostStack3DTimeMdio:
     """Tests for create_empty_mdio function."""
@@ -58,35 +102,51 @@ def _get_customized_v10_trace_header_spec(cls) -> HeaderSpec:
     @classmethod
     def _validate_empty_mdio_dataset(cls, ds: xr_Dataset, has_headers: bool) -> None:
         """Validate an empty MDIO dataset structure and content."""
+        assert ds.name == "PostStack3DVelocityTime"
         # 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_xr_variable(ds, "inline", {"inline": 345}, UNITS_NONE, np.int32, range(1, 346), get_values)
+        validate_xr_variable(ds, "crossline", {"crossline": 188}, UNITS_NONE, np.int32, range(1, 189), get_values)
+        validate_xr_variable(ds, "time", {"time": 1501}, UNITS_SECOND, 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)
+        validate_xr_variable(ds, "cdp_x", {"inline": 345, "crossline": 188}, UNITS_METER, np.float64, None, None)
+        validate_xr_variable(ds, "cdp_y", {"inline": 345, "crossline": 188}, UNITS_METER, np.float64, None, 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_customized_v10_trace_header_spec().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)
+            validate_xr_variable(ds, "headers", {"inline": 345, "crossline": 188}, UNITS_NONE, header_dtype, None, None)
+            validate_xr_variable(
+                ds,
+                "segy_file_header",
+                dims={},
+                units=UNITS_NONE,
+                data_type=np.dtype("U1"),
+                expected_values=None,
+                actual_value_generator=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", (345, 188), ("inline", "crossline"), np.bool_, None, None)
+        validate_xr_variable(ds, "trace_mask", {"inline": 345, "crossline": 188}, UNITS_NONE, 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)
+        validate_xr_variable(
+            ds,
+            "velocity",
+            {"inline": 345, "crossline": 188, "time": 1501},
+            UNITS_METER_PER_SECOND,
+            np.float32,
+            None,
+            None,
+        )
 
     @classmethod
     def _create_empty_mdio(cls, create_headers: bool, output_path: Path, overwrite: bool = True) -> None:
@@ -101,8 +161,10 @@ def _create_empty_mdio(cls, create_headers: bool, output_path: Path, overwrite:
         # 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 an empty MDIO v1 metric post-stack 3D time velocity dataset
         create_empty(
-            mdio_template_name="PostStack3DTime",
+            mdio_template=PostStack3DVelocityTemplate(data_domain="time", is_metric=True),
             dimensions=dims,
             output_path=output_path,
             headers=headers,
@@ -138,7 +200,7 @@ def test_dataset_metadata(self, mdio_with_headers: Path) -> None:
         # Check basic metadata attributes
         expected_attrs = {
             "apiVersion": __version__,
-            "name": "PostStack3DTime",
+            "name": "PostStack3DVelocityTime",
         }
         actual_attrs_json = ds.attrs
 
@@ -159,7 +221,7 @@ def test_dataset_metadata(self, mdio_with_headers: Path) -> None:
         assert attributes is not None
         assert len(attributes) == 3
         # Validate all attributes provided by the abstract template
-        assert attributes["defaultVariableName"] == "amplitude"
+        assert attributes["defaultVariableName"] == "velocity"
         assert attributes["surveyType"] == "3D"
         assert attributes["gatherType"] == "stacked"
 
@@ -203,9 +265,9 @@ def test_overwrite_behavior(self, empty_mdio_dir: Path) -> None:
 
     def test_populate_empty_dataset(self, mdio_with_headers: Path) -> None:
         """Test showing how to populate empty dataset."""
-        # Open an empty PostStack3DTime dataset with SEG-Y 1.0 headers
+        # Open an empty PostStack3DVelocityTime dataset with SEG-Y 1.0 headers
         # NOTES:
-        # When this empty dataset was created from the 'PostStack3DTime' template and dimensions,
+        # When this empty dataset was created from the 'PostStack3DVelocityTime' template and dimensions,
         # * 'inline', 'crossline', and 'time' dimension coordinate variables were created and pre-populated
         # * 'cdp_x', 'cdp_y' non-dimensional coordinate variables were created
         # * 'amplitude' variable was created (the name of this variable is specified in the template)
diff --git a/tests/integration/test_segy_roundtrip_teapot.py b/tests/integration/test_segy_roundtrip_teapot.py
@@ -14,11 +14,15 @@
 from segy.standards import get_segy_standard
 from tests.integration.testing_helpers import get_inline_header_values
 from tests.integration.testing_helpers import get_values
-from tests.integration.testing_helpers import validate_variable
+from tests.integration.testing_helpers import validate_xr_variable
 
 from mdio import __version__
 from mdio import mdio_to_segy
 from mdio.api.io import open_mdio
+from mdio.builder.schemas.v1.units import LengthUnitEnum
+from mdio.builder.schemas.v1.units import LengthUnitModel
+from mdio.builder.schemas.v1.units import TimeUnitEnum
+from mdio.builder.schemas.v1.units import TimeUnitModel
 from mdio.builder.template_registry import TemplateRegistry
 from mdio.converters.segy import segy_to_mdio
 from mdio.segy.file import SegyFileWrapper
@@ -148,6 +152,11 @@ def raw_binary_header_teapot_dome() -> str:
     )
 
 
+UNITS_NONE = None
+UNITS_METER = LengthUnitModel(length=LengthUnitEnum.METER)
+UNITS_SECOND = TimeUnitModel(time=TimeUnitEnum.SECOND)
+
+
 class TestTeapotRoundtrip:
     """Tests for Teapot Dome data ingestion and export."""
 
@@ -163,9 +172,13 @@ def test_teapot_import(
 
         NOTE: This test must be executed before the 'TestReader' and 'TestExport' tests.
         """
+        unit_aware_template = TemplateRegistry().get("PostStack3DTime")
+        unit_aware_template.add_units({"time": UNITS_SECOND})
+        unit_aware_template.add_units({"cdp_x": UNITS_METER})
+        unit_aware_template.add_units({"cdp_y": UNITS_METER})
         segy_to_mdio(
             segy_spec=teapot_segy_spec,
-            mdio_template=TemplateRegistry().get("PostStack3DTime"),
+            mdio_template=unit_aware_template,
             input_path=segy_input,
             output_path=zarr_tmp,
             overwrite=True,
@@ -224,38 +237,38 @@ def test_grid(self, zarr_tmp: Path, teapot_segy_spec: SegySpec) -> None:
         ds = open_mdio(zarr_tmp)
 
         # 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_xr_variable(ds, "inline", {"inline": 345}, UNITS_NONE, np.int32, range(1, 346), get_values)
+        validate_xr_variable(ds, "crossline", {"crossline": 188}, UNITS_NONE, np.int32, range(1, 189), get_values)
+        validate_xr_variable(ds, "time", {"time": 1501}, UNITS_SECOND, np.int32, range(0, 3002, 2), get_values)
 
         # Validate the non-dimensional coordinate variables
-        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)
+        validate_xr_variable(ds, "cdp_x", {"inline": 345, "crossline": 188}, UNITS_METER, np.float64, None, None)
+        validate_xr_variable(ds, "cdp_y", {"inline": 345, "crossline": 188}, UNITS_METER, np.float64, None, None)
 
         # Validate the headers
         # We have a custom set of headers since we used customize_segy_specs()
         segy_spec = teapot_segy_spec
         data_type = segy_spec.trace.header.dtype
 
-        validate_variable(
+        validate_xr_variable(
             ds,
             "headers",
-            (345, 188),
-            ("inline", "crossline"),
+            {"inline": 345, "crossline": 188},
+            UNITS_NONE,
             data_type.newbyteorder("native"),  # mdio saves with machine endian, spec could be different endian
             range(1, 346),
             get_inline_header_values,
         )
 
         # Validate the trace mask
-        validate_variable(ds, "trace_mask", (345, 188), ("inline", "crossline"), np.bool, None, None)
+        validate_xr_variable(ds, "trace_mask", {"inline": 345, "crossline": 188}, UNITS_NONE, np.bool, None, None)
 
         # validate the amplitude data
-        validate_variable(
+        validate_xr_variable(
             ds,
             "amplitude",
-            (345, 188, 1501),
-            ("inline", "crossline", "time"),
+            {"inline": 345, "crossline": 188, "time": 1501},
+            UNITS_NONE,
             np.float32,
             None,
             None,
diff --git a/tests/integration/testing_helpers.py b/tests/integration/testing_helpers.py
@@ -4,7 +4,8 @@
 
 import numpy as np
 import xarray as xr
-from numpy.typing import DTypeLike
+
+from mdio.builder.schemas.v1.units import AllUnitModel
 
 
 def get_values(arr: xr.DataArray) -> np.ndarray:
@@ -17,35 +18,40 @@ def get_inline_header_values(dataset: xr.Dataset) -> np.ndarray:
     return dataset["inline"].values
 
 
-def validate_variable(  # noqa PLR0913
+def validate_xr_variable(  # noqa PLR0913
     dataset: xr.Dataset,
     name: str,
-    shape: tuple[int, ...],
-    dims: tuple[str, ...],
-    data_type: DTypeLike,
+    dims: dict[int],
+    units: AllUnitModel,
+    data_type: np.dtype,
     expected_values: range | None,
     actual_value_generator: Callable[[xr.DataArray], np.ndarray] | None = None,
 ) -> None:
     """Validate the properties of a variable in an Xarray dataset."""
-    arr = dataset[name]
-    assert shape == arr.shape
-    assert set(dims) == set(arr.dims)
+    v = dataset[name]
+    assert v is not None
+    assert v.sizes == dims
     if hasattr(data_type, "fields") and data_type.fields is not None:
         # The following assertion will fail because of differences in offsets
         # assert data_type == arr.dtype
 
         # Compare field names
         expected_names = list(data_type.names)
-        actual_names = list(arr.dtype.names)
+        actual_names = list(v.dtype.names)
         assert expected_names == actual_names
 
         # Compare field types
         expected_types = [data_type[name] for name in data_type.names]
-        actual_types = [arr.dtype[name] for name in arr.dtype.names]
+        actual_types = [v.dtype[name] for name in v.dtype.names]
         assert expected_types == actual_types
     else:
-        assert data_type == arr.dtype
+        assert data_type == v.dtype
+
+    if units is not None:
+        assert v.attrs == {"unitsV1": units.model_dump(mode="json")}
+    else:
+        assert "unitsV1" not in v.attrs
 
     if expected_values is not None and actual_value_generator is not None:
-        actual_values = actual_value_generator(arr)
+        actual_values = actual_value_generator(v)
         assert np.array_equal(expected_values, actual_values)
