Implement working test for DuplicateTraces override

Author: BrianMichell

src/mdio/builder/dataset_builder.py

@@ -134,12 +134,9 @@ def push_dimension(self, dimension: NamedDimension, position: int, new_dim_chunk
         def propogate_dimension(variable: Variable, position: int, new_dim_chunk_size: int) -> Variable:
             """Propogates the dimension to the variable or coordinate."""
             from mdio.builder.schemas.chunk_grid import RegularChunkGrid, RegularChunkShape
-            if len(variable.dimensions) <= position:
+            if len(variable.dimensions) + 1 <= position:
                 # Don't do anything if the new dimension is not within the Variable's domain
                 return variable
-            if variable.name == "trace_mask":
-                # Special case for trace_mask. Don't do anything.
-                return variable
             new_dimensions = variable.dimensions[:position] + [dimension] + variable.dimensions[position:]
 
             # Get current chunk shape from metadata

src/mdio/segy/utilities.py

@@ -26,6 +26,46 @@
 logger = logging.getLogger(__name__)
 
 
+def _create_delayed_trace_dimension_transform(headers_subset: HeaderArray, position: int) -> callable:
+    """Create a delayed transform function that adds a trace dimension and its coordinate.
+    
+    This function creates a closure that captures the headers_subset and position,
+    but defers the actual computation until the transform is executed by the dataset builder.
+    The transform adds both the trace dimension and a corresponding coordinate.
+    
+    Args:
+        headers_subset: The header array containing trace information
+        position: The position where the trace dimension should be inserted
+        
+    Returns:
+        A callable that can be used as a transform function
+    """
+    def delayed_transform(builder):
+        from mdio.builder.schemas.dtype import ScalarType
+        
+        # Calculate the trace dimension size at execution time
+        if "trace" in headers_subset.dtype.names:
+            trace_size = int(np.max(headers_subset["trace"]))
+        else:
+            # Fallback: if trace field doesn't exist, we need to determine size differently
+            raise ValueError("Trace field not found in headers_subset when executing delayed transform")
+        
+        # Add the trace dimension
+        trace_dimension = NamedDimension(name="trace", size=trace_size)
+        builder.push_dimension(trace_dimension, position=position, new_dim_chunk_size=1)
+        
+        # Add the corresponding coordinate for the trace dimension
+        builder.add_coordinate(
+            "trace",
+            dimensions=("trace",),
+            data_type=ScalarType.INT32,
+        )
+        
+        return builder
+    
+    return delayed_transform
+
+
 def get_grid_plan(  # noqa:  C901
     segy_file: SegyFile,
     chunksize: tuple[int, ...] | None,
@@ -70,7 +110,8 @@ def get_grid_plan(  # noqa:  C901
 
     if grid_overrides.get("HasDuplicates", False):
         pos = len(template.dimension_names) - 1  # TODO: Implement the negative position case...
-        template._queue_transform(lambda builder: builder.push_dimension(NamedDimension(name="trace", size=np.max(headers_subset["trace"])), position=pos, new_dim_chunk_size=1))
+        # Use the delayed transform function instead of a simple lambda
+        template._queue_transform(_create_delayed_trace_dimension_transform(headers_subset, pos))
         horizontal_dimensions = (*horizontal_dimensions, "trace")
 
     dimensions = []

tests/integration/test_import_streamer_grid_overrides.py

@@ -52,8 +52,7 @@ def test_import_4d_segy(  # noqa: PLR0913
             segy_spec=segy_spec,
             mdio_template=TemplateRegistry().get("PreStackShotGathers3DTime"),
             input_path=segy_path,
-            # output_path=zarr_tmp,
-            output_path="test_has_duplicates.mdio",
+            output_path=zarr_tmp,
             overwrite=True,
             grid_overrides=grid_override,
         )
@@ -69,10 +68,9 @@ 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")
         expected = list(range(1, np.amax(receivers_per_cable) + 1))
         xrt.assert_duckarray_equal(ds["trace"], expected)