[ESSREDUCE] fix: support rectilinear grids in the numba interpolator

packages/essreduce/src/ess/reduce/unwrap/interpolator_numba.py

@@ -4,6 +4,21 @@
 from numba import njit, prange
 
 
+@njit(boundscheck=False, cache=True)
+def _locate(grid: np.ndarray, value: float, one_over_step: float, uniform: bool) -> int:
+    """Index of the cell of ``grid`` containing ``value``.
+
+    ``value`` is assumed to lie within the grid. A value on the upper edge is
+    assigned to the last cell, so that the caller can index ``grid[i + 1]``.
+    """
+    n = len(grid)
+    if value == grid[n - 1]:
+        return n - 2
+    if uniform:
+        return int((value - grid[0]) * one_over_step)
+    return np.searchsorted(grid, value, side='right') - 1
+
+
 @njit(boundscheck=False, cache=True, fastmath=False, parallel=True)
 def interpolate(
     x: np.ndarray,
@@ -14,17 +29,19 @@ def interpolate(
     xoffset: np.ndarray | None,
     deltax: float,
     fill_value: float,
+    x_uniform: bool,
+    y_uniform: bool,
     out: np.ndarray,
 ):
     """
-    Linear interpolation of data on a 2D regular grid.
+    Linear interpolation of data on a 2D rectilinear grid.
 
     Parameters
     ----------
     x:
-        1D array of grid edges along the x-axis (size nx). They must be linspaced.
+        1D array of grid points along the x-axis (size nx), strictly increasing.
     y:
-        1D array of grid edges along the y-axis (size ny). They must be linspaced.
+        1D array of grid points along the y-axis (size ny), strictly increasing.
     values:
         2D array of values on the grid. The shape must be (ny, nx).
     xp:
@@ -37,6 +54,11 @@ def interpolate(
         Multiplier to apply to the integer offsets (i.e. the step size).
     fill_value:
         Value to use for points outside of the grid.
+    x_uniform:
+        Whether ``x`` is equally spaced, which allows the containing cell to be
+        computed by division instead of searched for. See :class:`Interpolator`.
+    y_uniform:
+        Whether ``y`` is equally spaced.
     out:
         1D array where the interpolated values will be stored (size N).
     """
@@ -50,11 +72,10 @@ def interpolate(
     xmax = x[nx - 1]
     ymin = y[0]
     ymax = y[ny - 1]
-    dx = x[1] - xmin
-    dy = y[1] - ymin
 
-    one_over_dx = 1.0 / dx
-    one_over_dy = 1.0 / dy
+    one_over_dx = 1.0 / (x[1] - xmin)
+    one_over_dy = 1.0 / (y[1] - ymin)
+    both_uniform = x_uniform and y_uniform
     norm = one_over_dx * one_over_dy
 
     for i in prange(npoints):
@@ -65,8 +86,8 @@ def interpolate(
             out[i] = fill_value
 
         else:
-            ix = nx - 2 if xx == xmax else int((xx - xmin) * one_over_dx)
-            iy = ny - 2 if yy == ymax else int((yy - ymin) * one_over_dy)
+            ix = _locate(x, xx, one_over_dx, x_uniform)
+            iy = _locate(y, yy, one_over_dy, y_uniform)
 
             x1 = x[ix]
             x2 = x[ix + 1]
@@ -81,10 +102,22 @@ def interpolate(
             x2mxx = x2 - xx
             xxmx1 = xx - x1
 
+            # A uniform grid normalizes by the same cell area everywhere, which
+            # is worth hoisting out of the loop; a rectilinear one does not.
+            cell = norm if both_uniform else 1.0 / ((x2 - x1) * (y2 - y1))
+
             out[i] = (
                 (y2 - yy) * (x2mxx * a11 + xxmx1 * a21)
                 + (yy - y1) * (x2mxx * a12 + xxmx1 * a22)
-            ) * norm
+            ) * cell
+
+
+def _is_uniform(grid: np.ndarray) -> bool:
+    """Whether ``grid`` is equally spaced, to within floating-point noise."""
+    if len(grid) < 3:
+        return True
+    steps = np.diff(grid)
+    return bool(np.allclose(steps, steps[0], rtol=1.0e-9, atol=0.0))
 
 
 class Interpolator:
@@ -96,23 +129,40 @@ def __init__(
         fill_value: float = np.nan,
     ):
         """
-        Interpolator for 2D regular grid data (Numba implementation).
+        Interpolator for 2D rectilinear grid data (Numba implementation).
+
+        The axes need not be equally spaced: a lookup table may sample distance
+        densely where components sit and not at all in between. Uniformity is
+        detected here, once, because it decides how the containing cell is
+        found — by division for a uniform axis, by binary search otherwise —
+        and that is a per-point cost in the interpolation loop.
 
         Parameters
         ----------
         time_edges:
-            1D array of time edges.
+            1D array of time grid points, strictly increasing.
         distance_edges:
-            1D array of distance edges.
+            1D array of distance grid points, strictly increasing.
         values:
             2D array of values on the grid. The shape must be (ny, nx).
         fill_value:
             Value to use for points outside of the grid.
         """
+        for name, grid in (
+            ('time_edges', time_edges),
+            ('distance_edges', distance_edges),
+        ):
+            if len(grid) < 2:
+                raise ValueError(
+                    f"Interpolator: {name} has {len(grid)} point(s); at least two "
+                    "are needed to bracket a value."
+                )
         self.time_edges = time_edges
         self.distance_edges = distance_edges
         self.values = values
         self.fill_value = fill_value
+        self.time_uniform = _is_uniform(time_edges)
+        self.distance_uniform = _is_uniform(distance_edges)
 
     def __call__(
         self,
@@ -131,6 +181,8 @@ def __call__(
             xoffset=pulse_index,
             deltax=pulse_period,
             fill_value=self.fill_value,
+            x_uniform=self.time_uniform,
+            y_uniform=self.distance_uniform,
             out=out,
         )
         return out

packages/essreduce/tests/unwrap/interpolator_test.py

@@ -2,6 +2,7 @@
 # Copyright (c) 2025 Scipp contributors (https://github.com/scipp)
 
 import numpy as np
+import pytest
 
 from ess.reduce.unwrap.interpolator_numba import (
     Interpolator as InterpolatorNumba,
@@ -100,3 +101,85 @@ def test_numba_and_scipy_interpolators_yield_same_results_with_values_on_edges()
     numba_result = numba_interp(times, distances)
     scipy_result = scipy_interp(times, distances)
     assert np.allclose(numba_result, scipy_result, equal_nan=True)
+
+
+def _make_rectilinear_interpolators():
+    """Interpolators on a distance axis that is dense in two places only.
+
+    The shape a lookup table has when it samples the beamline where components
+    sit and not at all in between.
+    """
+    time_edges = np.linspace(0, 71, 101)
+    distance_edges = np.concatenate(
+        [np.linspace(6.4, 6.8, 5), np.linspace(72.0, 72.4, 5)]
+    )
+    time_g, distance_g = np.meshgrid(
+        time_edges, distance_edges, indexing='ij', sparse=True
+    )
+    values = _f(time_g, distance_g).T
+
+    numba_interp = InterpolatorNumba(
+        time_edges=time_edges, distance_edges=distance_edges, values=values
+    )
+    scipy_interp = InterpolatorScipy(
+        time_edges=time_edges, distance_edges=distance_edges, values=values
+    )
+    return numba_interp, scipy_interp
+
+
+def test_numba_and_scipy_interpolators_yield_same_results_on_rectilinear_grid():
+    numba_interp, scipy_interp = _make_rectilinear_interpolators()
+
+    rng = np.random.default_rng(seed=42)
+    npoints = 1000
+    times = rng.uniform(0, 71, npoints)
+    distances = np.concatenate(
+        [
+            rng.uniform(6.4, 6.8, npoints // 2),
+            rng.uniform(72.0, 72.4, npoints - npoints // 2),
+        ]
+    )
+
+    numba_result = numba_interp(times, distances)
+    scipy_result = scipy_interp(times, distances)
+
+    assert np.allclose(numba_result, scipy_result)
+
+
+def test_numba_and_scipy_interpolators_agree_across_a_gap_in_the_grid():
+    # Between the dense regions the grid has one wide cell, which both
+    # implementations interpolate across rather than treating as out of bounds.
+    numba_interp, scipy_interp = _make_rectilinear_interpolators()
+
+    times = np.array([10.0, 35.0, 60.0])
+    distances = np.array([20.0, 40.0, 60.0])
+
+    assert np.allclose(numba_interp(times, distances), scipy_interp(times, distances))
+
+
+def test_uniform_grid_is_detected():
+    # The uniform case takes a division rather than a search to locate a cell,
+    # so it must still be recognized as uniform.
+    numba_interp, _ = _make_interpolators()
+
+    assert numba_interp.time_uniform
+    assert numba_interp.distance_uniform
+
+
+def test_rectilinear_grid_is_not_reported_as_uniform():
+    numba_interp, _ = _make_rectilinear_interpolators()
+
+    assert numba_interp.time_uniform
+    assert not numba_interp.distance_uniform
+
+
+@pytest.mark.parametrize('npoints', [0, 1])
+def test_grid_too_short_to_bracket_a_value_raises(npoints: int):
+    # A single grid point cannot bracket anything: it used to be read past the
+    # end of the array.
+    with pytest.raises(ValueError, match='at least two'):
+        InterpolatorNumba(
+            time_edges=np.linspace(0, 71, 101),
+            distance_edges=np.linspace(40, 70, npoints),
+            values=np.zeros((npoints, 101)),
+        )