Checkpoint

Author: BrianMichell

src/mdio/converters/segy.py

@@ -178,6 +178,17 @@ def _scan_for_headers(
         )
         template._dim_names = actual_spatial_dims + (template.trace_domain,)
 
+    # Handle NonBinned: move non-binned dimensions to coordinates
+    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",
+                non_binned_dims,
+            )
+            # Add non-binned dimensions as logical coordinates
+            template._logical_coord_names = template._logical_coord_names + non_binned_dims
+
     return segy_dimensions, segy_headers
 
 

src/mdio/segy/geometry.py

@@ -421,19 +421,51 @@ def transform_chunksize(
 
 
 class NonBinned(DuplicateIndex):
-    """Automatically index traces in a single specified axis - trace."""
+    """Handle non-binned dimensions by converting them to a trace dimension with coordinates.
+
+    This override takes dimensions that are not regularly sampled (non-binned) and converts
+    them into a single 'trace' dimension. The original non-binned dimensions become coordinates
+    indexed by the trace dimension.
+
+    Example:
+        Template with dimensions [shot_point, cable, channel, azimuth, offset, sample]
+        and non_binned_dims=['azimuth', 'offset'] becomes:
+        - dimensions: [shot_point, cable, channel, trace, sample]
+        - coordinates: azimuth and offset with dimensions [shot_point, cable, channel, trace]
+
+    Attributes:
+        required_keys: No required keys for this override.
+        required_parameters: Set containing 'chunksize' and 'non_binned_dims'.
+    """
 
     required_keys = None
-    required_parameters = {"chunksize"}
+    required_parameters = {"chunksize", "non_binned_dims"}
+
+    def validate(self, index_headers: HeaderArray, grid_overrides: dict[str, bool | int]) -> None:
+        """Validate if this transform should run on the type of data."""
+        self.check_required_params(grid_overrides)
+
+        # Validate that non_binned_dims is a list
+        non_binned_dims = grid_overrides.get("non_binned_dims", [])
+        if not isinstance(non_binned_dims, list):
+            msg = f"non_binned_dims must be a list, got {type(non_binned_dims)}"
+            raise ValueError(msg)
+
+        # Validate that all non-binned dimensions exist in headers
+        missing_dims = set(non_binned_dims) - set(index_headers.dtype.names)
+        if missing_dims:
+            msg = f"Non-binned dimensions {missing_dims} not found in index headers"
+            raise ValueError(msg)
 
     def transform_chunksize(
         self,
         chunksize: Sequence[int],
         grid_overrides: dict[str, bool | int],
     ) -> Sequence[int]:
-        """Perform the transform of chunksize."""
+        """Insert chunksize for trace dimension at N-1 position."""
         new_chunks = list(chunksize)
-        new_chunks.insert(-1, grid_overrides["chunksize"])
+        trace_chunksize = grid_overrides["chunksize"]
+        new_chunks.insert(-1, trace_chunksize)
         return tuple(new_chunks)
 
 
@@ -544,6 +576,9 @@ def get_allowed_parameters(self) -> set:
 
             parameters.update(command.required_parameters)
 
+        # Add optional parameters that are not strictly required but are valid
+        parameters.add("non_binned_dims")
+
         return parameters
 
     def run(

src/mdio/segy/utilities.py

@@ -73,11 +73,19 @@ def get_grid_plan(  # noqa:  C901, PLR0913
         grid_overrides=grid_overrides,
         template=template,
     )
+
+    # Determine which dimensions are non-binned (converted to coordinates)
+    non_binned_dims = set()
+    if "NonBinned" in grid_overrides and "non_binned_dims" in grid_overrides:
+        non_binned_dims = set(grid_overrides["non_binned_dims"])
+
     # 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)
+    # 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"
+        name
+        for name in horizontal_coordinates
+        if (name in horizontal_dimensions or name == "trace") and name not in non_binned_dims
     ]
 
     dimensions = []

tests/integration/test_import_streamer_grid_overrides.py

@@ -27,9 +27,6 @@
 os.environ["MDIO__IMPORT__SAVE_SEGY_FILE_HEADER"] = "true"
 
 
-# TODO(BrianMichell): Add non-binned back
-# https://github.com/TGSAI/mdio-python/issues/612
-# @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:
@@ -69,13 +66,18 @@ def test_import_4d_segy(  # noqa: PLR0913
         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")
+        # HasDuplicates 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
+        trace_chunks = ds["amplitude"].chunksizes.get("trace", None)
+        if trace_chunks is not None:
+            assert all(chunk == 1 for chunk in trace_chunks)
+
 
 @pytest.mark.parametrize("grid_override", [{"AutoChannelWrap": True}, None])
 @pytest.mark.parametrize("chan_header_type", [StreamerShotGeometryType.A, StreamerShotGeometryType.B])