Repository navigation
refactor: consolidate rectangular meshes — kernel takes plain names, retire linear/spline/rotated #402
Description
Overview
PyAutoArray has accumulated 8 public rectangular mesh classes, 5 interpolators and 2 geometries from three generations of gradient work (linear empirical-CDF → spline-CDF → kernel-density-CDF, plus the PCA-rotated multi-node experiment). Only the kernel meshes (the certified gradient path, PR #374) and RectangularUniform (HowToLens methodology illustration) are needed. This task consolidates the family: the kernel implementation takes over the plain names RectangularAdaptDensity/RectangularAdaptImage (zero churn in workspaces, prior YAMLs and saved results), and the linear, spline and rotated variants are deleted from source (recoverable via git history + PyAutoGut condemned refs).
Plan
- Merge the kernel-CDF interpolator into
interpolator/rectangular.py;InterpolatorRectangularbecomes the kernel implementation (name + module path survive → zero forced downstream edits). - Delete the spline pair, rotated mesh/interpolator/geometry,
density_components.py, theRectangularKernel*names and the linear empirical-CDF path. RectangularAdaptDensity/AdaptImagegainbandwidth/n_knotskwargs;RectangularUniformunchanged behaviorally.- Redistribute the 17 kernel tests under the plain names; delete spline/rotated tests; replace kernel-vs-linear comparisons with analytic invariants.
- Downstream (PyAutoGalaxy/PyAutoLens): expected zero-diff — verify by test suite.
- Workspace leg: user workspaces need zero code edits; update
autolens_workspace_testjax_grad variants to the plain names (FD certification is the gate); stashautolens_workspace_developerspline/rotated experiments in PyAutoGut.
Autodiff evidence for keeping kernel (gradient audit #87 + certified jax_grad harness): linear is exactly piecewise-constant in mass/shear at os=1 and on the interferometer sparse path (no usable gradients; 1–3% FD contamination at os=4); spline was never certified (eager-vs-JIT LL gap ~15, erratic FD; superseded); kernel is C∞, strict FD-certified on ALL params in every configuration (imaging os=1/os=4/production shape, interferometer sparse) with FoM parity 2.7e-5–6.3e-4 vs linear; uniform is exactly smooth with no transform.
Detailed implementation plan
Affected Repositories
- PyAutoArray (primary)
- PyAutoGalaxy, PyAutoLens (verify zero-diff; refresh linear-CDF prose if found)
- autolens_workspace_test (jax_grad variant consolidation — verification gate)
- autolens_workspace_developer (experiment stash → PyAutoGut; BLOCKED on blackjax-smc-gradient-kernel claim)
- PyAutoGut (condemned refs for deleted material)
Branch Survey (2026-07-23)
| Repository | Current Branch | Dirty? |
|---|---|---|
| ./PyAutoArray | main | clean |
| ./PyAutoGalaxy | main | clean |
| ./PyAutoLens | main | clean |
| ./autolens_workspace_test | main | clean |
| ./autogalaxy_workspace_test | main | clean |
| ./autolens_workspace_developer | main | dirty (3 files) + claimed by blackjax-smc-gradient-kernel |
| ./PyAutoGut | main | clean |
Suggested branch: feature/rectangular-mesh-consolidation
Implementation Steps (library leg)
- Merge
autoarray/inversion/mesh/interpolator/rectangular_kernel.pyINTOinterpolator/rectangular.py(delete kernel file). Keep from the linear module onlyreverse_interp/reverse_interp_np(kernel inverse table). Deleteforward_interp*, linearcreate_transforms, linear transformed-grid/mappings helpers, the linearInterpolatorRectangularclass body. Rename kernel functions to plain names (create_transforms_kernel→create_transforms,adaptive_rectangular_*_from_kernel→adaptive_rectangular_*_from);InterpolatorRectangularKernel→InterpolatorRectangular(AbstractInterpolator)with ctor(mesh, mesh_grid, data_grid, mesh_weight_map, adapt_data=None, bandwidth=KERNEL_CDF_DEFAULT_BANDWIDTH, n_knots=KERNEL_CDF_DEFAULT_KNOTS, xp=np). Carry over the_mappings_sizes_weights_splitpassthrough explicitly. Rewrite module docstring. - Delete
interpolator/rectangular_spline.py(incl.InvertPolySpline+ its pytree registration — the only pytree in the family),interpolator/rectangular_rotated_spline.py,interpolator/density_components.py(unwired). - Geometry:
mesh_geometry/abstract.pydropsspline_deg(keeps None-ablekernel_bandwidth/kernel_knots→ Delaunay + pickle tests untouched).mesh_geometry/rectangular.py: delete module-level linearadaptive_rectangular_areas_from(~390–418); collapse 3-branchareas_transformed/edges_transformedto a single kernel call with default-resolved bandwidth/knots (keeps the no-kernel-args uniform path working); keep neighbors/edges helpers +rectangular_edge_pixel_list_from(downstream import) + the issue-fix: rectangular-mesh pcolormesh plots misaligned with mapper node convention #372 node-midpoint edge logic exactly as is. Deletemesh_geometry/rectangular_rotated.py. - Mesh classes:
RectangularAdaptDensity.__init__gainsbandwidth/n_knots+ aninterpolator_kwargsproperty hook consumed by the baseinterpolator_from(collapses the duplicated kernelinterpolator_fromcopies).RectangularAdaptImageforwards kwargs, keepsmesh_weight_map_fromverbatim.RectangularUniformoverridesinterpolator_kwargs→{}(its interpolator ctor has no bandwidth/knots — would TypeError). Deleterectangular_kernel_adapt_*.py×2,rectangular_spline_adapt_*.py×2,rectangular_rotated_adapt_image.py. mesh/mesh/__init__.pyshrinks to Mesh, RectangularAdaptDensity, RectangularAdaptImage, RectangularUniform, Delaunay, KNearestNeighbor, KNNBarycentric — no compatibility aliases.autoarray/__init__.pyunchanged (optional: drop duplicate line-81 import).plot/inversion.pyisinstance dispatch works with zero edits (uniform does NOT subclass InterpolatorRectangular — imshow-vs-pcolormesh split preserved).- Tests: delete
mesh/test_rectangular_spline.py(12) +mesh/test_rectangular_rotated_adapt_image.py(14). Redistribute the 17 tests ofmesh/test_rectangular_kernel.py: construction/weight-map →mesh/test_rectangular.pyunder plain names (RectangularAdaptImage's first dedicated coverage); transform/helper →interpolator/test_rectangular.py(currently an emptypassstub). Kernel-vs-linear comparisons → analytic invariants (areas sum to bounding-box area; large-bandwidth → uniform limit). Recompute the hard-coded linear value inmesh_geometry/test_rectangular.py::test__areas_transformed; port the fix: rectangular-mesh pcolormesh plots misaligned with mapper node convention #372 regression test's helper imports.
Behavioural risks (flag in PR body)
- Plain-name results shift numerically (linear → kernel CDF; certified FoM parity ~e-4).
- Uniform-mesh
areas_for_magnificationlatently switches to kernel defaults (no consumer found). - JAX re-certification only via workspace_test jax_grad (library unit tests numpy-only).
- One confirmed amplifier of the JIT-only branch flips (research: root-cause the measure-thin solver branch flips in pixelized likelihoods #377) is the LINEAR rank-CDF bounding-box discontinuity — deleting linear removes an amplifier.
Workspace leg (after library merge)
- User workspaces (autolens_workspace, autogalaxy_workspace, HowToLens, HowToGalaxy, euclid pipeline): zero code edits; spot-check tutorial prose describing CDF mechanics.
- autolens_workspace_test: merge jax_grad kernel variants (E/F/G in
scripts/jax_grad/imaging_pixelization.py, D ininterferometer.py) into the plain-named variants; drop linear-only variants; FD certification stays green. - autolens_workspace_developer (wait on blackjax claim): stash
rect_adapt_duo/,searches_minimal/probe_grad_pix*.py,jax_profiling/misc/pixelization_spline_vs_linear.py+_fit_comparison.pyin PyAutoGut; rename kernel-named scripts to plain names; updatejax_profiling/gradient/README.mdrows.
Key Files
autoarray/inversion/mesh/interpolator/rectangular.py— becomes the kernel implementationautoarray/inversion/mesh/interpolator/rectangular_kernel.py— merged in, deletedautoarray/inversion/mesh/mesh_geometry/rectangular.py— branch collapseautoarray/inversion/mesh/mesh/rectangular_adapt_density.py/rectangular_adapt_image.py/rectangular_uniform.py— surviving meshestest_autoarray/inversion/pixelization/mesh/test_rectangular_kernel.py— redistributed
Verification
python -m pytest test_autoarray/green; grep sweepRectangularKernel|RectangularSpline|RectangularRotated|InvertPolySpline|create_transforms_kernel|spline_deg|density_components|forward_interpreturns zero hits in autoarray/ + test_autoarray/.- PyAutoGalaxy + PyAutoLens full pytest against the new autoarray (expected zero-diff).
- autolens_workspace_test jax_grad strict-FD assertions green on the plain-named meshes (os=1, os=4, production shape, interferometer sparse).
- Aggregator round-trip: old plain-name result loads under the new classes.
Original Prompt
Click to expand starting prompt
In PYAutoArray we have ended up with loads of Rectangular mesh classes and associated other classes, and its confusing and unecessary. I think only two are actually use: The kernel one which actually does gradients ok, and the uniform one which does not use interpolation which is used in HowToLens to illustrate the methodology. Can you do a census of all these different classes, confirm only these are useful (but keep both adaptdensity and adaptimage) and refactor and tidy up this part of the source code? In turn make sure all workspaces only use this gradient one which is the one users will actually use. Double check none others have any use. We also did some work on a more adaptive rectangular mesh which for example adapts to multiple nodes. This work may come back (e.g. once we have the basic rectangular working with gradients well) but lets move it out the source code, if its thre, for now, and stash workspace work in PyAutoGut. im not usre what rectangular_rotated_adapt_image.py is for and if we still need it, and if the split / kernel separrate methods both have uses but lets get down to just one
Activity
Jammy2211 commentedon Jul 23, 2026
Library PR Created
pending-release)af.Modelprior counts unchanged (2/4/2) · Heart YELLOW acknowledged (2 pre-existing reasons, unrelated).API changes (summary):
RectangularAdaptDensity/AdaptImageare now the kernel-CDF implementation with optionalbandwidth/n_knotskwargs;RectangularKernel*,RectangularSpline*,RectangularRotatedAdaptImage, the linear empirical-CDF path anddensity_componentsare removed.RectangularUniform,InterpolatorRectangular(name/path), prior-config keys and downstream imports unchanged.Workspace impact: exactly two files reference removed names —
autolens_workspace_test/scripts/jax_grad/imaging_pixelization.py(variants E/F/G) andjax_grad/interferometer.py(variant D). No user workspace, HowTo, euclid or config references. Next:/start_workspaceto consolidate the jax_grad variants onto the plain names and re-run the strict-FD certification (the merge gate for this PR). Theautolens_workspace_developercleanup + PyAutoGut stash is deferred while that repo is claimed byblackjax-smc-gradient-kernel.Jammy2211 commentedon Jul 23, 2026
Workspace leg blocked (serialised):
autolens_workspace_testis claimed by taskenv-inline-declarations(PyAutoHands#187), whose coordination note pre-agreed to serialise on that repo. The jax_grad variant consolidation + strict-FD certification (the merge gate for #403) resumes when that claim releases. Theautolens_workspace_developercleanup + PyAutoGut stash remains separately deferred on theblackjax-smc-gradient-kernelclaim. PR #403 stays open,pending-release, un-merged until certification.Jammy2211 commentedon Jul 23, 2026
Enzi 2026 citation leg (docs) shipped: PyAutoLens#644 + PyAutoGalaxy#520 (both
pending-release) add a Rectangular Mesh section todocs/general/citations.md— adaptive rectangular meshes (RectangularAdaptDensity/AdaptImage) must cite Enzi et al. 2026 (arXiv:2606.30620, RTU grids) under keyEnzi2026;RectangularUniformexempt; regularization noted as differing from the paper's GP prior — plus the matchingfiles/citations.bibentry. The workspace-example paragraphs (pixelization feature examples + HowTo* chapter 4: paper link + brief regularization caveat) are folded into the already-queued workspace legs behind theenv-inline-declarationsclaims.Jammy2211 commentedon Jul 23, 2026
All three PRs merged (human-directed): PyAutoArray#403 (after a trivial main-merge: version bump 2026.7.23.1 + the branch's duplicate-import removal; note main's bump commit had introduced CRLF into
autoarray/__init__.py— the merge restores LF per repo rules), PyAutoLens#644, PyAutoGalaxy#520. Remaining: autolens_workspace_test jax_grad migration + strict-FD re-certification (claim now released — proceeding), HowTo* citation paragraphs, then the developer-repo/Gut leg.Jammy2211 commentedon Jul 23, 2026
HowTo citation PRs merged on green CI: HowToLens#53 + HowToGalaxy#43. All five PRs of this arc are now in (PyAutoArray#403, PyAutoLens#644, PyAutoGalaxy#520, HowToLens#53, HowToGalaxy#43). Remaining legs (all claim-blocked): autolens_workspace_test jax_grad migration + strict-FD re-certification (urgent — main's jax_grad scripts import the deleted RectangularKernel* names), user-workspace pixelization citation paragraphs, developer-repo cleanup + Gut stash.
Jammy2211 commentedon Jul 23, 2026
jax_grad migration re-certified + PR open: autolens_workspace_test#209 (claim override user-authorized; only the 2 jax_grad files touched, ENV lines byte-identical). Both scripts pass in full against merged main: strict FD-step-sweep on ALL parameters — imaging os=1 (bandwidth 0.1, documented Einstein-radius branch-flip exclusion), os=4, full production shape, and the interferometer sparse path — mass/shear live everywhere, eager-vs-JIT consistent. The linear staircase variants are retired exactly as their own docstrings prescribed; FoM-parity assertions retired with their linear reference (historical measurements remain in the audit README). Merging on green CI. Remaining after this: user-workspace citation paragraphs + developer-repo/Gut stash (both still claim-blocked).
Jammy2211 commentedon Jul 23, 2026
Smoke fallout found + fixed on wst#209: the first CI run caught what the census grep could not —
jax_likelihood_functions/imaging/rectangular*.pyhard-code reference LLs from the linear mesh. Re-derived (2 configs improve at default bandwidth; dspl pins bandwidth=0.1 with rationale — default over-smooths its adapt weights). Bandwidth-defaults research prompt filed:draft/research/autoarray/rectangular_kernel_bandwidth_defaults.md. Merging #209 on green CI.Jammy2211 commentedon Jul 23, 2026
wst#209 MERGED (human-directed; PR-gating CI was green on the content commit, final head carried only the env-docstring merge). The consolidation's certification leg is complete: strict-FD green on all params in all configs, references re-derived, dspl bandwidth pinned + documented. Still open on this issue: user-workspace pixelization citation paragraphs (behind env-declaration-docstring-form claims) and the developer-repo cleanup + PyAutoGut stash (behind blackjax-smc-gradient-kernel). Bandwidth-defaults research prompt filed.
Jammy2211 commentedon Jul 24, 2026
Wrapping up — consolidation complete, released-pending
All six PRs merged: PyAutoArray#403 (library), PyAutoLens#644 + PyAutoGalaxy#520 (Enzi citations page + bib), HowToLens#53 + HowToGalaxy#43 (ch4 tutorial paragraph), autolens_workspace_test#209 (jax_grad migration + reference re-derivation). Verified 892/1006/448 tests green; downstream zero-diff;
af.Modelprior counts unchanged. Kernel-CDF meshes now hold the plain namesRectangularAdaptDensity/AdaptImage; linear/spline/rotated +density_componentsdeleted (recoverable via history). The surviving mesh implements the RTU formulation of Enzi et al. 2026 (arXiv:2606.30620), now a documented citation requirement.Completion record:
PyAutoMind/complete/2026/07/rectangular-mesh-consolidation.md. Worktree + all six feature branches cleaned up.Two in-scope follow-ups were blocked on other tasks' worktree claims and split to tracked drafts (they will be issued when the claims release):
autolens_workspace/autogalaxy_workspacepixelization examples —draft/docs/workspaces/rectangular_mesh_enzi_citation_examples.md(blocked on env-declaration-docstring-form).autolens_workspace_developercleanup + PyAutoGut stash of the spline/rotated experiments —draft/maintenance/autolens_workspace_developer/rectangular_experiments_gut_stash.md(blocked on blackjax-smc-gradient-kernel).Plus a research follow-up:
draft/research/autoarray/rectangular_kernel_bandwidth_defaults.md(fit-vs-set the kernel bandwidth; auto-scaled default). Closing.