[Jammy2211]

Summary

Phase 5 (final) of the point-source solved-likelihood series (#657): differentiable point-source likelihoods through the PointSolver, via an implicit fixed-point custom_jvp — the same mechanism the source paper's own code (gravity.jl, Lombardi 2024 arXiv:2406.15280, Eq. 30) uses: differentiate at the solved positions (A dθ = dα|_θ + dβ), never through the triangle-refinement iteration. Gradient searches (af.MultiStartProdigy) can now run on image-plane point-source fits.

• New autolens/point/solver/implicit_diff.py: solve_padded_factory (a custom_jvp closure over the solver — PointSolver is deliberately not a registered pytree) with the tangent linear algebra factored into the numpy-testable implicit_tangents_from. Tangent rule is linear in tangents, so reverse mode (value_and_grad) transposes automatically. Padded (inf, inf) sentinel rows get zero tangent; near-critical det(A) → 0 divergence is surfaced, never clamped (binding phase-1 rule).
• PointSolver.solve JAX path routes through the wrapper; the raw solve is factored into _solve_array (numpy path byte-identical). A compatibility gate (tracer_is_jax_compatible) keeps hand-built/unregistered tracers (simulator scripts) on the plain forward path with unchanged behaviour.
• FitFluxes.model_data vectorized (drops a per-element traced list comprehension; matches the FitFluxesSolved pattern) — unblocks nested autodiff through the Hessian magnifications (third potential derivatives, which the paper declares analytically out of reach).
• FitPositionsImagePairRepeat docstring documents subgradient semantics (min-selection is piecewise; policy masks are piecewise-constant).

Differentiability contract (module docstring): exact between image-count events up to the solver residual; at caustic-crossing/threshold/containment flips the likelihood itself is discontinuous (measure-zero seams, same class as Delaunay re-wiring); the forward solve quantizes at pixel_scale_precision, so the computed likelihood is a staircase whose envelope derivative the implicit rule returns — naive FD below the stair width reads exactly zero (certification methodology below).

Known limitation (documented, follow-up filed): a FREE cosmology parameter cannot cross the custom_jvp boundary (Tracer registers cosmology as no_flatten aux → stale tracer → UnexpectedTracerError). Physically lossless for 2-plane positions fits; multi-plane free-cosmology gradient fits need the cosmology flattened into the Tracer pytree.

API Changes

No public API changes — PointSolver.solve signature and all values are unchanged (numpy path byte-identical; JAX regression literals in autolens_workspace_test re-verified unchanged). The new gradient behaviour engages only under JAX autodiff on the registered-model path, where gradients were previously identically zero.

Test Plan

• New numpy-only unit tests (test_autolens/point/triangles/test_implicit_diff.py): tangent linear algebra vs explicit solve, padded-row zeroing, unclamped near-critical divergence, analytic ±1 dθ/dθ_E for the on-axis spherical-isothermal images, numpy-path never touches implicit_diff.
• Full test_autolens/ suite: 492 passed.
• Downstream certification (workspace_test PR, linked): FD-certified gradients through FitPositionsImagePairAllSolved + FitPositionsImagePairAll, subgradient liveness for PairRepeatSolved, nested-autodiff fluxes+time-delays; all four jax_likelihood regression literals unchanged.

Generated by the PyAutoLabs agent workflow.