From 2bfa08ee7ec70324b7c8855298b6d8bc66fea489 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 21 Aug 2026 01:34:23 +0000 Subject: [PATCH] perf: windowed numba fast path for the kernel-CDF forward transform (numpy branch) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The RectangularAdaptDensity kernel-CDF forward transform dominates the numba CPU likelihood (euclid 1.66 s = 55% of the eval; hst 27 s = 89%): an O(M x N) blocked erf broadcast with ~126 MB per-block temporaries, rebuilt every evaluation. The numpy branch of F_raw now evaluates each dimension's weighted kernel CDF with a numba kernel over sorted points + weight prefix sums and a +-9-bandwidth saturation window — same values to ~1e-13 (dropped tail terms < 1e-19 of the weight sum), 3.0-3.4x on the step. The blocked numpy implementation remains the JAX branch and the differentiable reference; the sort is internal evaluation order on the gradient-free numpy path only, preserving the module's no-sorts differentiability guarantee where it matters. Validated: autolens_profiling pixelization_numba pins PASS at euclid + hst (euclid eval 3.08 -> 1.17 s, hst ~30 -> 10.1 s on a 4-core container, stacked with the merged #453/#455 wins); test_autoarray 1036 passed (+2 new dense-reference equivalence tests, weighted and unweighted, with out-of-range queries exercising the saturated tails). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Vcc7MUBMnNU6n8qqS9ioVZ --- .../mesh/interpolator/rectangular.py | 72 +++++++++++++++++-- .../interpolator/test_rectangular.py | 44 ++++++++++++ 2 files changed, 109 insertions(+), 7 deletions(-) diff --git a/autoarray/inversion/mesh/interpolator/rectangular.py b/autoarray/inversion/mesh/interpolator/rectangular.py index 581990a9f..47a01a461 100644 --- a/autoarray/inversion/mesh/interpolator/rectangular.py +++ b/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). diff --git a/test_autoarray/inversion/pixelization/interpolator/test_rectangular.py b/test_autoarray/inversion/pixelization/interpolator/test_rectangular.py index 51d1e6e2f..ba8bd160c 100644 --- a/test_autoarray/inversion/pixelization/interpolator/test_rectangular.py +++ b/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)