From 9ab6ac7fab7ca061908f359397fcd17c2bd98bba Mon Sep 17 00:00:00 2001 From: Simon Heybrock Date: Tue, 25 Aug 2026 13:07:18 +0000 Subject: [PATCH] fix: support rectilinear grids in the numba interpolator The numba interpolator located a cell by dividing by a single global step, which assumes an equally spaced grid. On a grid that is not equally spaced it therefore read the wrong cell -- and with boundscheck=False, often past the end of the array: numba: [ 2131.86 54175.32 -1.3e+254 ] scipy: [ 2131.86 54175.32 265739.87 ] The scipy implementation handles such a grid correctly, so which answer you got depended on whether numba was installed. Nothing said the grid had to be equally spaced except a line in the docstring. Uniformity is now detected once per axis when the Interpolator is constructed and passed to the kernel, which divides when the axis is equally spaced and binary-searches when it is not. The uniform path keeps its hoisted cell-area normalization and is unchanged in cost. This matters for a lookup table that samples distance where components sit and not at all in between -- dense across the detectors, a few rows at each monitor -- which spans a beamline in a few tens of rows instead of thousands. Also raises when an axis has fewer than two points. Such a grid cannot bracket a value: it used to be read past the end of the array as well. --- .../ess/reduce/unwrap/interpolator_numba.py | 78 ++++++++++++++--- .../tests/unwrap/interpolator_test.py | 83 +++++++++++++++++++ 2 files changed, 148 insertions(+), 13 deletions(-) diff --git a/packages/essreduce/src/ess/reduce/unwrap/interpolator_numba.py b/packages/essreduce/src/ess/reduce/unwrap/interpolator_numba.py index 58a64b410..69f098187 100644 --- a/packages/essreduce/src/ess/reduce/unwrap/interpolator_numba.py +++ b/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 diff --git a/packages/essreduce/tests/unwrap/interpolator_test.py b/packages/essreduce/tests/unwrap/interpolator_test.py index b7f14426e..f4e2367e5 100644 --- a/packages/essreduce/tests/unwrap/interpolator_test.py +++ b/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)), + )