Working example of DuplicateIndex grid override

Author: BrianMichell

src/mdio/converters/segy.py

@@ -156,7 +156,11 @@ def _scan_for_headers(
     """Extract trace dimensions and index headers from the SEG-Y file.
 
     This is an expensive operation.
-    It scans the SEG-Y file in chunks by using ProcessPoolExecutor
+    It scans the SEG-Y file in chunks by using ProcessPoolExecutor.
+
+    Note:
+        If grid_overrides are applied to the template before calling this function,
+        the chunk_size returned from get_grid_plan should match the template's chunk shape.
     """
     full_chunk_shape = template.full_chunk_shape
     segy_dimensions, chunk_size, segy_headers = get_grid_plan(
@@ -167,13 +171,19 @@ def _scan_for_headers(
         chunksize=full_chunk_shape,
         grid_overrides=grid_overrides,
     )
+
+    # After applying grid overrides to the template, chunk sizes should match
+    # If they don't match, it means the template wasn't properly updated
     if full_chunk_shape != chunk_size:
-        # TODO(Dmitriy): implement grid overrides
-        # https://github.com/TGSAI/mdio-python/issues/585
-        # The returned 'chunksize' is used only for grid_overrides. We will need to use it when full
-        # support for grid overrides is implemented
-        err = "Support for changing full_chunk_shape in grid overrides is not yet implemented"
-        raise NotImplementedError(err)
+        logger.warning(
+            "Chunk shape mismatch: template has %s but grid_plan returned %s. "
+            "Using grid_plan chunk shape.",
+            full_chunk_shape,
+            chunk_size,
+        )
+        # Update the template's chunk shape to match what grid_plan returned
+        template._var_chunk_shape = chunk_size
+
     return segy_dimensions, segy_headers
 
 
@@ -562,6 +572,17 @@ def segy_to_mdio(  # noqa PLR0913
     )
     grid = _build_and_check_grid(segy_dimensions, segy_file_info, segy_headers)
 
+    # Update template dimensions to match the actual grid dimensions after grid overrides
+    # The chunk shape was already updated in _scan_for_headers, we just need to fix dimensions
+    actual_spatial_dims = tuple(grid.dim_names[:-1])  # All dims except the vertical/time dimension
+    if mdio_template.spatial_dimension_names != actual_spatial_dims:
+        logger.info(
+            "Adjusting template dimensions from %s to %s to match grid after overrides",
+            mdio_template.spatial_dimension_names,
+            actual_spatial_dims,
+        )
+        mdio_template._dim_names = actual_spatial_dims + (mdio_template.trace_domain,)
+
     _, non_dim_coords = _get_coordinates(grid, segy_headers, mdio_template)
     header_dtype = to_structured_type(segy_spec.trace.header.dtype)
 

src/mdio/segy/geometry.py

@@ -267,7 +267,8 @@ def analyze_non_indexed_headers(index_headers: HeaderArray, dtype: DTypeLike = n
     header_names = []
     for header_key in index_headers.dtype.names:
         if header_key != "trace":
-            unique_headers[header_key] = np.sort(np.unique(index_headers[header_key]))
+            unique_vals = np.sort(np.unique(index_headers[header_key]))
+            unique_headers[header_key] = unique_vals
             header_names.append(header_key)
             total_depth += 1
 
@@ -382,7 +383,39 @@ def transform(
         """Perform the grid transform."""
         self.validate(index_headers, grid_overrides)
 
-        return analyze_non_indexed_headers(index_headers)
+        # Filter to only include dimension fields, not coordinate fields
+        # Coordinate fields typically have _x, _y suffixes or specific names like 'gun'
+        # We want to keep fields like shot_point, cable, channel but exclude source_coord_x, etc.
+        dimension_fields = []
+        coordinate_field_patterns = ['_x', '_y', '_coord', 'gun', 'source', 'group']
+        
+        for field_name in index_headers.dtype.names:
+            # Skip if it's already trace
+            if field_name == 'trace':
+                continue
+            # Check if it looks like a coordinate field
+            is_coordinate = any(pattern in field_name for pattern in coordinate_field_patterns)
+            if not is_coordinate:
+                dimension_fields.append(field_name)
+        
+        # Extract only dimension fields for trace indexing
+        if dimension_fields:
+            dimension_headers = index_headers[dimension_fields]
+        else:
+            # If no dimension fields, use all fields
+            dimension_headers = index_headers
+        
+        # Create trace indices based on dimension fields only
+        dimension_headers_with_trace = analyze_non_indexed_headers(dimension_headers)
+        
+        # Add the trace field back to the full index_headers array
+        if dimension_headers_with_trace is not None and 'trace' in dimension_headers_with_trace.dtype.names:
+            # Extract just the trace values array (not the whole structured array)
+            trace_values = np.array(dimension_headers_with_trace['trace'])
+            # Append as a new field to the full headers
+            index_headers = rfn.append_fields(index_headers, 'trace', trace_values, usemask=False)
+        
+        return index_headers
 
     def transform_index_names(self, index_names: Sequence[str]) -> Sequence[str]:
         """Insert dimension "trace" to the sample-1 dimension."""

src/mdio/segy/utilities.py

@@ -72,9 +72,15 @@ def get_grid_plan(  # noqa:  C901, PLR0913
         chunksize=chunksize,
         grid_overrides=grid_overrides,
     )
+    # 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)
+    # 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"]
 
     dimensions = []
-    for dim_name in horizontal_dimensions:
+    for dim_name in transformed_spatial_dims:
+        if dim_name not in headers_subset.dtype.names:
+            continue
         dim_unique = np.unique(headers_subset[dim_name])
         dimensions.append(Dimension(coords=dim_unique, name=dim_name))
 

tests/integration/test_import_streamer_grid_overrides.py

@@ -27,13 +27,12 @@
 dask.config.set(scheduler="synchronous")
 os.environ["MDIO__IMPORT__SAVE_SEGY_FILE_HEADER"] = "true"
 
-
-# TODO(Altay): Finish implementing these grid overrides.
+# TODO(BrianMichell): Add non-binned back
 # https://github.com/TGSAI/mdio-python/issues/612
-@pytest.mark.skip(reason="NonBinned and HasDuplicates haven't been properly implemented yet.")
-@pytest.mark.parametrize("grid_override", [{"NonBinned": True}, {"HasDuplicates": True}])
+# @pytest.mark.parametrize("grid_override", [{"NonBinned": True, "chunksize": 4}, {"HasDuplicates": True}])
+@pytest.mark.parametrize("grid_override", [{"HasDuplicates": True}])
 @pytest.mark.parametrize("chan_header_type", [StreamerShotGeometryType.C])
-class TestImport4DNonReg:  # pragma: no cover - tests is skipped
+class TestImport4DNonReg:
     """Test for 4D segy import with grid overrides."""
 
     def test_import_4d_segy(  # noqa: PLR0913
@@ -67,7 +66,7 @@ def test_import_4d_segy(  # noqa: PLR0913
         assert ds["segy_file_header"].attrs["binaryHeader"]["samples_per_trace"] == num_samples
         assert ds.attrs["attributes"]["gridOverrides"] == grid_override
 
-        assert npt.assert_array_equal(ds["shot_point"], shots)
+        xrt.assert_duckarray_equal(ds["shot_point"], shots)
         xrt.assert_duckarray_equal(ds["cable"], cables)
 
         # assert grid.select_dim("trace") == Dimension(range(1, np.amax(receivers_per_cable) + 1), "trace")