Performance enhancements for geometry calculation

Author: BrianMichell

src/mdio/segy/geometry.py

@@ -631,13 +631,11 @@ def transform(
         else:
             # Type A: shot points are already unique per gun, create 0-based index from unique values
             for line_val in unique_lines:
-                line_idxs = np.where(index_headers[line_field][:] == line_val)[0]
+                line_idxs = np.where(index_headers[line_field][:] == line_val)
                 shot_points = index_headers["shot_point"][line_idxs]
-                unique_shots = np.sort(np.unique(shot_points))
-                # Map each shot_point to its 0-based index
-                shot_to_idx = {sp: i for i, sp in enumerate(unique_shots)}
-                for i, idx in enumerate(line_idxs):
-                    index_headers["shot_index"][idx] = shot_to_idx[shot_points[i]]
+                # np.unique returns sorted values; searchsorted maps each shot_point to its 0-based index
+                unique_shots = np.unique(shot_points)
+                index_headers["shot_index"][line_idxs] = np.searchsorted(unique_shots, shot_points)
 
         return index_headers
 

tests/integration/conftest.py

@@ -382,3 +382,36 @@ def segy_mock_obn_multiline_type_a(fake_segy_tmp: Path) -> Path:
         components=components,
         filename_suffix="multiline_type_a",
     )
+
+
+@pytest.fixture(scope="module")
+def segy_mock_obn_multiline_type_a_sparse(fake_segy_tmp: Path) -> Path:
+    """Generate mock OBN SEG-Y file with Type A geometry and sparse shot points.
+
+    Variant of segy_mock_obn_multiline_type_a using non-contiguous shot point
+    values per gun. Exercises the vectorized Type A shot_index mapping
+    (np.searchsorted over np.unique) on sparse values to ensure indices are
+    assigned positionally within the sorted unique set, not by value.
+    """
+    num_samples = 25
+    receivers = [101, 102]
+    shot_lines = [1, 2]
+    guns = [1, 2]
+    components = [1]
+
+    # Sparse, non-contiguous shot points; same set per gun keeps geometry Type A
+    shot_points_per_gun = {
+        1: [10, 50, 100],
+        2: [10, 50, 100],
+    }
+
+    return create_segy_mock_obn(
+        fake_segy_tmp,
+        num_samples=num_samples,
+        receivers=receivers,
+        shot_lines=shot_lines,
+        guns=guns,
+        shot_points_per_gun=shot_points_per_gun,
+        components=components,
+        filename_suffix="multiline_type_a_sparse",
+    )

tests/integration/test_import_obn_grid_overrides.py

@@ -6,6 +6,7 @@
 from typing import TYPE_CHECKING
 
 import dask
+import numpy as np
 import pytest
 import xarray.testing as xrt
 from tests.integration.conftest import get_segy_mock_obn_spec
@@ -223,3 +224,39 @@ def test_import_obn_multiline_type_a_all_lines_processed(
         # Verify shot_point is preserved as a coordinate
         assert "shot_point" in ds.coords
         assert ds["shot_point"].dims == ("shot_line", "gun", "shot_index")
+
+    def test_import_obn_multiline_type_a_sparse_shot_points(
+        self,
+        segy_mock_obn_multiline_type_a_sparse: Path,
+        zarr_tmp: Path,
+    ) -> None:
+        """Test Type A shot_index mapping with sparse, non-contiguous shot points.
+
+        Guards the vectorized Type A path (np.searchsorted over np.unique) by
+        using shot points that are not 0/1-based or contiguous. shot_index must
+        be a dense 0-based sequence over the sorted unique shot points, and the
+        original shot_point values must be preserved as a coordinate.
+        """
+        segy_spec = get_segy_mock_obn_spec(include_component=True)
+        grid_override = {"CalculateShotIndex": True}
+
+        segy_to_mdio(
+            segy_spec=segy_spec,
+            mdio_template=TemplateRegistry().get("ObnReceiverGathers3D"),
+            input_path=segy_mock_obn_multiline_type_a_sparse,
+            output_path=zarr_tmp,
+            overwrite=True,
+            grid_overrides=grid_override,
+        )
+
+        ds = open_mdio(zarr_tmp)
+
+        # Sparse shot points [10, 50, 100] map to dense 0-based indices [0, 1, 2]
+        expected_shot_index = [0, 1, 2]
+        xrt.assert_duckarray_equal(ds["shot_index"], expected_shot_index)
+
+        # Original shot_point values preserved as coordinate, not remapped
+        assert "shot_point" in ds.coords
+        assert ds["shot_point"].dims == ("shot_line", "gun", "shot_index")
+        unique_shot_points = np.unique(ds["shot_point"].values)
+        xrt.assert_duckarray_equal(unique_shot_points, [10, 50, 100])