Skip to content

perf: windowed numba fast path for the RectangularAdaptDensity kernel-CDF transform #457

Description

@Jammy2211

Overview

The kernel-CDF forward transform of RectangularAdaptDensity dominates the numba CPU likelihood on the default Rectangular route: 1.66 s = 55% of a euclid eval, 27 s = 89% at hst (PyAutoLabs/autolens_profiling#151) — an O(M×N) blocked erf broadcast with ~126 MB per-block temporaries, rebuilt every evaluation (the traced grid moves with the mass model, so no cross-eval caching applies). Phase 3 of the CPU speed campaign, per the PyAutoMind kernel-CDF prompt.

Plan

  • Numba kernel _kernel_cdf_dim_windowed in inversion/mesh/interpolator/rectangular.py: per-dimension weighted kernel CDF over sorted points + weight prefix sums with a ±9-bandwidth saturation window (dropped tails < 1e-19 of the weight sum → values match the exact sum to ~1e-13).
  • The numpy branch of F_raw uses it; the blocked numpy implementation stays as the JAX branch and the differentiable reference. The module's "no sorts" differentiability guarantee is preserved — the sort is internal evaluation order on the gradient-free numpy path only.
  • Beyond ~3.4× the transform is irreducibly O(M×N) for exact values; the measured interpolated-CDF option (K=8192 → Δlnℒ ≤ +4e-3 / ≤ 2e-3 differential, 18–55×) and the default-mesh question are tracked separately (Inference programme: no fast default CPU mesh — Rectangular is the default and slow on CPU (new-user hazard) autolens_profiling#153, programme Phase 14).

Results (4-core cloud container)

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions