Merge pull request #231 from scipp/interpolator_optmizations

Time-of-flight interpolator optmizations

Author: nvaytet

src/ess/reduce/time_of_flight/eto_to_tof.py

@@ -308,7 +308,11 @@ def __init__(self, lookup: sc.DataArray, distance_unit: str, time_unit: str):
         )
 
     def __call__(
-        self, ltotal: sc.Variable, event_time_offset: sc.Variable
+        self,
+        ltotal: sc.Variable,
+        event_time_offset: sc.Variable,
+        pulse_period: sc.Variable,
+        pulse_index: sc.Variable | None = None,
     ) -> sc.Variable:
         if ltotal.unit != self._distance_unit:
             raise sc.UnitError(
@@ -326,7 +330,12 @@ def __call__(
 
         return sc.array(
             dims=out_dims,
-            values=self._interpolator(times=event_time_offset, distances=ltotal),
+            values=self._interpolator(
+                times=event_time_offset,
+                distances=ltotal,
+                pulse_index=pulse_index.values if pulse_index is not None else None,
+                pulse_period=pulse_period.value,
+            ),
             unit=self._time_unit,
         )
 
@@ -359,7 +368,11 @@ def _time_of_flight_data_histogram(
     interp = TofInterpolator(lookup, distance_unit=ltotal.unit, time_unit=eto_unit)
 
     # Compute time-of-flight of the bin edges using the interpolator
-    tofs = interp(ltotal=ltotal.broadcast(sizes=etos.sizes), event_time_offset=etos)
+    tofs = interp(
+        ltotal=ltotal.broadcast(sizes=etos.sizes),
+        event_time_offset=etos,
+        pulse_period=pulse_period,
+    )
 
     return rebinned.assign_coords(tof=tofs)
 
@@ -418,11 +431,13 @@ def _guess_pulse_stride_offset(
         values=event_time_offset.values[inds],
         unit=event_time_offset.unit,
     )
-    pulse_period = pulse_period.to(unit=etos.unit)
     for i in range(pulse_stride):
         pulse_inds = (pulse_index + i) % pulse_stride
         tofs[i] = interp(
-            ltotal=ltotal, event_time_offset=etos + pulse_inds * pulse_period
+            ltotal=ltotal,
+            event_time_offset=etos,
+            pulse_index=pulse_inds,
+            pulse_period=pulse_period,
         )
     # Find the entry in the list with the least number of nan values
     return sorted(tofs, key=lambda x: sc.isnan(tofs[x]).sum())[0]
@@ -446,12 +461,12 @@ def _time_of_flight_data_events(
     ltotal = sc.bins_like(etos, ltotal).bins.constituents["data"]
     etos = etos.bins.constituents["data"]
 
-    # Compute a pulse index for every event: it is the index of the pulse within a
-    # frame period. When there is no pulse skipping, those are all zero. When there is
-    # pulse skipping, the index ranges from zero to pulse_stride - 1.
-    if pulse_stride == 1:
-        pulse_index = sc.zeros(sizes=etos.sizes)
-    else:
+    pulse_index = None
+    pulse_period = pulse_period.to(unit=eto_unit)
+
+    if pulse_stride > 1:
+        # Compute a pulse index for every event: it is the index of the pulse within a
+        # frame period. The index ranges from zero to pulse_stride - 1.
         etz_unit = 'ns'
         etz = (
             da.bins.coords["event_time_zero"]
@@ -495,7 +510,9 @@ def _time_of_flight_data_events(
     # Compute time-of-flight for all neutrons using the interpolator
     tofs = interp(
         ltotal=ltotal,
-        event_time_offset=etos + pulse_index * pulse_period.to(unit=eto_unit),
+        event_time_offset=etos,
+        pulse_index=pulse_index,
+        pulse_period=pulse_period,
     )
 
     parts = da.bins.constituents

src/ess/reduce/time_of_flight/interpolator_numba.py

@@ -11,6 +11,8 @@ def interpolate(
     values: np.ndarray,
     xp: np.ndarray,
     yp: np.ndarray,
+    xoffset: np.ndarray | None,
+    deltax: float,
     fill_value: float,
     out: np.ndarray,
 ):
@@ -29,6 +31,10 @@ def interpolate(
         1D array of x-coordinates where to interpolate (size N).
     yp:
         1D array of y-coordinates where to interpolate (size N).
+    xoffset:
+        1D array of integer offsets to apply to the x-coordinates (size N).
+    deltax:
+        Multiplier to apply to the integer offsets (i.e. the step size).
     fill_value:
         Value to use for points outside of the grid.
     out:
@@ -52,7 +58,7 @@ def interpolate(
     norm = one_over_dx * one_over_dy
 
     for i in prange(npoints):
-        xx = xp[i]
+        xx = xp[i] + (xoffset[i] * deltax if xoffset is not None else 0.0)
         yy = yp[i]
 
         if (xx < xmin) or (xx > xmax) or (yy < ymin) or (yy > ymax):
@@ -108,14 +114,22 @@ def __init__(
         self.values = values
         self.fill_value = fill_value
 
-    def __call__(self, times: np.ndarray, distances: np.ndarray) -> np.ndarray:
+    def __call__(
+        self,
+        times: np.ndarray,
+        distances: np.ndarray,
+        pulse_period: float = 0.0,
+        pulse_index: np.ndarray | None = None,
+    ) -> np.ndarray:
         out = np.empty_like(times)
         interpolate(
             x=self.time_edges,
             y=self.distance_edges,
             values=self.values,
             xp=times,
             yp=distances,
+            xoffset=pulse_index,
+            deltax=pulse_period,
             fill_value=self.fill_value,
             out=out,
         )

src/ess/reduce/time_of_flight/interpolator_scipy.py

@@ -39,16 +39,21 @@ def __init__(
         from scipy.interpolate import RegularGridInterpolator
 
         self._interp = RegularGridInterpolator(
-            (
-                distance_edges,
-                time_edges,
-            ),
+            (distance_edges, time_edges),
             values,
             method=method,
             bounds_error=bounds_error,
             fill_value=fill_value,
             **kwargs,
         )
 
-    def __call__(self, times: np.ndarray, distances: np.ndarray) -> np.ndarray:
+    def __call__(
+        self,
+        times: np.ndarray,
+        distances: np.ndarray,
+        pulse_period: float = 0.0,
+        pulse_index: np.ndarray | None = None,
+    ) -> np.ndarray:
+        if pulse_index is not None:
+            times = times + (pulse_index * pulse_period)
         return self._interp((distances, times))

tests/time_of_flight/interpolator_test.py

@@ -53,6 +53,22 @@ def test_numba_and_scipy_interpolators_yield_same_results():
     assert np.allclose(numba_result, scipy_result)
 
 
+def test_numba_and_scipy_interpolators_yield_same_results_with_pulse_offset():
+    numba_interp, scipy_interp = _make_interpolators()
+
+    rng = np.random.default_rng(seed=42)
+    npoints = 1000
+    times = rng.uniform(0, 71, npoints)
+    distances = rng.uniform(40, 70, npoints)
+    offsets = rng.uniform(0, 2, npoints)
+    period = 1.0
+
+    numba_result = numba_interp(times, distances, period, offsets)
+    scipy_result = scipy_interp(times, distances, period, offsets)
+
+    assert np.allclose(numba_result, scipy_result, equal_nan=True)
+
+
 def test_numba_and_scipy_interpolators_yield_same_results_with_out_of_bounds():
     numba_interp, scipy_interp = _make_interpolators()