From 03070e7424d5dc165f908810dd55599e9ee78637 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 11 Sep 2026 01:44:11 +0000 Subject: [PATCH] refactor: einstein_radius_jit_from finds its own seed inside the JIT trace MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `init_guess` is now optional. When omitted, the new `LensCalc._seed_via_coarse_grid_argmin` evaluates the tangential eigen value on a coarse uniform grid (25x25, ±3 arcsec by default) from the raw Hessian and returns the cell of minimum |eigen value| as the Newton seed for `jax_zero_contour.ZeroSolver`. It is xp-generic, masks non-finite cells (the odd-shaped grid puts a cell on the singular centre), and uses no skimage and no Python control flow on traced values, so the jit path no longer needs a caller-supplied seed. `seed_grid_shape` / `seed_grid_extent` let off-centre or cluster-scale callers widen the search without leaving the JIT-friendly path. The explicit `init_guess` path, the closure cache and the jax_zero_contour-missing soft-fail are unchanged. The signature-lock test flips to assert the argument is optional; numpy-path unit tests cover the centred, off-centre and radial cases and the soft-fail without a seed. Closes #614. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014JqpPWfk6nUAqnyv4P5oNM --- autogalaxy/operate/lens_calc.py | 162 +++++++++++++++++++- test_autogalaxy/operate/test_deflections.py | 125 ++++++++++++++- 2 files changed, 272 insertions(+), 15 deletions(-) diff --git a/autogalaxy/operate/lens_calc.py b/autogalaxy/operate/lens_calc.py index 80d93c38..b1cc4df9 100644 --- a/autogalaxy/operate/lens_calc.py +++ b/autogalaxy/operate/lens_calc.py @@ -1301,6 +1301,120 @@ def _init_guess_from_coarse_grid( seeds = [curve[len(curve) // 2] for curve in coarse_curves] return jnp.array(seeds) + def _seed_via_coarse_grid_argmin( + self, + kind="tangential", + grid_shape=(25, 25), + grid_extent=3.0, + pixel_scales=(0.05, 0.05), + xp=np, + ): + """ + Returns a single Newton seed position near the critical curve, found by + an argmin over a coarse grid. + + This is the JIT-traceable replacement for ``_init_guess_from_coarse_grid`` + on the jit path. That method runs marching squares (``skimage``) over the + coarse eigen-value map and returns one seed per distinct contour segment, + which is strictly more informative -- but marching squares is Python + control flow over concrete array values and cannot be traced, so it is + unusable inside ``jax.jit``. + + Here the eigen value is evaluated on the same coarse uniform grid and the + cell whose ``|eigen value|`` is smallest is returned as the seed. The grid + itself is a static NumPy constant even under jit, because ``grid_shape`` + and ``grid_extent`` are Python values; only the eigen values evaluated on + it are traced. ``xp.argmin`` and indexing a constant array with the + resulting traced scalar are both jit-safe, so the whole body traces with + no Python branch on a traced value. + + The eigen value is computed from the raw Hessian tuple returned by + ``hessian_from`` rather than through ``tangential_eigen_value_from``, + because the latter returns an ``aa.Array2D`` under NumPy and a raw array + under ``jax.numpy``; the raw-Hessian route is ``xp``-generic in one body. + The shear terms are symmetrised exactly as in ``_make_eigen_fn``. + + Cells whose eigen value is not finite are masked out before the argmin. + The default ``25 x 25`` grid has an odd shape, so one cell sits exactly on + the origin, where an isothermal centre is singular; without the mask the + argmin would follow that NaN instead of the critical curve. + + Limitations + ----------- + Exactly **one** seed is returned -- the global minimum of + ``|eigen value|`` over the coarse grid. A model with several distinct + tangential critical curves (a cluster-scale lens, a multi-component + deflector) will therefore be seeded on only one of them. Callers with + such models should pass ``init_guess`` explicitly. + + Cost is ``grid_shape[0] * grid_shape[1]`` Hessian evaluations (625 by + default), which is negligible next to the contour trace that follows. + + Parameters + ---------- + kind + ``"tangential"`` (eigen value = ``1 - kappa - |gamma|``) or + ``"radial"`` (``1 - kappa + |gamma|``). + grid_shape + Number of pixels along each axis of the coarse evaluation grid. + grid_extent + Half-width of the coarse grid in arc-seconds. This is the knob for + off-centre or cluster-scale lenses: the critical curve must fall + inside the grid for the argmin to find it. + pixel_scales + Accepted for signature symmetry with ``einstein_radius_jit_from``, + but **not used**: the JAX Hessian path takes its pixel scales from + the coarse grid itself (``_hessian_via_jax`` reads + ``grid.pixel_scales``), and the NumPy path uses finite differences + with its own adaptive step. + xp + The array module (``numpy`` or ``jax.numpy``), forwarded to + ``hessian_from``. + + Returns + ------- + Array of shape ``(1, 2)`` + The ``(y, x)`` arc-second coordinate of the coarse cell closest to + the critical curve, ready to pass as ``init_guess``. + """ + grid = aa.Grid2D.uniform( + shape_native=grid_shape, + pixel_scales=( + 2.0 * grid_extent / grid_shape[0], + 2.0 * grid_extent / grid_shape[1], + ), + ) + + if xp is np: + # The Richardson finite-difference path warns on the singular + # centre cell, which is exactly the cell the finite mask below + # discards. Expected, and not worth surfacing to the caller. + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + hessian_yy, hessian_xy, hessian_yx, hessian_xx = self.hessian_from( + grid=grid, xp=xp + ) + else: + hessian_yy, hessian_xy, hessian_yx, hessian_xx = self.hessian_from( + grid=grid, xp=xp + ) + + convergence = 0.5 * (hessian_yy + hessian_xx) + gamma_1 = 0.5 * (hessian_xx - hessian_yy) + gamma_2 = 0.5 * (hessian_xy + hessian_yx) # symmetrised + shear = xp.sqrt(gamma_1**2 + gamma_2**2) + + sign = -1.0 if kind == "tangential" else 1.0 + eigen = 1.0 - convergence + sign * shear + + score = xp.where(xp.isfinite(eigen), xp.abs(eigen), xp.inf) + idx = xp.argmin(score) + + grid_values = grid.array if hasattr(grid, "array") else grid + seed = xp.asarray(grid_values)[idx] + + return xp.reshape(seed, (1, 2)) + def _critical_curve_list_via_zero_contour( self, kind: str, @@ -1684,12 +1798,14 @@ def einstein_radius_via_zero_contour_from( def einstein_radius_jit_from( self, - init_guess, + init_guess=None, delta: float = 0.05, N: int = 500, pixel_scales: Tuple[float, float] = (0.05, 0.05), tol: float = 1e-6, max_newton: int = 5, + seed_grid_shape: Tuple[int, int] = (25, 25), + seed_grid_extent: float = 3.0, ): """ JIT-friendly Einstein radius from the tangential critical curve. @@ -1699,8 +1815,18 @@ def einstein_radius_jit_from( ``LATENT_BATCH_MODE='jit'``). Unlike ``einstein_radius_via_zero_contour_from``, this method: - - Requires an explicit ``init_guess`` — no marching-squares seed search - (which would require ``skimage`` and break the JAX trace). + - Finds its own Newton seed inside the trace when ``init_guess`` is not + given: the eigen value is evaluated on a coarse + ``seed_grid_shape`` grid of half-width ``seed_grid_extent`` arc-seconds + and the cell with the smallest ``|tangential eigen value|`` is taken as + the seed (``_seed_via_coarse_grid_argmin``). No ``skimage``, no + marching squares, and no Python control flow on traced values — unlike + ``_init_guess_from_coarse_grid``, which cannot be traced. The seed + search returns a **single** seed (the global minimum), so a model with + several distinct tangential critical curves is seeded on only one of + them; pass ``init_guess`` explicitly for those. It costs + ``seed_grid_shape[0] * seed_grid_shape[1]`` ``jacfwd`` Hessian + evaluations (625 by default), negligible next to the contour trace. - Skips ``ZeroSolver.path_reduce`` (variable-length output, not jit-able). - Computes the enclosed area directly from the raw NaN-padded paths array via a JAX-vectorised shoelace formula. @@ -1716,11 +1842,13 @@ def einstein_radius_jit_from( Parameters ---------- init_guess - JAX or NumPy array of shape ``(n_seeds, 2)`` — seed positions - (y, x) near the expected critical curve. For typical galaxy-scale - lenses centred on the image, a single seed at ``[[1.0, 0.0]]`` - works; for clusters or off-centre lenses pass a small fan of - seeds (e.g. four at cardinal positions). + Optional JAX or NumPy array of shape ``(n_seeds, 2)`` — seed + positions (y, x) near the expected critical curve. ``None`` (the + default) runs the in-trace coarse-grid argmin seed search described + above, which is correct for any single-critical-curve model whose + curve falls inside ``seed_grid_extent``. Pass seeds explicitly for a + model with several distinct tangential critical curves (e.g. four at + cardinal positions), which the single-seed search cannot cover. delta Arc-second step size along the contour. Forwarded to ``ZeroSolver``. N @@ -1732,6 +1860,15 @@ def einstein_radius_jit_from( Newton's method convergence tolerance. max_newton Maximum Newton iterations per step. + seed_grid_shape + Number of pixels along each axis of the coarse grid used by the + in-trace seed search. Ignored when ``init_guess`` is given. + seed_grid_extent + Half-width in arc-seconds of that coarse grid. This is the knob for + off-centre or cluster-scale lenses: the critical curve must fall + inside the grid for the argmin to find it (the default ±3" covers a + galaxy-scale lens centred near the origin). Ignored when + ``init_guess`` is given. Returns ------- @@ -1747,6 +1884,15 @@ def einstein_radius_jit_from( from jax_zero_contour import ZeroSolver import jax.numpy as jnp + if init_guess is None: + init_guess = self._seed_via_coarse_grid_argmin( + kind="tangential", + grid_shape=seed_grid_shape, + grid_extent=seed_grid_extent, + pixel_scales=pixel_scales, + xp=jnp, + ) + init_guess = jnp.atleast_2d(jnp.asarray(init_guess)) # Reuse the (f, ZeroSolver) pair from PR #434's cache so the warm diff --git a/test_autogalaxy/operate/test_deflections.py b/test_autogalaxy/operate/test_deflections.py index 70a86fb3..5365f657 100644 --- a/test_autogalaxy/operate/test_deflections.py +++ b/test_autogalaxy/operate/test_deflections.py @@ -700,8 +700,10 @@ def test__einstein_radius_jit_from__method_exists_with_expected_signature(): """ `LensCalc.einstein_radius_jit_from` is the JIT-friendly Einstein-radius helper added for `Analysis.LATENT_BATCH_MODE='jit'` (see PyAutoFit). It - must remain on the class with `init_guess` as the only required - argument — pipelines pass a static `jnp.array([[1.0, 0.0]])` etc. + must remain on the class, and `init_guess` must be *optional* — the jit + path now finds its own Newton seed in-trace via + `_seed_via_coarse_grid_argmin`, so callers no longer have to supply a + static seed array. Callers with multi-curve models may still pass one. Runtime JAX behaviour is exercised in the workspace_test integration suite (no JAX in library unit tests per project policy). @@ -715,14 +717,94 @@ def test__einstein_radius_jit_from__method_exists_with_expected_signature(): sig = inspect.signature(LensCalc.einstein_radius_jit_from) params = sig.parameters assert "init_guess" in params - assert params["init_guess"].default is inspect.Parameter.empty, ( - "init_guess must be a required argument — there is no JAX-friendly " - "default (the legacy `_init_guess_from_coarse_grid` uses skimage)" - ) - for kw in ("delta", "N", "pixel_scales", "tol", "max_newton"): + assert params["init_guess"].default is None, ( + "init_guess must be optional — the jit path finds its own seed " + "in-trace with `_seed_via_coarse_grid_argmin` (a JAX-native argmin " + "over a coarse eigen-value grid, no skimage)" + ) + for kw in ( + "delta", + "N", + "pixel_scales", + "tol", + "max_newton", + "seed_grid_shape", + "seed_grid_extent", + ): assert kw in params, f"einstein_radius_jit_from missing kwarg {kw!r}" +def test__seed_via_coarse_grid_argmin__centred_isothermal(): + """ + The in-trace seed finder returns the coarse-grid cell whose + |tangential eigen value| is smallest. For an SIS of Einstein radius 1.0" + centred on the origin the tangential critical curve is the circle of + radius 1.0", so the seed must land within one coarse cell width of it. + + With the defaults the cell width is `2 * 3.0 / 25 = 0.24"`. + + The centre cell of the (odd) 25x25 grid sits exactly on the singular + profile centre; the finite-mask in the helper is what stops the argmin + following that NaN instead of the critical curve. + """ + mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.0) + od = LensCalc.from_mass_obj(mp) + + seed = od._seed_via_coarse_grid_argmin(xp=np) + + assert seed.shape == (1, 2) + assert np.all(np.isfinite(seed)) + + cell_width = 2.0 * 3.0 / 25 + radius = np.hypot(seed[0, 0], seed[0, 1]) + assert abs(radius - 1.0) < cell_width + + +def test__seed_via_coarse_grid_argmin__off_centre_isothermal(): + """ + The seed search follows the model, not a hardcoded position. This is the + case a fixed ±1" cardinal fan of seeds cannot handle: an SIE of Einstein + radius 1.0" centred at (2.0, -1.5) has its tangential critical curve + nowhere near any of `[[1,0], [0,1], [-1,0], [0,-1]]`, so a fan-seeded + Newton solve starts from points 1.5-3.5" off the curve. + + `grid_extent=5.0` widens the coarse grid enough to contain the curve; + the cell width is then `2 * 5.0 / 25 = 0.4"`. + """ + mp = ag.mp.Isothermal(centre=(2.0, -1.5), ell_comps=(0.0, 0.0), einstein_radius=1.0) + od = LensCalc.from_mass_obj(mp) + + seed = od._seed_via_coarse_grid_argmin(grid_extent=5.0, xp=np) + + assert seed.shape == (1, 2) + assert np.all(np.isfinite(seed)) + + cell_width = 2.0 * 5.0 / 25 + radius = np.hypot(seed[0, 0] - 2.0, seed[0, 1] - (-1.5)) + assert abs(radius - 1.0) < cell_width + + +def test__seed_via_coarse_grid_argmin__radial_kind__returns_finite_seed(): + """ + `kind="radial"` selects the `1 - kappa + |gamma|` eigen value. An SIS has + no radial critical curve at all (its radial eigen value never reaches + zero away from the singular centre), so there is no position to assert + against — the point of this test is only that the radial branch runs, + masks the singular centre cell, and returns a finite `(1, 2)` seed + rather than raising or handing back NaN. + + A model with a genuine radial curve (e.g. `ag.mp.PowerLaw` with a core, + or an NFW) is exercised on the JAX path in the integration suite. + """ + mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.0) + od = LensCalc.from_mass_obj(mp) + + seed = od._seed_via_coarse_grid_argmin(kind="radial", xp=np) + + assert seed.shape == (1, 2) + assert np.all(np.isfinite(seed)) + + def test__analysis_dataset__latent_batch_mode_is_jit(): """ PyAutoGalaxy's `AnalysisDataset` overrides PyAutoFit's default @@ -819,6 +901,35 @@ def test__einstein_radius_jit_from__missing_jax_zero_contour__returns_nan_and_wa assert len(matching) == 1 +def test__einstein_radius_jit_from__missing_dep__no_init_guess__returns_nan_and_warns( + monkeypatch, caplog +): + """ + The soft-fail early-out precedes the seed search, so a call with no + `init_guess` must behave exactly as the seeded call does when + `jax_zero_contour` is missing: NaN back, one warning, and no attempt to + import JAX or run `_seed_via_coarse_grid_argmin` (625 Hessians thrown + away would be a silly way to reach the same NaN). + """ + _lens_calc_module._OPTIONAL_DEP_WARNED.discard("einstein_radius_jit_from") + _patch_missing_jax_zero_contour(monkeypatch) + + mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=2.0) + od = LensCalc.from_mass_obj(mp) + + def _fail(*args, **kwargs): + raise AssertionError("seed search ran despite the missing-dependency early-out") + + monkeypatch.setattr(od, "_seed_via_coarse_grid_argmin", _fail) + + with caplog.at_level(logging.WARNING, logger=_lens_calc_module.__name__): + result = od.einstein_radius_jit_from() + + assert math.isnan(result) + matching = [r for r in caplog.records if "einstein_radius_jit_from" in r.message] + assert len(matching) == 1 + + def test__tangential_critical_curve_list_via_zero_contour__missing_dep__returns_empty( monkeypatch, caplog ):