Fully functional demo

Author: BrianMichell

src/mdio/builder/dataset_builder.py

@@ -150,12 +150,32 @@ def add_coordinate(  # noqa: PLR0913
             msg = "Adding coordinate with the same name twice is not allowed"
             raise ValueError(msg)
 
-        # Validate that all referenced dimensions are already defined
+        # Resolve referenced dimensions strictly, except allow a single substitution with 'trace' if present.
         named_dimensions = []
+        trace_dim = _get_named_dimension(self._dimensions, "trace")
+        resolved_dim_names: list[str] = []
+        trace_used = False
+        missing_dims: list[str] = []
         for dim_name in dimensions:
             nd = _get_named_dimension(self._dimensions, dim_name)
+            if nd is not None:
+                if dim_name not in resolved_dim_names:
+                    resolved_dim_names.append(dim_name)
+                continue
+            if trace_dim is not None and not trace_used and "trace" not in resolved_dim_names:
+                resolved_dim_names.append("trace")
+                trace_used = True
+            else:
+                missing_dims.append(dim_name)
+
+        if missing_dims:
+            msg = f"Pre-existing dimension named {missing_dims[0]!r} is not found"
+            raise ValueError(msg)
+
+        for resolved_name in resolved_dim_names:
+            nd = _get_named_dimension(self._dimensions, resolved_name)
             if nd is None:
-                msg = f"Pre-existing dimension named {dim_name!r} is not found"
+                msg = f"Pre-existing dimension named {resolved_name!r} is not found"
                 raise ValueError(msg)
             named_dimensions.append(nd)
 
@@ -174,7 +194,7 @@ def add_coordinate(  # noqa: PLR0913
         self.add_variable(
             name=coord.name,
             long_name=coord.long_name,
-            dimensions=dimensions,  # dimension names (list[str])
+            dimensions=tuple(resolved_dim_names),  # resolved dimension names
             data_type=coord.data_type,
             compressor=compressor,
             coordinates=[name],  # Use the coordinate name as a reference

src/mdio/builder/templates/base.py

@@ -89,6 +89,21 @@ def build_dataset(
         self._builder = MDIODatasetBuilder(name=name, attributes=attributes)
         self._add_dimensions()
         self._add_coordinates()
+        # Ensure any coordinates declared on the template but not added by subclass overrides
+        # are materialized with generic defaults. This keeps templates override-agnostic while
+        # allowing runtime-augmented coordinate lists to be respected.
+        for coord_name in self.coordinate_names:
+            try:
+                self._builder.add_coordinate(
+                    name=coord_name,
+                    dimensions=self.spatial_dimension_names,
+                    data_type=ScalarType.FLOAT64,
+                    compressor=compressors.Blosc(cname=compressors.BloscCname.zstd),
+                    metadata=CoordinateMetadata(units_v1=self.get_unit_by_key(coord_name)),
+                )
+            except ValueError as exc:  # coordinate may already exist from subclass override
+                if "same name twice" not in str(exc):
+                    raise
         self._add_variables()
         self._add_trace_mask()
 

src/mdio/converters/segy.py

@@ -178,16 +178,19 @@ def _scan_for_headers(
         )
         template._dim_names = actual_spatial_dims + (template.trace_domain,)
 
-    # Handle NonBinned: move non-binned dimensions to coordinates
+    # If using NonBinned override, expose non-binned dims as logical coordinates on the template instance
     if grid_overrides and "NonBinned" in grid_overrides and "non_binned_dims" in grid_overrides:
         non_binned_dims = tuple(grid_overrides["non_binned_dims"])
         if non_binned_dims:
             logger.debug(
-                "NonBinned grid override: moving dimensions %s to coordinates",
+                "NonBinned grid override: exposing non-binned dims as coordinates: %s",
                 non_binned_dims,
             )
-            # Add non-binned dimensions as logical coordinates
-            template._logical_coord_names = template._logical_coord_names + non_binned_dims
+            # Append any missing names; keep existing order and avoid duplicates
+            existing = set(template.coordinate_names)
+            to_add = tuple(n for n in non_binned_dims if n not in existing)
+            if to_add:
+                template._logical_coord_names = template._logical_coord_names + to_add
 
     return segy_dimensions, segy_headers
 

src/mdio/segy/utilities.py

@@ -82,11 +82,28 @@ def get_grid_plan(  # noqa:  C901, PLR0913
     # Use the spatial dimension names from horizontal_coordinates (which may have been modified by grid overrides)
     # Extract only the dimension names (not including non-dimension coordinates or non-binned dimensions)
     # After grid overrides, trace might have been added to horizontal_coordinates
-    transformed_spatial_dims = [
-        name
-        for name in horizontal_coordinates
-        if (name in horizontal_dimensions or name == "trace") and name not in non_binned_dims
-    ]
+    # Compute transformed spatial dims: drop non-binned dims, insert trace if present in headers
+    transformed_spatial_dims = []
+    for name in horizontal_coordinates:
+        if name in non_binned_dims:
+            continue
+        if name == "trace" or name in horizontal_dimensions:
+            transformed_spatial_dims.append(name)
+
+    # Recompute chunksize to match transformed dims
+    original_spatial_dims = list(template.spatial_dimension_names)
+    original_chunks = list(template.full_chunk_shape)
+    new_spatial_chunks: list[int] = []
+    # Insert trace chunk size at N-1 when present, otherwise map remaining dims
+    for dim_name in transformed_spatial_dims:
+        if dim_name == "trace":
+            chunk_val = int(grid_overrides.get("chunksize", 1)) if "NonBinned" in grid_overrides else 1
+            new_spatial_chunks.append(chunk_val)
+            continue
+        if dim_name in original_spatial_dims:
+            idx = original_spatial_dims.index(dim_name)
+            new_spatial_chunks.append(original_chunks[idx])
+    chunksize = tuple(new_spatial_chunks + [original_chunks[-1]])
 
     dimensions = []
     for dim_name in transformed_spatial_dims:

tests/integration/test_import_streamer_grid_overrides.py

@@ -27,7 +27,8 @@
 os.environ["MDIO__IMPORT__SAVE_SEGY_FILE_HEADER"] = "true"
 
 
-@pytest.mark.parametrize("grid_override", [{"HasDuplicates": True}])
+# @pytest.mark.parametrize("grid_override", [{"NonBinned": True, "chunksize": 4, "non_binned_dims": ["channel"]}, {"HasDuplicates": True}])
+@pytest.mark.parametrize("grid_override", [{"NonBinned": True, "chunksize": 4, "non_binned_dims": ["channel"]}])
 @pytest.mark.parametrize("chan_header_type", [StreamerShotGeometryType.C])
 class TestImport4DNonReg:
     """Test for 4D segy import with grid overrides."""
@@ -43,11 +44,14 @@ def test_import_4d_segy(  # noqa: PLR0913
         segy_spec: SegySpec = get_segy_mock_4d_spec()
         segy_path = segy_mock_4d_shots[chan_header_type]
 
+        path = "tmp.mdio"
+        print(f"Running test with grid override: {grid_override}")
+
         segy_to_mdio(
             segy_spec=segy_spec,
             mdio_template=TemplateRegistry().get("PreStackShotGathers3DTime"),
             input_path=segy_path,
-            output_path=zarr_tmp,
+            output_path=path,
             overwrite=True,
             grid_overrides=grid_override,
         )
@@ -58,25 +62,31 @@ def test_import_4d_segy(  # noqa: PLR0913
         cables = [0, 101, 201, 301]
         receivers_per_cable = [1, 5, 7, 5]
 
-        ds = open_mdio(zarr_tmp)
+        ds = open_mdio(path)
 
         assert ds["segy_file_header"].attrs["binaryHeader"]["samples_per_trace"] == num_samples
         assert ds.attrs["attributes"]["gridOverrides"] == grid_override
 
         xrt.assert_duckarray_equal(ds["shot_point"], shots)
         xrt.assert_duckarray_equal(ds["cable"], cables)
 
-        # HasDuplicates should create a trace dimension
+        # Both HasDuplicates and NonBinned should create a trace dimension
         expected = list(range(1, np.amax(receivers_per_cable) + 1))
         xrt.assert_duckarray_equal(ds["trace"], expected)
 
         times_expected = list(range(0, num_samples, 1))
         xrt.assert_duckarray_equal(ds["time"], times_expected)
 
-        # HasDuplicates uses chunksize of 1 for trace dimension
+        # Check trace chunk size based on grid override
         trace_chunks = ds["amplitude"].chunksizes.get("trace", None)
         if trace_chunks is not None:
-            assert all(chunk == 1 for chunk in trace_chunks)
+            if "NonBinned" in grid_override:
+                # NonBinned uses specified chunksize for trace dimension
+                expected_chunksize = grid_override.get("chunksize", 1)
+                assert all(chunk == expected_chunksize for chunk in trace_chunks)
+            else:
+                # HasDuplicates uses chunksize of 1 for trace dimension
+                assert all(chunk == 1 for chunk in trace_chunks)
 
 
 @pytest.mark.parametrize("grid_override", [{"AutoChannelWrap": True}, None])

tests/unit/test_segy_grid_overrides.py

@@ -103,10 +103,10 @@ def test_duplicates(self, mock_streamer_headers: dict[str, npt.NDArray]) -> None
     def test_non_binned(self, mock_streamer_headers: dict[str, npt.NDArray]) -> None:
         """Test the NonBinned Grid Override command."""
         index_names = ("shot_point", "cable")
-        grid_overrides = {"NonBinned": True, "chunksize": 4}
+        grid_overrides = {"NonBinned": True, "chunksize": 4, "non_binned_dims": ["channel"]}
 
-        # Remove channel header
-        streamer_headers = mock_streamer_headers[list(index_names)]
+        # Keep channel header for non-binned processing
+        streamer_headers = mock_streamer_headers
         chunksize = (4, 4, 8)
 
         new_headers, new_names, new_chunks = run_override(
@@ -123,7 +123,9 @@ def test_non_binned(self, mock_streamer_headers: dict[str, npt.NDArray]) -> None
 
         assert_array_equal(dims[0].coords, SHOTS)
         assert_array_equal(dims[1].coords, CABLES)
-        assert_array_equal(dims[2].coords, RECEIVERS)
+        # Trace coords are the unique channel values (1-20)
+        expected_trace_coords = np.arange(1, 21, dtype="int32")
+        assert_array_equal(dims[2].coords, expected_trace_coords)
 
 
 class TestStreamerGridOverrides: