Merge pull request #458 from PyAutoLabs/feature/kernel-cdf-numba-fast-path

perf: windowed numba fast path for the kernel-CDF forward transform

Author: Jammy2211

autoarray/inversion/mesh/interpolator/rectangular.py

@@ -31,12 +31,16 @@
 constants, gradients flow smoothly through the table values.
 """
 
+import math
+
 import numpy as np
 from functools import partial
 from typing import Optional
 
 from autonerves import cached_property
 
+from autoarray import numba_util
+
 from autoarray.inversion.mesh.interpolator.abstract import AbstractInterpolator
 
 
@@ -84,6 +88,43 @@ def reverse_interp_np(xp, yp, x):
 
 _SQRT2 = np.sqrt(2.0)
 
+# Phi(t) saturates to exactly-representable 0/1 contributions in fp64 well
+# inside |t| = 9: the dropped tail terms are < 1e-19 of the weight sum, far
+# below the ~1e-13 accumulation noise of the blocked numpy sum the windowed
+# numba kernel replaces (measured max deviation 1e-13 on the hst fiducial).
+_KERNEL_CDF_SATURATION_T = 9.0
+
+
+@numba_util.jit()
+def _kernel_cdf_dim_windowed(p_sorted, w_sorted, w_prefix, h_d, q, T):
+    """
+    Exact 1D weighted kernel CDF ``F(q) = sum_i w_i Phi((q - p_i) / h)`` for
+    one dimension, evaluated with a saturation window over sorted points.
+
+    Points below ``q - T h`` contribute exactly their weight (prefix sum);
+    points above ``q + T h`` contribute zero; only the window is summed with
+    ``erfc``. Replaces the O(M x N) blocked numpy broadcast on the numpy path
+    — same values to ~1e-13 (see ``_KERNEL_CDF_SATURATION_T``) at ~3x the
+    speed and none of the ~126 MB per-block temporaries; the blocked numpy
+    implementation remains the JAX-path/differentiable reference.
+
+    Sorting note: the module docstring's "no sorts anywhere" invariant is a
+    JAX-differentiability guarantee. This kernel runs only on the ``xp is
+    np`` branch, which carries no gradients — the sort is an internal
+    evaluation order and the returned VALUES are those of the sort-free sum.
+    """
+    out = np.empty(q.shape[0])
+    inv = 1.0 / (h_d * 1.4142135623730951)
+    for m in range(q.shape[0]):
+        qm = q[m]
+        a = np.searchsorted(p_sorted, qm - T * h_d)
+        b = np.searchsorted(p_sorted, qm + T * h_d)
+        acc = w_prefix[a]
+        for i in range(a, b):
+            acc += w_sorted[i] * 0.5 * math.erfc((p_sorted[i] - qm) * inv)
+        out[m] = acc
+    return out
+
 
 def _norm_cdf(t, xp):
     """Standard normal CDF, xp-aware (scipy erf on numpy, jax.scipy under jax)."""
@@ -161,15 +202,32 @@ def F_raw(q):
             return out.reshape(n_blocks * KERNEL_FORWARD_BLOCK, 2)[:M]
 
     else:
+        # numpy fast path: per-dimension sorted points + weight prefix sums,
+        # evaluated by the windowed numba kernel. Same values as the blocked
+        # broadcast above to ~1e-13; that implementation stays as the JAX
+        # branch and the differentiable reference.
+        _p_sorted = []
+        _w_sorted = []
+        _w_prefix = []
+        for _d in range(2):
+            _order = np.argsort(points[:, _d], kind="stable")
+            _p_sorted.append(np.ascontiguousarray(np.asarray(points)[_order, _d]))
+            _w_sorted.append(np.ascontiguousarray(np.asarray(w)[_order]))
+            _w_prefix.append(np.concatenate([[0.0], np.cumsum(_w_sorted[_d])]))
 
         def F_raw(q):
-            return np.concatenate(
-                [
-                    F_raw_block(q[i : i + KERNEL_FORWARD_BLOCK])
-                    for i in range(0, q.shape[0], KERNEL_FORWARD_BLOCK)
-                ],
-                axis=0,
-            )
+            q = np.asarray(q)
+            out = np.empty_like(q)
+            for d in range(2):
+                out[:, d] = _kernel_cdf_dim_windowed(
+                    _p_sorted[d],
+                    _w_sorted[d],
+                    _w_prefix[d],
+                    float(h[d]),
+                    np.ascontiguousarray(q[:, d]),
+                    _KERNEL_CDF_SATURATION_T,
+                )
+            return out
 
     # The unit square maps onto the data bounding box exactly (the kernel
     # tails outside [lo, hi] are absorbed by the rescale).

test_autoarray/inversion/pixelization/interpolator/test_rectangular.py

@@ -226,3 +226,47 @@ def __getattr__(self, item):
     assert areas.shape == (36,)
     assert np.all(np.isfinite(areas))
     assert np.all(areas > 0.0)
+
+
+# ---------------------------------------------------------------------------
+# Windowed numba fast path (numpy branch)
+# ---------------------------------------------------------------------------
+
+
+def _dense_reference_forward(data_grid, mesh_pixels, weights, q):
+    """The pre-fast-path definition: dense O(M x N) normal-CDF sum, rescaled
+    so the data bounding box maps onto the unit square, clipped to [0, 1]."""
+    from scipy.special import erf
+
+    points = data_grid
+    N = points.shape[0]
+    w = np.full(N, 1.0 / N) if weights is None else weights / weights.sum()
+    lo, hi = points.min(axis=0), points.max(axis=0)
+    h = 1.0 * (hi - lo) / mesh_pixels
+
+    def F_raw(qq):
+        t = (qq[:, None, :] - points[None, :, :]) / h[None, None, :]
+        return np.sum(w[None, :, None] * (0.5 * (1.0 + erf(t / np.sqrt(2.0)))), axis=1)
+
+    F_lo = F_raw(lo[None, :])[0]
+    F_hi = F_raw(hi[None, :])[0]
+    return np.clip((F_raw(q) - F_lo[None, :]) / (F_hi - F_lo)[None, :], 0.0, 1.0)
+
+
+@pytest.mark.parametrize("weighted", [False, True])
+def test__forward_transform__windowed_numba_matches_dense_reference(weighted):
+    data_grid, data_grid_over, weights = _seeded_inputs(M=300, K=500, seed=7)
+
+    # queries beyond the data bounding box exercise the saturated tails
+    q = np.concatenate([data_grid_over, data_grid.min(axis=0) - 1.0 + np.zeros((1, 2)),
+                        data_grid.max(axis=0) + 1.0 + np.zeros((1, 2))])
+
+    fwd, _ = create_transforms(
+        data_grid, mesh_pixels=16, mesh_weight_map=weights if weighted else None, xp=np
+    )
+
+    reference = _dense_reference_forward(
+        data_grid, 16, weights if weighted else None, q
+    )
+
+    np.testing.assert_allclose(fwd(q), reference, rtol=0.0, atol=1e-12)