Add RectangularSplineAdapt{Density,Image} meshes for gradient-based samplers#289
Merged
Add RectangularSplineAdapt{Density,Image} meshes for gradient-based samplers#289
Conversation
New spline-CDF variants of the adaptive rectangular meshes intended for
gradient-based samplers. Replaces the C⁰ empirical-CDF linear-interp
transform with a degree-11 polynomial fit to the inverse CDF plus a
cubic-Hermite spline inverter — C¹ continuous and JAX-autograd-compatible.
xp-aware so numpy and jax.numpy share one algorithm.
- `interpolator/rectangular_spline.py`: `InvertPolySpline`, transforms
factory, adaptive mappings/areas helpers, `InterpolatorRectangularSpline`
(subclass of `InterpolatorRectangular` so plot-dispatch isinstance
checks work without further changes).
- `mesh/rectangular_spline_adapt_{density,image}.py`: two new mesh
classes subclassing the existing density / image adapt meshes.
- `mesh_geometry/rectangular.py` + `abstract.py`: `MeshGeometryRectangular`
takes `spline_deg` and dispatches `areas_transformed` /
`edges_transformed` to the spline helpers when set. Linear path
unchanged.
- `mesh/__init__.py`: export the two new classes.
- 12 new unit tests; 60/60 existing pixelization tests still pass.
Co-Authored-By: Claude Opus 4.7 <[email protected]>
Collaborator
Author
|
Workspace follow-up: https://github.com/PyAutoLabs/autolens_workspace_developer/pull/30 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
RectangularSplineAdaptDensityandRectangularSplineAdaptImagemeshes using a polynomial + Hermite-spline CDF transform (deg=11 default, tunable) instead of the empirical-CDF linear interp used by the existing adapt meshes.isinstance(InterpolatorRectangular)dispatch inplot/inversion.pyroutes the spline source-plane reconstruction through_plot_rectangularautomatically.MeshGeometryRectangularcarries aspline_degand dispatchesareas_transformed/edges_transformedthrough the spline helpers — required forAdaptiveBrightnessregularization consistency.Why
User is implementing gradient-based samplers. The empirical-CDF linear interp gives piecewise-constant gradients with
1/Δknotspikes where traced points crowd together — noisy under HMC symplectic integrators. The spline CDF is C¹ end-to-end and passes throughjax.gradcleanly. Ported fromz_staging/rect_adap_spline_invert_jax (1).ipynb(RSE prototype) in xp-aware form.API changes
aa.mesh.RectangularSplineAdaptDensity(shape, spline_deg=11)aa.mesh.RectangularSplineAdaptImage(shape, weight_power=1.0, weight_floor=0.0, spline_deg=11)MeshGeometryRectangular.__init__now accepts an optionalspline_degkwarg (defaultNone→ existing linear-CDF behaviour).Scripts Changed
None in this repo. See
autolens_workspace_developerPR (linked after that one is opened) for two new benchmark scripts.Test plan
pytest test_autoarray/inversion/pixelization/mesh/test_rectangular_spline.py— 12/12 green (pure numpy, per no-JAX-in-unit-tests rule).pytest test_autoarray/inversion/pixelization/— 60/60 green (no regression in existing rectangular / delaunay / border tests).Known limitations
Under JIT with
use_positive_only_solver=Falsethe spline log_L shows small oscillations (~1e5 in std of first-difference) across a ±5% einstein_radius sweep where the linear path is monotone. Tracked inautolens_workspace_developerfollow-up prompt (linked in the workspace PR) — shipping now as opt-in while the gradient investigation continues.🤖 Generated with Claude Code