Affected Repositories

• PyAutoGalaxy (primary)
• PyAutoLens (follow-on, one-function change)

Work Classification

Library / Library (both legs are libraries; no workspace change — see Overview).

Sizing note

The prompt is declared Difficulty: too-large / Unattended: needs-slicing, and the Feature Agent proposed a 4-phase split on that header. The code survey does not support it: the library change is one new private helper (~40 lines), a signature change plus docstring on one method, and test updates; the PyAutoLens leg is a one-function edit. Treated as a single task with two sequenced PRs; the header should be re-sized to medium when the prompt is retired.

Branch Survey

worktree_check_conflict against PyAutoMind/active.md: no claim on PyAutoGalaxy or PyAutoLens (exit 0).

Suggested branch: feature/einstein-radius-jit-seed-finder
Worktree root: ~/Code/PyAutoLabs-wt/einstein-radius-jit-seed-finder/ (created by /start_library on a local checkout; a web session works on its clones instead)

Implementation Steps

1. autogalaxy/operate/lens_calc.py — new helper LensCalc._seed_via_coarse_grid_argmin(kind="tangential", grid_shape=(25, 25), grid_extent=3.0, pixel_scales=(0.05, 0.05), xp=np) placed next to _init_guess_from_coarse_grid (line ~1250).

   • Build the coarse grid as _init_guess_from_coarse_grid does: aa.Grid2D.uniform(shape_native=grid_shape, pixel_scales=(2*grid_extent/grid_shape[0], 2*grid_extent/grid_shape[1])). The grid is a static numpy constant even under jit (shape/extent are Python values).
   • Compute the eigen value from self.hessian_from(grid=grid, xp=xp) directly (hessian_yy, hessian_xy, hessian_yx, hessian_xx), using the same formula as _make_eigen_fn: convergence = 0.5*(Hyy+Hxx), gamma_1 = 0.5*(Hxx-Hyy), gamma_2 = 0.5*(Hxy+Hyx), eigen = 1 - convergence -/+ sqrt(gamma_1²+gamma_2²) (sign by kind). Working from the raw Hessian tuple avoids the aa.Array2D-vs-raw-array return of tangential_eigen_value_from, so one body serves both xp=np (unit tests) and xp=jnp (the trace).
   • Mask non-finite cells: score = xp.where(xp.isfinite(eigen), xp.abs(eigen), xp.inf). Needed because 25×25 is odd, so a cell sits exactly on the origin, where an isothermal centre is singular; without the guard argmin follows the NaN.
   • idx = xp.argmin(score); seed = xp.asarray(grid.array)[idx] (slim row idx is that cell's (y, x)), returned with shape (1, 2). Under jit idx is traced; indexing a constant array with a traced scalar is fine.
   • Docstring: single global-minimum seed, so a model with several tangential critical curves finds only one — callers with multi-curve models keep passing init_guess explicitly (accepted limitation from the prompt).

2. einstein_radius_jit_from (line ~1685) — signature becomes (self, init_guess=None, delta=0.05, N=500, pixel_scales=(0.05, 0.05), tol=1e-6, max_newton=5, seed_grid_shape=(25, 25), seed_grid_extent=3.0).

   • Keep the _maybe_optional_dep_warn early-out first (returns float("nan") before any JAX import).
   • After 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); then the existing jnp.atleast_2d(jnp.asarray(init_guess)) and the rest of the body unchanged (cache key, zero_contour_finder, shoelace area, max over seeds).
   • Rewrite the docstring: init_guess optional; describe the seed finder, its cost (625 jacfwd Hessian evaluations, negligible next to the contour trace), the single-seed limitation, and that seed_grid_extent is the knob for cluster-scale / off-centre lenses. Keep the "do not vmap" warning.
   • Do not touch _init_guess_from_coarse_grid, einstein_radius_via_zero_contour_from or the plotter path (out of scope per the prompt).

3. test_autogalaxy/operate/test_deflections.py

   • test__einstein_radius_jit_from__method_exists_with_expected_signature (line 699): flip the assertion — params["init_guess"].default is None; add seed_grid_shape and seed_grid_extent to the required-kwarg loop; fix the docstring that says pipelines pass a static seed.
   • New test__seed_via_coarse_grid_argmin__centred_isothermal: IsothermalSph(centre=(0,0), einstein_radius=1.0), default grid → seed radius within one cell width (0.24") of 1.0, finite, shape (1, 2).
   • New test__seed_via_coarse_grid_argmin__off_centre_isothermal: Isothermal(centre=(2.0, -1.5), einstein_radius=1.0, ...) with grid_extent=5.0 → distance from (2.0, -1.5) within one cell width of 1.0. This is the case the hardcoded fan cannot handle.
   • New test__seed_via_coarse_grid_argmin__radial_kind (cheap): kind="radial" returns a finite seed.
   • Extend test__einstein_radius_jit_from__missing_jax_zero_contour__returns_nan_and_warns with a no-init_guess call: still NaN + one warning (the early-out precedes the seed search).
   • No JAX in these tests (project policy: runtime JAX behaviour lives in workspace_test).

4. JAX verification (implementation session, not committed to the library): with jax + jax_zero_contour installed, an ad-hoc script checks (a) _seed_via_coarse_grid_argmin(xp=jnp) under jax.jit traces and returns the same cell as xp=np; (b) einstein_radius_jit_from() with no seed matches einstein_radius_jit_from(init_guess=fan) on an SIE to rel=1e-3, and matches einstein_radius_from(grid) to the tolerance PR feat: JIT-friendly Einstein radius helper + AnalysisDataset jit latent mode #435 used; (c) the off-centre SIE at (2.0, -1.5) converges with seed_grid_extent=5.0 where the ±1 fan does not. Record the numbers in the PR body. The Phase-B Euclid baseline (~2.10 arcsec, PYAUTO_TEST_MODE=1) is re-checked via the Euclid pipeline's tests/test_compute_latent_variable.py::test_effective_einstein_radius once the PyAutoLens leg is in.

5. PyAutoLens follow-on — autolens/analysis/latent.py::effective_einstein_radius (line ~418): delete the jnp.array([[1,0],[0,1],[-1,0],[0,-1]]) fan and call lens_calc.einstein_radius_jit_from(); update the docstring ("JAX path: einstein_radius_jit_from() — seeds found in-trace"). Second PR, merged after the PyAutoGalaxy PR (library-first gate); the fan remains valid against the released PyAutoGalaxy so PyAutoLens CI stays green either way.

6. Docs: no .rst/.md references to einstein_radius_jit_from outside lens_calc.py; the docstring is the documentation.

Key Files

• autogalaxy/operate/lens_calc.py — _init_guess_from_coarse_grid (~1250, skimage; unchanged), _make_eigen_fn (~1188, formula to mirror), hessian_from / _hessian_via_jax (~383/574, the xp plumbing the helper relies on), einstein_radius_jit_from (~1685, the change)
• test_autogalaxy/operate/test_deflections.py — signature lock at ~699, soft-fail tests at ~797
• autogalaxy/analysis/analysis/dataset.py — LATENT_BATCH_MODE = "jit" comment references this method (leave as is)
• PyAutoLens autolens/analysis/latent.py — effective_einstein_radius (~418), the only production caller
• euclid_strong_lens_modeling_pipeline config/latent.yaml — effective_einstein_radius: true; tests/test_compute_latent_variable.py — end-to-end check of the latent value

Trade-offs

• Single argmin seed vs. a top-k fan: argmin is deterministic, trace-friendly and sufficient for the unique tangential curve of galaxy-scale lenses; a fan would need a spatial de-duplication step inside the trace for little gain. Multi-curve models keep the explicit init_guess.
• Raw-Hessian eigen vs. tangential_eigen_value_from: the raw route duplicates four lines of formula but is xp-generic in one body and unit-testable without JAX.
• Default extent ±3 arcsec: covers the Euclid range (0.5–2 arcsec) with margin; clusters override via seed_grid_extent.

Testing

cd PyAutoGalaxy && python3 -m pytest test_autogalaxy/operate/test_deflections.py -q -n auto
cd PyAutoGalaxy && python3 -m pytest test_autogalaxy -q -n auto

einstein_radius_jit_from: replace static init_guess with a JAX-native seed finder

Type: refactor
Target: PyAutoGalaxy
Themes:

• jax-compile
  Difficulty: too-large
  Autonomy: supervised
  Priority: high
  Status: formalised
  Consequence: judge
  Review-minutes: 25
  Unattended: needs-slicing
  Filed: 2026-05-21 (backfilled from git)

Small follow-up to the Phase B work (PyAutoGalaxy #435, PyAutoFit #1288,
pipeline #15). Drops the requirement that callers pass a static
init_guess to LensCalc.einstein_radius_jit_from(...).

What init_guess does today

einstein_radius_jit_from(init_guess, ...) traces the tangential critical
curve using jax_zero_contour.ZeroSolver.zero_contour_finder. ZeroSolver
needs starting positions for Newton's method — points near the
expected zero-crossing of the tangential eigen-value. From each starting
point ("seed"), Newton iterates onto the curve and then walks along it.
init_guess is a JAX array of shape (n_seeds, 2) giving (y, x) arc-sec
coordinates for those seed points.

In the existing einstein_radius_via_zero_contour_from path (non-jit),
the seeds are discovered automatically by _init_guess_from_coarse_grid:
evaluate the eigen-value on a coarse 25×25 grid, run skimage's marching-
squares contour finder on the result, take the midpoint of each curve
segment. That approach uses skimage which is not JAX-traceable, so it
breaks compute_latent_samples' JIT trace — which is why
einstein_radius_jit_from requires the caller to provide init_guess
explicitly.

In the Euclid pipeline today, the workspace passes a hardcoded 4-seed
fan at ±1 arcsec from origin (util.py):

init_guess = jnp.array([[1.0, 0.0], [0.0, 1.0], [-1.0, 0.0], [0.0, -1.0]])

This works for the Euclid use case (preprocessed lenses are centred and
have Einstein radii roughly 0.5–2 arcsec, all within Newton's basin from
any of those seeds). It fails for:

• Off-centre lenses (e.g. group / cluster fits where the lens isn't at
  origin).
• Lenses with Einstein radii outside the basin from (±1, 0) / (0, ±1)
  (very-strong or very-weak lenses).
• Non-Euclid pipelines that don't know they need to pass init_guess.

What to do

Replace the required init_guess with a JAX-native seed finder that
runs inside the jit trace and doesn't need skimage.

Sketch:

def einstein_radius_jit_from(
    self,
    init_guess=None,
    delta=0.05, N=500,
    pixel_scales=(0.05, 0.05),
    tol=1e-6, max_newton=5,
    seed_grid_shape=(25, 25),
    seed_grid_extent=3.0,
):
    import jax.numpy as jnp

    if init_guess is None:
        # JAX-native seed search: argmin |eigen_value| on a coarse uniform grid.
        # No skimage, no Python control flow that depends on traced values.
        pixel_scale = 2.0 * seed_grid_extent / seed_grid_shape[0]
        grid = aa.Grid2D.uniform(
            shape_native=seed_grid_shape, pixel_scales=(pixel_scale, pixel_scale),
        )
        eigen = self.tangential_eigen_value_from(grid=grid, xp=jnp)
        # Flatten + argmin on |eigen|
        flat_idx = jnp.argmin(jnp.abs(eigen.native if hasattr(eigen, "native") else eigen))
        iy, ix = jnp.unravel_index(flat_idx, seed_grid_shape)
        y = -seed_grid_extent + (iy + 0.5) * pixel_scale
        x = -seed_grid_extent + (ix + 0.5) * pixel_scale
        init_guess = jnp.array([[y, x]])
    init_guess = jnp.atleast_2d(jnp.asarray(init_guess))
    # ... rest of the existing body unchanged

Implementation details to nail down:

• tangential_eigen_value_from(grid=grid, xp=jnp) — verify it threads
  xp=jnp correctly all the way down through convergence_2d_via_hessian_from
  and shear_yx_2d_via_hessian_from. The Euclid validation we did during
  the parent task suggests this should work; double-check by running the
  pipeline's start_here.py with the new default.
• eigen.native vs raw — tangential_eigen_value_from returns an aa.Array2D
  under xp=np but a raw jax.Array under xp=jnp (per the existing
  if xp is np guard in that function). The helper above should handle
  both shapes.
• Single seed vs fan — argmin returns ONE seed (the global minimum).
  For most models that's sufficient (tangential critical curve is
  unique). For pathological models with multiple critical curves we'd
  miss the secondary ones — accept this limitation; multi-curve cases
  can still pass init_guess explicitly.
• Seed-grid extent and shape — pick a sensible default. 3 arcsec half-
  width covers Euclid range; cluster fits would need wider. Could expose
  as args (seed_grid_extent, seed_grid_shape) so callers can override
  without leaving the JIT-friendly path.

Workspace cleanup

After the library lands, drop the hardcoded init_guess from
euclid_strong_lens_modeling_pipeline/util.py. The dispatch becomes:

if self._use_jax:
    effective_einstein_radius = lens_calc.einstein_radius_jit_from()
else:
    effective_einstein_radius = lens_calc.einstein_radius_from(
        grid=self.dataset.grids.lp,
    )

Verification

• Unit test: argmin-based seed matches the _init_guess_from_coarse_grid
  result to ~1 grid-cell width on an SIE tracer.
• End-to-end: re-run the Euclid pipeline's JAX-branch latent under
  PYAUTO_TEST_MODE=1. The latent.effective_einstein_radius value
  should match the Phase B baseline (~2.10 arcsec on the test dataset) to
  reasonable tolerance.
• Cluster check: when an off-centre lens model is used (manually set
  lens.mass.centre = (2.0, -1.5) or similar), confirm the new seed
  finder converges. The hardcoded (±1, 0) fan would fail this; the new
  default should succeed.

Out of scope

• Replacing the legacy _init_guess_from_coarse_grid skimage call —
  that path stays for the non-JIT einstein_radius_via_zero_contour_from
  and the plotter, and is fine where it lives.
• Vmap compatibility — the new helper is still jit-only by design (per
  upstream jax_zero_contour ZeroSolver vmap warning).

References

• PyAutoGalaxy feat: JIT-friendly Einstein radius helper + AnalysisDataset jit latent mode #435 — parent task, added einstein_radius_jit_from(init_guess, ...).
• complete.md::euclid-einstein-radius-zero-contour — context on why
  init_guess was required in the first place (skimage / find_contours
  blocking the JAX trace).
• feedback_jax_closure_cache_busts (memory) — relevant if benchmarking
  the new default's warm-call latency.