refactor: einstein_radius_jit_from finds its own seed inside the JIT trace

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

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
 ):