Skip to content

feat: kernel-density CDF meshes RectangularKernelAdaptDensity/Image#374

Merged
Jammy2211 merged 3 commits into
mainfrom
feature/rectangular-kernel-cdf-mesh
Jul 10, 2026
Merged

feat: kernel-density CDF meshes RectangularKernelAdaptDensity/Image#374
Jammy2211 merged 3 commits into
mainfrom
feature/rectangular-kernel-cdf-mesh

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

Adds opt-in kernel-density-CDF variants of the adaptive rectangular meshes: RectangularKernelAdaptDensity / RectangularKernelAdaptImage (issue #373). The per-axis transform is the smooth kernel CDF F(x) = Σᵢ wᵢ·Φ((x−xᵢ)/h) (Enzi et al. arXiv:2606.30620 RTU formulation) instead of the empirical point-rank CDF, which the 2026-07-09 JAX gradient audit proved makes the likelihood exactly piecewise-constant in mass/shear at pixelization over-sampling 1 and on the interferometer sparse path.

Properties by construction: strictly monotone (no jitter hack), C^∞ in interp queries and traced-point positions (no ranks, no sorts, no argsort anywhere), duplicate-safe (no 1/Δknot terms). The forward transform is evaluated exactly and maps the data bounding box onto the unit square (matching the linear convention); the inverse is a fixed n_knots table lookup whose queries are unit-grid constants, so gradients flow smoothly through the table values. bandwidth defaults to the mesh pixel scale and is exposed as a mesh kwarg. The linear and spline meshes are untouched.

Function-level JAX sanity (local, x64): eager == JIT to 1e-13; AD grad matches central FD to ~5e-10 relative through both the bilinear-weights path and the inverse-table path.

API Changes

Added only — two new opt-in mesh classes (aa.mesh.RectangularKernelAdaptDensity, aa.mesh.RectangularKernelAdaptImage), their interpolator, and additive default-None kwargs on AbstractMeshGeometry. No existing symbol removed, renamed, or changed in behaviour.
See full details below.

Test Plan

  • pytest test_autoarray/ — 892 passed (16 new tests in test_rectangular_kernel.py: monotonicity, bounding-box mapping, round-trip, duplicate safety, large-bandwidth limit, mapper shape/consistency vs linear, areas totals vs linear, interpolator/geometry plumbing)
  • Function-level JAX check: eager-vs-JIT consistency + AD-vs-FD on forward and inverse paths
  • Certification (workspace leg, gating merge): autolens_workspace_test/scripts/jax_grad/imaging_pixelization.py kernel variants at os_pix=1 AND 4 with strict FD on ALL params; interferometer.py sparse path; eager-vs-JIT; figure-of-merit parity within a few e-4 of the linear meshes
Full API Changes (for automation & release notes)

Added

  • aa.mesh.RectangularKernelAdaptDensity(shape=, bandwidth=1.0, n_knots=64) — density-adaptive rectangular mesh with kernel-density CDF transform (opt-in).
  • aa.mesh.RectangularKernelAdaptImage(shape=, weight_power=1.0, weight_floor=0.0, bandwidth=1.0, n_knots=64) — image-adaptive variant; adapt-image weights become the kernel weights.
  • autoarray.inversion.mesh.interpolator.rectangular_kernel — new module: create_transforms_kernel(), adaptive_rectangular_mappings_weights_via_interpolation_from_kernel(), adaptive_rectangular_areas_from_kernel(), adaptive_rectangular_transformed_grid_from_kernel(), InterpolatorRectangularKernel, KERNEL_CDF_DEFAULT_BANDWIDTH, KERNEL_CDF_DEFAULT_KNOTS.

Changed Signature (additive, backwards-compatible)

  • AbstractMeshGeometry.__init__(..., kernel_bandwidth=None, kernel_knots=None) — new optional kwargs; MeshGeometryRectangular.areas_transformed / edges_transformed route through the kernel helpers when kernel_bandwidth is set (mirrors the existing spline_deg dispatch).

Migration

  • None required — all changes are additive and opt-in.

Coordination note

Adjacent to #372 (rect-adapt, parked at sign-off): both touch mesh_geometry/rectangular.py in distinct hunks (its edges_transformed body vs this PR's dispatch branches) — clean merge in either order; if #372 lands first the kernel path inherits the corrected edges automatically.

Merge gate: do not merge until the workspace-leg certification PR (jax_grad strict-FD assertions) is green — that is this task's success criterion.

🤖 Generated with Claude Code

…ants

Opt-in adaptive rectangular meshes whose per-axis CDF transform is a smooth
kernel-density CDF F(x) = sum_i w_i * Phi((x - x_i)/h) (Enzi et al.
arXiv:2606.30620 RTU formulation) instead of the empirical point-rank CDF.
Strictly monotone by construction, C-infinity in interp queries and traced
point positions (no ranks, no sorts), duplicate-safe — so mass/shear
gradients are smooth in every configuration, including imaging at
pixelization over-sampling 1 and the interferometer sparse path where the
empirical CDF makes the likelihood exactly piecewise-constant.

Forward transform is evaluated exactly (unit square <-> data bounding box,
matching the linear convention); inverse is a fixed n_knots table lookup
whose queries are unit-grid constants, so gradients flow through the table
values. Bandwidth defaults to the mesh pixel scale and is exposed as a mesh
kwarg. Linear and spline meshes are untouched.

Closes nothing; certification via autolens_workspace_test jax_grad scripts
follows on the workspace leg (#373).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label Jul 10, 2026
Jammy2211 and others added 2 commits July 10, 2026 10:25
The bilinear interpolation mapper places reconstruction value (row r,
col c) at node U_y=(n_y-r-1)/(n_y-3), U_x=(c-1)/(n_x-3) in unit CDF space
(guard ring at the border, rows flipped by the mapper's row = n - i), but
edges_transformed drew a uniform [0,1] partition — pcolormesh plots put
flux up to ~1.5 mesh pixels from where the mapper scatters it (the "small
offsets" of issue #372, confirmed by a delta-function reproduction). The
edges are now the node midpoints pushed through the same CDF machinery
(guard-safe denominator for degenerate n<=3 meshes; border guard cells
clamp to the data span). Both consumers inherit the fix: the pcolormesh
plot path and rectangular_rotated's warped-centre reconstruction; spline
and linear CDFs share the upstream index convention. Plot geometry only —
the mapper is untouched, so inversions and likelihoods are byte-identical.

#372

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@Jammy2211

Copy link
Copy Markdown
Collaborator Author

Folded in feature/rect-adapt (c177391, PR #375) by merge — the kernel meshes' plot geometry inherits the #372 edge fix through the shared edges_transformed dispatch. Full suite re-run on the combined tree: 893 passed. Verification of the fix (incl. on the kernel meshes) is on #372. If #375 merges first this PR needs nothing; if this merges first, #375 becomes a no-op diff.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release PR queued for the next release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant