From 61dbd31cba04e4361676eae21b7a058e0d43f26a Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 31 Jul 2026 11:54:19 +0100 Subject: [PATCH] =?UTF-8?q?feat:=20PointSolver=20implicit-diff=20gradients?= =?UTF-8?q?=20=E2=80=94=20custom=5Fjvp=20fixed-point=20rule=20(phase=205?= =?UTF-8?q?=20of=20#657)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 --- autolens/point/fit/fluxes.py | 8 +- .../point/fit/positions/image/pair_repeat.py | 8 + autolens/point/solver/implicit_diff.py | 198 ++++++++++++++++++ autolens/point/solver/point_solver.py | 117 +++++++++-- .../point/triangles/test_implicit_diff.py | 118 +++++++++++ 5 files changed, 422 insertions(+), 27 deletions(-) create mode 100644 autolens/point/solver/implicit_diff.py create mode 100644 test_autolens/point/triangles/test_implicit_diff.py diff --git a/autolens/point/fit/fluxes.py b/autolens/point/fit/fluxes.py index 79deccf33..5052b3a92 100644 --- a/autolens/point/fit/fluxes.py +++ b/autolens/point/fit/fluxes.py @@ -9,6 +9,7 @@ is computed against the observed flux values and noise map, contributing to the total ``FitPointDataset`` log likelihood. """ + import numpy as np from typing import Optional @@ -115,12 +116,7 @@ def model_data(self): are used. """ return aa.ArrayIrregular( - values=self._xp.array( - [ - magnification * self.profile.flux - for magnification in self.magnifications_at_positions - ] - ) + values=self.magnifications_at_positions.array * self.profile.flux ) @property diff --git a/autolens/point/fit/positions/image/pair_repeat.py b/autolens/point/fit/positions/image/pair_repeat.py index f37621f0b..a691725c4 100644 --- a/autolens/point/fit/positions/image/pair_repeat.py +++ b/autolens/point/fit/positions/image/pair_repeat.py @@ -66,6 +66,14 @@ class FitQuiet(al.FitPositionsImagePairRepeat): Penalty terms are normalized by the mean of the position noise-map, and all policy computations are fixed-shape (NaN-padded model positions are masked, never dropped), so the fit remains JAX-compilable. + **Gradient (subgradient) semantics**: under JAX autodiff the solver positions carry implicit-rule + gradients (see ``autolens.point.solver.implicit_diff``), but this fit's likelihood is only + piecewise-smooth in them: nearest-neighbour pairing (``min`` over the distance matrix) and the + unmatched-model policy masks are piecewise-constant selections, so autodiff returns the subgradient of + the currently-selected pairing and the policy mask contributes zero derivative. Gradients are exact + between pairing/policy flips and undefined at the measure-zero flip events — the same contract as the + all-pairs mixture fit, which is smooth in the pairings and preferred for gradient-based searches. + Point source fitting uses name pairing, whereby the `name` of the `Point` object is paired to the name of the point source dataset to ensure that point source datasets are fitted to the correct point source. diff --git a/autolens/point/solver/implicit_diff.py b/autolens/point/solver/implicit_diff.py new file mode 100644 index 000000000..263f01694 --- /dev/null +++ b/autolens/point/solver/implicit_diff.py @@ -0,0 +1,198 @@ +""" +Implicit differentiation of the ``PointSolver`` triangle-tiling solve. + +The solver's refinement iteration is not differentiable: triangle containment is a hard +boolean test, so autodiff through the solve is identically zero (this was the correct +verdict of the 2026-07 gradient audit). The solved image positions themselves, however, +are smooth functions of the lens parameters between image-count events — each solved +image ``θᵢ`` satisfies the lens equation ``β = θᵢ − α(θᵢ, p)`` at fixed source-plane +position ``β``, so the inverse function theorem gives the exact tangent + + ``Aᵢ dθᵢ = dα|_{θᵢ} + dβ`` with ``Aᵢ = ∂β/∂θ|_{θᵢ} = I − ∂α/∂θ|_{θᵢ}`` + +This is the same mechanism as gravity.jl (Lombardi 2024, arXiv:2406.15280, Eq. 30): +differentiate *at* the solution, never *through* the solver iteration. It is applied here +as a ``jax.custom_jvp`` whose rule is linear in the tangents, so JAX transposes it +automatically and reverse mode (``value_and_grad``) works. + +Differentiability contract (mirrors the frozen-Delaunay-tables contract): + +- **Between image-count events** the rule is exact up to the solver's residual: the + forward positions are quantized at ``pixel_scale_precision``, so the computed + likelihood is a staircase whose plateaus straddle the smooth exact-solve envelope. + The implicit gradient is the derivative of that envelope — the quantity a gradient + search wants. Finite differences below the stair width read exactly zero; certify + against a fine-precision solver with a per-parameter FD step sweep + (``autolens_workspace_test/scripts/point_source/jax_grad/gradient.py``). +- **At image-count events** (caustic crossings, magnification-threshold flips, triangle + containment flips) the likelihood itself is discontinuous and no method has a + gradient — the same measure-zero seam class as Delaunay re-wiring events; a sampler + almost surely never lands on one. +- **Near critical curves** ``det(Aᵢ) → 0`` and the tangent legitimately diverges. No + clamping or regularization is applied — the instability is surfaced (large or + non-finite gradients), never silenced. + +Padded rows (the ``inf`` sentinels of the fixed ``MAX_CONTAINING_SIZE`` output) are +constants of the output shape; their tangent is forced to zero so they cannot inject +NaNs into the batch. + +Known limitation — free cosmology parameters: ``Tracer`` is registered with +``cosmology`` as ``no_flatten`` aux, so a cosmology carrying traced parameters (a free +``H0``) crosses the ``custom_jvp`` boundary as a stale tracer and raises +``UnexpectedTracerError``. Positions-only 2-plane fits lose nothing (image positions do +not depend on H0); multi-plane gradient fits with free cosmology require flattening the +cosmology into the Tracer pytree — recorded as a follow-up on PyAutoLens#657. + +JAX is imported inside the factory, never at module level, per the workspace JAX rules. +""" + +from typing import Tuple + +import autoarray as aa + + +def implicit_tangents_from(jac_alpha, dalpha, dbeta, finite, xp): + """ + The pure tangent linear algebra of the implicit rule, shared by the JAX JVP rule and + the NumPy unit tests. + + Solves ``(I − ∂α/∂θ) dθ = dα + dβ`` row-by-row and zeroes the tangent of padded + (non-finite) rows. + + Parameters + ---------- + jac_alpha + The per-image Jacobian ``∂α/∂θ`` of shape ``(n, 2, 2)`` evaluated at the solved + positions. + dalpha + The tangent of the deflection field at the (fixed) solved positions, shape + ``(n, 2)``. + dbeta + The tangent of the source-plane coordinate, shape ``(2,)``. + finite + Boolean mask of shape ``(n,)`` — True for real solved images, False for padded + sentinel rows. + xp + The array module (``numpy`` or ``jax.numpy``). + + Returns + ------- + The tangent of the solved positions, shape ``(n, 2)``; zero on padded rows. Near a + critical curve ``det(I − ∂α/∂θ) → 0`` and the returned tangent diverges — this is + deliberate (see the module docstring). + """ + identity = xp.eye(2) + a_mat = identity[None, :, :] - jac_alpha + rhs = dalpha + dbeta[None, :] + dtheta = xp.linalg.solve(a_mat, rhs[..., None])[..., 0] + return xp.where(finite[:, None], dtheta, 0.0) + + +def tracer_is_jax_compatible(tracer) -> bool: + """ + Whether ``tracer`` can cross a ``jax.custom_jvp`` boundary: it (and everything it + holds) must flatten to genuine JAX-value leaves. + + True on the production model-fit path, where ``autofit.jax.register_model`` has + registered the model classes and ``AnalysisPoint`` has registered ``Tracer`` — the + leaves are arrays, scalars or JAX tracers. False for a hand-built ``Tracer`` whose + ``Galaxy``/profile objects are unregistered (e.g. simulator scripts calling + ``PointSolver.solve`` directly with ``use_jax=True``); the solve then falls back to + the plain forward path, whose behaviour (including its zero gradient) is unchanged. + """ + import numbers + + import numpy as np + from jax import tree_util + import jax + + for leaf in tree_util.tree_leaves(tracer): + if isinstance(leaf, (jax.Array, jax.core.Tracer, np.ndarray, np.generic)): + continue + if isinstance(leaf, numbers.Number): + continue + return False + return True + + +def solve_padded_factory(solver, plane_redshift, plane_index: int, xp): + """ + Build the ``jax.custom_jvp``-wrapped padded solve for one ``(solver, plane)`` pair. + + The solver instance and plane identifiers are closed over rather than passed as + arguments: ``PointSolver`` is not a registered pytree (the fit-class pytree + registrations deliberately carry it as ``no_flatten`` aux data), so it cannot cross + a ``custom_jvp`` boundary as a JAX value. The returned function takes only the two + differentiable inputs. + + Parameters + ---------- + solver + The ``PointSolver`` whose forward solve is wrapped. + plane_redshift + Forwarded verbatim to the forward solve (``None`` for the source plane). + plane_index + The plane index the deflections in the tangent rule are computed between + (``plane_i=0`` and this), derived from ``plane_redshift`` exactly as + ``AbstractSolver._plane_grid`` derives it — ``-1`` when ``plane_redshift`` is + ``None``. + xp + The JAX array module (``jax.numpy``); the factory is only reached on the JAX + path. + + Returns + ------- + A function ``(tracer, beta) -> (MAX_CONTAINING_SIZE, 2) array`` whose JVP applies + the implicit fixed-point rule. + """ + import jax + + def deflections_from(positions_array, tracer): + grid = aa.Grid2DIrregular(values=positions_array, xp=xp) + deflections = tracer.deflections_between_planes_from( + grid=grid, plane_i=0, plane_j=plane_index, xp=xp + ) + return deflections.array if hasattr(deflections, "array") else deflections + + @jax.custom_jvp + def solve_padded(tracer, beta: Tuple[float, float]): + return solver._solve_array( + tracer=tracer, + source_plane_coordinate=(beta[0], beta[1]), + xp=xp, + plane_redshift=plane_redshift, + remove_infinities=False, + ) + + @solve_padded.defjvp + def solve_padded_jvp(primals, tangents): + tracer, beta = primals + tracer_dot, beta_dot = tangents + + theta = solve_padded(tracer, beta) + finite = xp.isfinite(theta).all(axis=1) + theta_safe = xp.where(finite[:, None], theta, 0.0) + + def deflections_single(position, tracer_): + return deflections_from(position[None, :], tracer_)[0] + + jac_alpha = jax.vmap( + lambda position: jax.jacfwd(deflections_single, argnums=0)(position, tracer) + )(theta_safe) + + _, dalpha = jax.jvp( + lambda tracer_: deflections_from(theta_safe, tracer_), + (tracer,), + (tracer_dot,), + ) + + dtheta = implicit_tangents_from( + jac_alpha=jac_alpha, + dalpha=dalpha, + dbeta=xp.asarray(beta_dot), + finite=finite, + xp=xp, + ) + return theta, dtheta + + return solve_padded diff --git a/autolens/point/solver/point_solver.py b/autolens/point/solver/point_solver.py index 86ab41eb0..aa13a4af2 100644 --- a/autolens/point/solver/point_solver.py +++ b/autolens/point/solver/point_solver.py @@ -16,6 +16,7 @@ sentinel value ``inf`` for JAX compatibility — these ``inf`` entries are stripped by default but can be retained for use inside a ``jax.jit``-traced function. """ + import logging import os from typing import Tuple, Optional @@ -84,6 +85,13 @@ def solve( Notes ----- + Gradients (JAX path): the returned positions carry a ``jax.custom_jvp`` implicit + fixed-point rule (``dθ = A⁻¹ (dα + dβ)`` at the solved positions — the gravity.jl + / Lombardi 2024 Eq. 30 mechanism), so ``jax.grad`` through a solver-chained + likelihood is exact between image-count events instead of identically zero. See + ``autolens.point.solver.implicit_diff`` for the differentiability contract + (staircase quantization, caustic-crossing seams, near-critical divergence). + Smoke-test short-circuit (``PYAUTO_SMALL_DATASETS``): the triangle-tiling solve is the dominant cost in many simulator scripts and is meaningless on the downsized grids used for fast smoke tests. When ``PYAUTO_SMALL_DATASETS=1`` is @@ -111,6 +119,93 @@ def solve( if os.environ.get("PYAUTO_SMALL_DATASETS") == "1": return aa.Grid2DIrregular(values=[(1.0, 0.0), (0.0, 1.0)]) + if xp is not np: + from .implicit_diff import solve_padded_factory, tracer_is_jax_compatible + + if not tracer_is_jax_compatible(tracer): + # Unregistered tracer (e.g. a simulator script's hand-built Tracer): + # cannot cross the custom_jvp boundary. Plain forward solve, gradient + # behaviour unchanged (identically zero, as before this rule existed). + solution = self._solve_array( + tracer=tracer, + source_plane_coordinate=source_plane_coordinate, + xp=xp, + plane_redshift=plane_redshift, + remove_infinities=remove_infinities, + ) + return aa.Grid2DIrregular(solution) + + plane_index = ( + -1 + if plane_redshift is None + else tracer.plane_index_via_redshift_from(redshift=plane_redshift) + ) + solve_padded = solve_padded_factory( + solver=self, + plane_redshift=plane_redshift, + plane_index=plane_index, + xp=xp, + ) + beta = xp.stack( + [ + xp.asarray(source_plane_coordinate[0]), + xp.asarray(source_plane_coordinate[1]), + ] + ) + solution = solve_padded(tracer, beta) + + if remove_infinities: + solution = solution[~xp.isinf(solution).any(axis=1)] + + return aa.Grid2DIrregular(solution) + + solution = self._solve_array( + tracer=tracer, + source_plane_coordinate=source_plane_coordinate, + xp=xp, + plane_redshift=plane_redshift, + remove_infinities=remove_infinities, + ) + + # Warn on the *final* result rather than only on the empty branch inside + # `_solve_array`, because there are two distinct routes to an empty answer: + # + # 1. no triangle contained the coordinate -> `filtered_means` is already length 0 + # 2. every candidate failed the magnification threshold -> `_filter_low_magnification` + # preserves the length and writes NaN rows, which become `inf` and are then stripped + # by `remove_infinities`, landing here at length 0 + # + # Only reachable on the NumPy path: the JAX path keeps its padded static shape, so + # `len(solution)` is non-zero there and this reads no traced value. + if len(solution) == 0: + + logger.warning( + f"PointSolver.solve found no images for source-plane coordinate " + f"{tuple(source_plane_coordinate)}, so an empty grid is returned. This means " + f"either that the coordinate lies outside the region traced by the image-plane " + f"grid, or that every candidate image was rejected by " + f"`magnification_threshold` (currently {self.magnification_threshold})." + ) + + return aa.Grid2DIrregular(solution) + + def _solve_array( + self, + tracer: Tracer, + source_plane_coordinate: Tuple[float, float], + xp, + plane_redshift: Optional[float] = None, + remove_infinities: bool = False, + ): + """ + The raw solve: triangle refinement, magnification filtering and sentinel padding, + returning the bare ``xp`` positions array (no ``Grid2DIrregular`` wrap, no logging). + + On the JAX path this is the primal of the ``custom_jvp`` wrapper built by + ``implicit_diff.solve_padded_factory`` — gradients are supplied by the implicit + fixed-point rule at the solved positions, never by differentiating through the + refinement iteration (see ``implicit_diff``'s module docstring for the contract). + """ kept_triangles = super().solve_triangles( tracer=tracer, shape=Point(*source_plane_coordinate), @@ -155,24 +250,4 @@ def solve( solution = solution[~xp.isinf(solution).any(axis=1)] - # Warn on the *final* result rather than only on the branch above, because there are two - # distinct routes to an empty answer and only one goes through it: - # - # 1. no triangle contained the coordinate -> `filtered_means` is already length 0 - # 2. every candidate failed the magnification threshold -> `_filter_low_magnification` - # preserves the length and writes NaN rows, which become `inf` and are then stripped - # by `remove_infinities`, landing here at length 0 - # - # Only reachable on the NumPy path: the JAX path keeps its padded static shape, so - # `len(solution)` is non-zero there and this reads no traced value. - if len(solution) == 0: - - logger.warning( - f"PointSolver.solve found no images for source-plane coordinate " - f"{tuple(source_plane_coordinate)}, so an empty grid is returned. This means " - f"either that the coordinate lies outside the region traced by the image-plane " - f"grid, or that every candidate image was rejected by " - f"`magnification_threshold` (currently {self.magnification_threshold})." - ) - - return aa.Grid2DIrregular(solution) + return solution diff --git a/test_autolens/point/triangles/test_implicit_diff.py b/test_autolens/point/triangles/test_implicit_diff.py new file mode 100644 index 000000000..c29718e74 --- /dev/null +++ b/test_autolens/point/triangles/test_implicit_diff.py @@ -0,0 +1,118 @@ +""" +NumPy-only tests of the implicit-diff tangent rule (`autolens.point.solver.implicit_diff`). + +The JAX behaviour (custom_jvp wiring, value_and_grad through a solver-chained likelihood, +finite-difference certification) is validated by the workspace_test scripts +(`scripts/point_source/jax_grad/gradient.py`) per the no-JAX-in-unit-tests rule; these +tests pin the pure linear algebra and the numpy-path invariants. +""" + +import numpy as np +import pytest + +import autogalaxy as ag +from autolens import PointSolver, Tracer +from autolens.point.solver import implicit_diff + + +def test_implicit_tangents_solve_the_linear_system(): + rng = np.random.default_rng(3) + n = 5 + jac_alpha = 0.3 * rng.standard_normal((n, 2, 2)) + dalpha = rng.standard_normal((n, 2)) + dbeta = rng.standard_normal(2) + finite = np.array([True, True, False, True, False]) + + dtheta = implicit_diff.implicit_tangents_from( + jac_alpha=jac_alpha, dalpha=dalpha, dbeta=dbeta, finite=finite, xp=np + ) + + a_mat = np.eye(2)[None] - jac_alpha + for i in range(n): + if finite[i]: + np.testing.assert_allclose( + a_mat[i] @ dtheta[i], dalpha[i] + dbeta, rtol=1e-12 + ) + else: + np.testing.assert_array_equal(dtheta[i], 0.0) + + +def test_implicit_tangents_near_critical_diverge_unclamped(): + # det(I - J) -> 0: the tangent must diverge with the true solve, never be clamped. + eps = 1e-12 + jac_alpha = np.array([[[1.0 - eps, 0.0], [0.0, 0.0]]]) + dalpha = np.array([[1.0, 0.0]]) + dbeta = np.zeros(2) + finite = np.array([True]) + + dtheta = implicit_diff.implicit_tangents_from( + jac_alpha=jac_alpha, dalpha=dalpha, dbeta=dbeta, finite=finite, xp=np + ) + # rtol reflects float64 conditioning of the near-singular solve, not the rule + np.testing.assert_allclose(dtheta[0, 0], 1.0 / eps, rtol=1e-4) + + +def test_isothermal_sph_on_axis_rule_is_analytic(): + """ + For a spherical isothermal lens the on-axis images of a source at radius b sit at + theta = +(theta_E + b) and -(theta_E - b): d theta / d theta_E is exactly +1 and -1. + The rule must reproduce this from the deflection field alone. + """ + theta_e = 1.3 + b = 0.07 + mass = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=theta_e) + + theta = np.array([[0.0, theta_e + b], [0.0, -(theta_e - b)]]) + + lens_calc = ag.LensCalc.from_mass_obj(mass_obj=mass) + jacobian = lens_calc.jacobian_from(grid=theta) + # jacobian_from returns A = I - d alpha / d theta in (x, y) ordering; convert to the + # rule's d alpha / d theta in (y, x) ordering (the solved.py re-index precedent). + a_xx, a_xy, a_yx, a_yy = ( + np.asarray(jacobian[i][j]) for i in (0, 1) for j in (0, 1) + ) + jac_alpha = np.stack( + [ + np.stack([1.0 - a_yy, -a_yx], axis=-1), + np.stack([-a_xy, 1.0 - a_xx], axis=-1), + ], + axis=-2, + ) + + # isothermal: alpha = theta_E * unit(theta), so d alpha / d theta_E = alpha / theta_E + alpha = np.asarray(mass.deflections_yx_2d_from(grid=ag.Grid2DIrregular(theta))) + dalpha = alpha / theta_e + + dtheta = implicit_diff.implicit_tangents_from( + jac_alpha=jac_alpha, + dalpha=dalpha, + dbeta=np.zeros(2), + finite=np.array([True, True]), + xp=np, + ) + + np.testing.assert_allclose(dtheta[0], [0.0, 1.0], atol=1e-4) + np.testing.assert_allclose(dtheta[1], [0.0, -1.0], atol=1e-4) + + +def test_numpy_solve_path_never_touches_implicit_diff(monkeypatch): + def _boom(*args, **kwargs): + raise AssertionError("implicit_diff must not be reached on the numpy path") + + monkeypatch.setattr(implicit_diff, "solve_padded_factory", _boom) + + grid = ag.Grid2D.uniform(shape_native=(50, 50), pixel_scales=0.1) + solver = PointSolver.for_grid( + grid=grid, pixel_scale_precision=0.01, magnification_threshold=0.1 + ) + tracer = Tracer( + galaxies=[ + ag.Galaxy( + redshift=0.5, + mass=ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.0), + ), + ag.Galaxy(redshift=1.0), + ] + ) + result = solver.solve(tracer=tracer, source_plane_coordinate=(0.03, 0.03)) + assert len(np.asarray(result.array)) > 0