Affected Repositories

• PyAutoArray (primary)
• autolens_workspace_test (certification scripts)
• autolens_workspace_developer (gradient README rows)

Branch Survey

Suggested branch: feature/rectangular-kernel-cdf-mesh
Worktree: ~/Code/PyAutoLabs-wt/rectangular-kernel-cdf-mesh/

Coordination note

The parked rect-adapt worktree (#372, uncommitted, awaiting ship sign-off) also edits autoarray/inversion/mesh/mesh_geometry/rectangular.py, but only the edges_transformed hunk (node-midpoint edges); this task adds dispatch in areas_transformed / transformed-grid / ctor — distinct hunks, clean merge expected in either order. If #372 merges first, the kernel path inherits the corrected edges automatically.

Implementation Steps

1. Stash preamble: diff PyAutoArray stash@{0} against current rectangular.py; salvage or reject (likely reject — it predates the audit and belongs to the rejected spline chain); comment the decision here; retire the parked.md rectangular-spline-cdf entry; keep the stash until the decision comment lands.
2. New module autoarray/inversion/mesh/interpolator/rectangular_kernel.py (mirror rectangular_spline.py's independently-auditable-copy layout):
   • create_transforms_kernel(traced_points, mesh_weight_map=None, bandwidth=None, n_knots=64, xp=np) — same contract as create_transforms (rectangular.py:70). Per axis: weights w_i = 1/N (density) or normalized mesh_weight_map (image); F(x) = Σ_i w_i · Φ((x − x_i)/h), Φ = Gaussian CDF via scipy.special.erf (np) / jax.scipy.special.erf (jax). Forward evaluated exactly at query points, rescaled by [F(lo), F(hi)] to fill [0,1]. Inverse via interp on a fixed K=n_knots knot grid spanning the data range padded by ~4h — inverse queries are fixed U-grid constants, so gradients flow smoothly through the table values. Bandwidth default h = bandwidth_factor × data_range_axis / mesh_pixels_axis, factor ~1.0, overridable via the mesh bandwidth kwarg.
   • adaptive_rectangular_mappings_weights_via_interpolation_from_kernel(...) — copy of steps 3–7 of the linear helper (floor/ceil, flatten, bilinear weights) routed through kernel transforms.
   • adaptive_rectangular_areas_from_kernel(...) + adaptive_rectangular_transformed_grid_from_kernel(...) — kernel versions of the geometry helpers.
   • InterpolatorRectangularKernel(InterpolatorRectangular) — overrides _mappings_sizes_weights and mesh_geometry, mirroring InterpolatorRectangularSpline.
3. Geometry dispatch in autoarray/inversion/mesh/mesh_geometry/rectangular.py: MeshGeometryRectangular already branches on spline_deg is not None (~lines 478/526); add parallel optional kernel_bandwidth / kernel_knots ctor kwargs (default None) routing areas_transformed and the transformed-grid property through the *_from_kernel helpers.
4. Mesh classes autoarray/inversion/mesh/mesh/rectangular_kernel_adapt_density.py / rectangular_kernel_adapt_image.py — RectangularKernelAdaptDensity(RectangularAdaptDensity), RectangularKernelAdaptImage(RectangularAdaptImage); template rectangular_spline_adapt_density.py; store bandwidth/n_knots, override interpolator_cls, forward kwargs in interpolator_from. Export from autoarray/inversion/mesh/mesh/__init__.py; mirror any registry mentions the spline classes have in autoarray/inversion/pixelization.py / autoarray/exc.py. Spline meshes stay untouched.
5. Unit tests (numpy-only): test_autoarray/inversion/pixelization/interpolator/test_rectangular_kernel.py — strict monotonicity; forward/inverse round-trip; duplicate-point safety; weight-map variant; mesh-class plumbing; areas positive + consistency vs linear variant on a smooth cloud. Full pytest test_autoarray/ green.
6. Certification (autolens_workspace_test/scripts/jax_grad/, same branch name):
   • imaging_pixelization.py: RectangularKernelAdaptDensity at os_pix=1 with strict FD on ALL params (no skip_indices — must NOT be a staircase); both kernel meshes at os_pix=4 strict FD; util.assert_eager_jit_consistent everywhere; eager figure_of_merit within a few e-4 relative of linear AdaptDensity/AdaptImage.
   • interferometer.py: kernel mesh on the sparse-operator path — strict FD on all mass/shear params + eager-vs-JIT.
   • Existing linear/spline variant assertions unchanged (no regression).
7. Docs row: autolens_workspace_developer/jax_profiling/gradient/README.md — kernel-mesh rows + staircase-section implications.

Key Files

• PyAutoArray/autoarray/inversion/mesh/interpolator/rectangular.py — linear transform contract being mirrored (create_transforms, bilinear helper).
• PyAutoArray/autoarray/inversion/mesh/interpolator/rectangular_spline.py — structural template for the new module (do not build on its math).
• PyAutoArray/autoarray/inversion/mesh/mesh_geometry/rectangular.py — geometry dispatch site.
• PyAutoArray/autoarray/inversion/mesh/mesh/rectangular_spline_adapt_density.py — mesh-class template.
• autolens_workspace_test/scripts/jax_grad/{imaging_pixelization,interferometer,util}.py — certification harness.
• autolens_workspace_developer/jax_profiling/gradient/README.md — audit table to update.

Risks

• Bandwidth trade-off: h too small re-approaches the staircase; h too large over-smooths geometry and drifts FoM parity. Default tied to mesh resolution + exposed kwarg.
• Inverse-table resolution: K=64 must resolve the CDF better than the mesh shape (~30×30 in jax_test configs); n_knots kwarg, bump if parity fails.

Kernel-density CDF transform for the adaptive rectangular mesh (differentiable everywhere)

Type: feature
Target: autoarray
Repos:

• PyAutoArray
• autolens_workspace_test
• autolens_workspace_developer
  Difficulty: medium
  Autonomy: supervised
  Priority: normal
  Status: formalised

Replace the empirical point-rank CDF in the adaptive rectangular mesh with a smooth
kernel-density CDF (the Enzi et al. arXiv:2606.30620 RTU formulation) so the mesh
carries correct smooth mass/shear gradients in every configuration — imaging at any
pixelization over-sampling and the interferometer sparse path (which has no
over-sampling and today has zero usable gradients).

Original request (verbatim): "it sounds like you think we can improve the rectangular
gradient method across all use cases? explain how [...]" → recommendation approved:
kernel-CDF leapfrog, one focused task, "drop the nuts stuff for now I only care that
the jax gradient workspace test script shows grad all works and can work on samples
down the line."

Background (from the 2026-07-09 JAX gradient audit, autolens_workspace_developer#87)

• The linear rank-CDF (create_transforms in
  PyAutoArray/autoarray/inversion/mesh/interpolator/rectangular.py) makes the
  likelihood exactly piecewise-constant in mass/shear whenever interp queries
  coincide with knots: imaging os_pix=1 and the interferometer sparse path
  (jax_grad/interferometer.py variant B documents the staircase). At os_pix=4
  gradients live but every rank swap leaves a micro-kink (~1-3% FD contamination).
• Prior attempt: RectangularSplineAdapt{Density,Image} (PyAutoArray PR Add RectangularSplineAdapt{Density,Image} meshes for gradient-based samplers #289,
  2026-04-22, opt-in, built for gradient samplers). Measured 2026-07-09: breaks the
  rank invariance (AD live on all 14 params at os_pix=1) BUT eager-vs-JIT LL gap
  ~15 and erratic FD (noisy surface — Add RectangularSplineAdapt{Density,Image} meshes for gradient-based samplers #289's shipped "oscillations" limitation).
  Suspects: _enforce_strict_monotone eps-jitter, deg-11 polyfit conditioning,
  Hermite inversion table. Do not build on this chain; keep the spline meshes
  untouched/opt-in.
• A 727-line jax-friendly spline-interpolator refactor is PARKED in PyAutoArray
  stash@{0} (parked.md rectangular-spline-cdf, 2026-05-08).

Task

1. Preamble: inspect PyAutoArray stash@{0} — salvage anything useful, then
   retire the parked.md entry either way (keep the stash until a decision is
   recorded on the issue).
2. Implement the kernel-density CDF transform as a new opt-in mesh variant (e.g.
   RectangularKernelAdaptDensity / ...AdaptImage), slotting into the existing
   transform/inv_transform machinery: per axis
   F(x) = sum_i w_i * Phi((x - x_i)/h) evaluated on a small fixed knot grid
   (K ~ 64), inverse via interpolation on the same grid. Strictly monotone by
   construction (no jitter hack), C-infinity in queries AND point positions,
   duplicate-safe (no 1/Delta-knot terms). Weight-map variant uses w_i from the
   adapt image. Bandwidth h defaults tied to mesh resolution; expose as a mesh
   kwarg.
3. Certification is the success criterion (no sampler work in this task):
   • Extend autolens_workspace_test/scripts/jax_grad/imaging_pixelization.py
     with kernel-mesh variants at os_pix=1 AND os_pix=4: strict FD assertions on
     ALL parameters (mass/shear included) using the existing util.py harness —
     the os_pix=1 variant must NOT be a staircase (that is the point).
   • Extend autolens_workspace_test/scripts/jax_grad/interferometer.py with the
     kernel mesh on the sparse-operator path: strict FD on all mass/shear params.
   • Eager-vs-JIT consistency (util.assert_eager_jit_consistent) must hold on
     every variant.
   • Fit-quality parity: eager figure_of_merit on the jax_test config within a few
     e-4 relative of the linear AdaptDensity/AdaptImage meshes (mesh geometry
     changes slightly; reconstruction quality must not degrade).
   • Update autolens_workspace_developer/jax_profiling/gradient/README.md rows +
     the staircase-section implications once certified.
4. Sampler trials (NUTS/HMC) are explicitly OUT of scope — future task once the
   jax_grad scripts are green.

Constraints

• PyAutoArray is claimed by nnls-solver-optimization (active.md) at time of
  filing — start when the claim releases or coordinate explicitly (different
  files, adjacent territory).
• Spline meshes (PR Add RectangularSplineAdapt{Density,Image} meshes for gradient-based samplers #289) stay as-is; no deletion/refactor of them in this task.
• Library unit tests numpy-only per repo rules; all JAX validation through the
  workspace_test jax_grad scripts.