Add test to cover prestack template.

Author: markspec

src/mdio/builder/template_registry.py

@@ -11,6 +11,7 @@
 from mdio.builder.templates.seismic_3d_prestack_cdp import Seismic3DPreStackCDPTemplate
 from mdio.builder.templates.seismic_3d_prestack_coca import Seismic3DPreStackCocaTemplate
 from mdio.builder.templates.seismic_3d_prestack_shot import Seismic3DPreStackShotTemplate
+from mdio.builder.templates.seismic_prestack import SeismicPreStackTemplate
 
 
 class TemplateRegistry:
@@ -85,6 +86,7 @@ def _register_default_templates(self) -> None:
         self.register(Seismic3DPreStackCocaTemplate("depth"))
 
         # Field (shot) data
+        self.register(SeismicPreStackTemplate("time"))
         self.register(Seismic2DPreStackShotTemplate("time"))
         self.register(Seismic3DPreStackShotTemplate("time"))
 

src/mdio/builder/templates/seismic_prestack.py

@@ -1,8 +1,11 @@
 """SeismicPreStackTemplate MDIO v1 dataset templates."""
 
-from mdio.schemas.dtype import ScalarType
-from mdio.schemas.metadata import UserAttributes
-from mdio.schemas.v1.templates.abstract_dataset_template import AbstractDatasetTemplate
+from typing import Any
+
+from mdio.builder.schemas.dtype import ScalarType
+from mdio.builder.schemas.v1.variable import CoordinateMetadata
+from mdio.builder.templates.abstract_dataset_template import AbstractDatasetTemplate
+from mdio.builder.templates.types import SeismicDataDomain
 
 
 class SeismicPreStackTemplate(AbstractDatasetTemplate):
@@ -14,8 +17,8 @@ class SeismicPreStackTemplate(AbstractDatasetTemplate):
         domain: The domain of the dataset.
     """
 
-    def __init__(self, domain: str = "time"):
-        super().__init__(domain=domain)
+    def __init__(self, data_domain: SeismicDataDomain):
+        super().__init__(data_domain=data_domain)
 
         self._coord_dim_names = [
             "shot_line",
@@ -24,54 +27,59 @@ def __init__(self, domain: str = "time"):
             "cable",
             "channel",
         ]  # Custom coordinates for shot gathers
-        self._dim_names = [*self._coord_dim_names, self._trace_domain]
-        self._coord_names = ["source_coord_x", "source_coord_y", "group_coord_x", "group_coord_y"]
+        self._dim_names = [*self._coord_dim_names, self._data_domain]
+        self._coord_names = [
+            "energy_source_point_number",
+            "source_coord_x",
+            "source_coord_y",
+            "group_coord_x",
+            "group_coord_y",
+        ]
         self._var_chunk_shape = [1, 1, 16, 1, 32, -1]
 
     @property
     def _name(self) -> str:
-        return f"PreStackGathers3D{self._trace_domain.capitalize()}"
+        return f"PreStackGathers3D{self._data_domain.capitalize()}"
 
-    def _load_dataset_attributes(self) -> UserAttributes:
-        return UserAttributes(
-            attributes={
-                "surveyDimensionality": "3D",
-                "ensembleType": "shot_point",
-                "processingStage": "pre-stack",
-            }
-        )
+    def _load_dataset_attributes(self) -> dict[str, Any]:
+        return {
+            "surveyDimensionality": "3D",
+            "ensembleType": "shot_point",
+            "processingStage": "pre-stack",
+        }
 
     def _add_coordinates(self) -> None:
         # Add dimension coordinates
         for name in self._dim_names:
-            self._builder.add_coordinate(
-                name,
-                dimensions=[name],
-                data_type=ScalarType.INT32,
-                metadata_info=None,
-            )
+            self._builder.add_coordinate(name, dimensions=(name,), data_type=ScalarType.INT32)
 
+        # Add non-dimension coordinates
+        self._builder.add_coordinate(
+            "energy_source_point_number",
+            dimensions=("shot_line", "gun", "shot_point"),
+            data_type=ScalarType.INT32,
+        )
         self._builder.add_coordinate(
             "source_coord_x",
-            dimensions=["shot_line", "gun", "shot_point", "cable", "channel"],
+            dimensions=("shot_line", "gun", "shot_point"),
             data_type=ScalarType.FLOAT64,
-            metadata_info=[self._horizontal_coord_unit],
+            metadata=CoordinateMetadata(units_v1=self._horizontal_coord_unit),
         )
         self._builder.add_coordinate(
             "source_coord_y",
-            dimensions=["shot_line", "gun", "shot_point", "cable", "channel"],
+            dimensions=("shot_line", "gun", "shot_point"),
             data_type=ScalarType.FLOAT64,
-            metadata_info=[self._horizontal_coord_unit],
+            metadata=CoordinateMetadata(units_v1=self._horizontal_coord_unit),
         )
         self._builder.add_coordinate(
             "group_coord_x",
-            dimensions=["shot_line", "gun", "shot_point", "cable", "channel"],
+            dimensions=("shot_line", "gun", "shot_point", "cable", "channel"),
             data_type=ScalarType.FLOAT64,
-            metadata_info=[self._horizontal_coord_unit],
+            metadata=CoordinateMetadata(units_v1=self._horizontal_coord_unit),
         )
         self._builder.add_coordinate(
             "group_coord_y",
-            dimensions=["shot_line", "gun", "shot_point", "cable", "channel"],
+            dimensions=("shot_line", "gun", "shot_point", "cable", "channel"),
             data_type=ScalarType.FLOAT64,
-            metadata_info=[self._horizontal_coord_unit],
+            metadata=CoordinateMetadata(units_v1=self._horizontal_coord_unit),
         )

tests/unit/v1/templates/test_seismic_prestack.py

@@ -0,0 +1,207 @@
+"""Unit tests for SeismicPreStackTemplate."""
+
+import pytest
+from tests.unit.v1.helpers import validate_variable
+
+from mdio.builder.schemas.chunk_grid import RegularChunkGrid
+from mdio.builder.schemas.compressors import Blosc
+from mdio.builder.schemas.compressors import BloscCname
+from mdio.builder.schemas.dtype import ScalarType
+from mdio.builder.schemas.dtype import StructuredType
+from mdio.builder.schemas.v1.dataset import Dataset
+from mdio.builder.schemas.v1.units import LengthUnitEnum
+from mdio.builder.schemas.v1.units import LengthUnitModel
+from mdio.builder.templates.seismic_prestack import SeismicPreStackTemplate
+
+UNITS_METER = LengthUnitModel(length=LengthUnitEnum.METER)
+
+
+def _validate_coordinates_headers_trace_mask(dataset: Dataset, headers: StructuredType, domain: str) -> None:
+    """Validate the coordinate, headers, trace_mask variables in the dataset."""
+    # Verify variables
+    # 6 dim coords + 5 non-dim coords + 1 data + 1 trace mask + 1 headers = 14 variables
+    assert len(dataset.variables) == 14
+
+    # Verify trace headers
+    validate_variable(
+        dataset,
+        name="headers",
+        dims=[("shot_line", 1), ("gun", 3), ("shot_point", 256), ("cable", 512), ("channel", 24)],
+        coords=["energy_source_point_number", "source_coord_x", "source_coord_y", "group_coord_x", "group_coord_y"],
+        dtype=headers,
+    )
+
+    validate_variable(
+        dataset,
+        name="trace_mask",
+        dims=[("shot_line", 1), ("gun", 3), ("shot_point", 256), ("cable", 512), ("channel", 24)],
+        coords=["energy_source_point_number", "source_coord_x", "source_coord_y", "group_coord_x", "group_coord_y"],
+        dtype=ScalarType.BOOL,
+    )
+
+    # Verify dimension coordinate variables
+    shot_line = validate_variable(
+        dataset,
+        name="shot_line",
+        dims=[("shot_line", 1)],
+        coords=["shot_line"],
+        dtype=ScalarType.INT32,
+    )
+    assert shot_line.metadata is None
+
+    gun = validate_variable(
+        dataset,
+        name="gun",
+        dims=[("gun", 3)],
+        coords=["gun"],
+        dtype=ScalarType.INT32,
+    )
+    assert gun.metadata is None
+
+    shot_point = validate_variable(
+        dataset,
+        name="shot_point",
+        dims=[("shot_point", 256)],
+        coords=["shot_point"],
+        dtype=ScalarType.INT32,
+    )
+    assert shot_point.metadata is None
+
+    cable = validate_variable(
+        dataset,
+        name="cable",
+        dims=[("cable", 512)],
+        coords=["cable"],
+        dtype=ScalarType.INT32,
+    )
+    assert cable.metadata is None
+
+    channel = validate_variable(
+        dataset,
+        name="channel",
+        dims=[("channel", 24)],
+        coords=["channel"],
+        dtype=ScalarType.INT32,
+    )
+    assert channel.metadata is None
+
+    domain_var = validate_variable(
+        dataset,
+        name=domain,
+        dims=[(domain, 2048)],
+        coords=[domain],
+        dtype=ScalarType.INT32,
+    )
+    assert domain_var.metadata is None
+
+    # Verify non-dimension coordinate variables
+    validate_variable(
+        dataset,
+        name="energy_source_point_number",
+        dims=[("shot_line", 1), ("gun", 3), ("shot_point", 256)],
+        coords=["energy_source_point_number"],
+        dtype=ScalarType.INT32,
+    )
+
+    source_coord_x = validate_variable(
+        dataset,
+        name="source_coord_x",
+        dims=[("shot_line", 1), ("gun", 3), ("shot_point", 256)],
+        coords=["source_coord_x"],
+        dtype=ScalarType.FLOAT64,
+    )
+    assert source_coord_x.metadata.units_v1.length == LengthUnitEnum.METER
+
+    source_coord_y = validate_variable(
+        dataset,
+        name="source_coord_y",
+        dims=[("shot_line", 1), ("gun", 3), ("shot_point", 256)],
+        coords=["source_coord_y"],
+        dtype=ScalarType.FLOAT64,
+    )
+    assert source_coord_y.metadata.units_v1.length == LengthUnitEnum.METER
+
+    group_coord_x = validate_variable(
+        dataset,
+        name="group_coord_x",
+        dims=[("shot_line", 1), ("gun", 3), ("shot_point", 256), ("cable", 512), ("channel", 24)],
+        coords=["group_coord_x"],
+        dtype=ScalarType.FLOAT64,
+    )
+    assert group_coord_x.metadata.units_v1.length == LengthUnitEnum.METER
+
+    group_coord_y = validate_variable(
+        dataset,
+        name="group_coord_y",
+        dims=[("shot_line", 1), ("gun", 3), ("shot_point", 256), ("cable", 512), ("channel", 24)],
+        coords=["group_coord_y"],
+        dtype=ScalarType.FLOAT64,
+    )
+    assert group_coord_y.metadata.units_v1.length == LengthUnitEnum.METER
+
+
+class TestSeismic3DPreStackShotTemplate:
+    """Unit tests for SeismicPreStackTemplate."""
+
+    def test_configuration(self) -> None:
+        """Unit tests for SeismicPreStackTemplate in time domain."""
+        t = SeismicPreStackTemplate(data_domain="time")
+
+        # Template attributes for prestack shot
+        assert t._data_domain == "time"
+        assert t._coord_dim_names == ["shot_line", "gun", "shot_point", "cable", "channel"]
+        assert t._dim_names == ["shot_line", "gun", "shot_point", "cable", "channel", "time"]
+        assert t._coord_names == ["energy_source_point_number", "source_coord_x", "source_coord_y", "group_coord_x", "group_coord_y"]
+        assert t._var_chunk_shape == [1, 1, 16, 1, 32, -1]
+
+        # Variables instantiated when build_dataset() is called
+        assert t._builder is None
+        assert t._dim_sizes == ()
+        assert t._horizontal_coord_unit is None
+
+        # Verify prestack shot attributes
+        attrs = t._load_dataset_attributes()
+        assert attrs == {"surveyDimensionality": "3D", "ensembleType": "shot_point", "processingStage": "pre-stack"}
+
+        assert t.name == "PreStackGathers3DTime"
+
+    def test_build_dataset(self, structured_headers: StructuredType) -> None:
+        """Unit tests for SeismicPreStackTemplate build in time domain."""
+        t = SeismicPreStackTemplate(data_domain="time")
+
+        assert t.name == "PreStackGathers3DTime"
+        dataset = t.build_dataset(
+            "North Sea 3D Shot Time",
+            sizes=(1, 3, 256, 512, 24, 2048),
+            horizontal_coord_unit=UNITS_METER,
+            header_dtype=structured_headers,
+        )
+
+        assert dataset.metadata.name == "North Sea 3D Shot Time"
+        assert dataset.metadata.attributes["surveyDimensionality"] == "3D"
+        assert dataset.metadata.attributes["ensembleType"] == "shot_point"
+        assert dataset.metadata.attributes["processingStage"] == "pre-stack"
+
+        _validate_coordinates_headers_trace_mask(dataset, structured_headers, "time")
+
+        # Verify seismic variable (prestack shot time data)
+        seismic = validate_variable(
+            dataset,
+            name="amplitude",
+            dims=[("shot_line", 1), ("gun", 3), ("shot_point", 256), ("cable", 512), ("channel", 24), ("time", 2048)],
+            coords=["energy_source_point_number", "source_coord_x", "source_coord_y", "group_coord_x", "group_coord_y"],
+            dtype=ScalarType.FLOAT32,
+        )
+        assert isinstance(seismic.compressor, Blosc)
+        assert seismic.compressor.cname == BloscCname.zstd
+        assert isinstance(seismic.metadata.chunk_grid, RegularChunkGrid)
+        assert seismic.metadata.chunk_grid.configuration.chunk_shape == (1, 1, 16, 1, 32, -1)
+        assert seismic.metadata.stats_v1 is None
+
+
+@pytest.mark.parametrize("data_domain", ["Time", "TiME"])
+def test_domain_case_handling(data_domain: str) -> None:
+    """Test that domain parameter handles different cases correctly."""
+    template = SeismicPreStackTemplate(data_domain=data_domain)
+    assert template._data_domain == data_domain.lower()
+    assert template.name.endswith(data_domain.capitalize())