refactor: einstein_radius_jit_from finds its own seed inside the JIT trace

[Jammy2211]

Summary

LensCalc.einstein_radius_jit_from required callers to pass static Newton seed positions (init_guess) for jax_zero_contour.ZeroSolver, because the non-jit seed search (_init_guess_from_coarse_grid) goes through skimage marching squares and cannot be traced. The only production caller (PyAutoLens' effective_einstein_radius latent) hardcodes a 4-seed fan at ±1 arcsec from the origin.

This PR makes init_guess optional. When omitted, a new private helper _seed_via_coarse_grid_argmin evaluates the tangential eigen value on a coarse uniform grid (25×25, ±3 arcsec by default) straight from the raw Hessian tuple, masks non-finite cells, and returns the cell of minimum |eigen value| as the seed. It is xp-generic (numpy for unit tests, jax.numpy inside the trace), uses no skimage and no Python control flow on traced values. seed_grid_shape / seed_grid_extent let off-centre or cluster-scale callers widen the search without leaving the JIT path.

Unchanged: the explicit init_guess path, the (f, ZeroSolver) closure cache, the jax_zero_contour-missing soft-fail (still the first thing in the body, so a seedless call also returns NaN with one warning), _init_guess_from_coarse_grid, the non-jit zero-contour methods and the plotter path.

Closes #614. Follow-on: PyAutoLabs/PyAutoLens#735 drops the hardcoded fan (merge after this PR).

API Changes

LensCalc.einstein_radius_jit_from: init_guess is now optional (default None → in-trace seed search); two new keyword arguments seed_grid_shape=(25, 25) and seed_grid_extent=3.0. Every existing call that passes init_guess behaves exactly as before. New private helper LensCalc._seed_via_coarse_grid_argmin.
See full details below.

Test Plan

• test_autogalaxy/operate/test_deflections.py: 41 passed (was 37). Signature-lock test flipped to assert init_guess is optional; new numpy-only tests for the seed helper on a centred SIS (seed within one cell width of the 1.0" curve), an off-centre SIE at (2.0, −1.5) with grid_extent=5.0, the radial branch, and the soft-fail with no seed (asserts the seed search is never entered).
• Full library suite: 1195 passed.
• JAX verification (ad-hoc script, jax 0.11.1 / jax_zero_contour 2.0.0, not committed):
  1. _seed_via_coarse_grid_argmin(xp=jnp) traces under jax.jit and returns the same cell as xp=np ([[0.24, -1.2]] for an SIE with einstein_radius=1.2).
  2. Centred SIE: no seed → 1.188489; cardinal fan → 1.188980; relative difference 4.1e-4. (einstein_radius_from on a 200×200/0.05" grid gives 1.19706.)
  3. Off-centre SIE at (2.0, −1.5), einstein_radius=1.0: no seed with seed_grid_extent=5.0 → 0.971 (2.9% low, from coarse-cell seed quantisation). Note: the ±1 fan also converged on this particular model (1.004) — Newton walked onto the off-centre curve from the fan points. The seed finder's case is that it needs no caller knowledge of where the lens is, not that the fan always fails.
  4. jax.jit over a function that builds IsothermalSph(einstein_radius=r) from a traced scalar and calls einstein_radius_jit_from() works: r=0.8 → 0.793, r=1.5 → 1.494.
  5. Cold call 4.0 s, warm call 0.05 s.
• After merge and refactor: effective_einstein_radius latent drops its hardcoded seed fan PyAutoLens#735: Euclid pipeline tests/test_compute_latent_variable.py::test_effective_einstein_radius re-checks the latent value against the Phase B baseline.

Full API Changes (for automation & release notes)

Added

• LensCalc._seed_via_coarse_grid_argmin(kind="tangential", grid_shape=(25, 25), grid_extent=3.0, pixel_scales=(0.05, 0.05), xp=np) — private; JIT-traceable single-seed finder (argmin of |eigen value| on a coarse grid, non-finite cells masked). pixel_scales accepted for signature symmetry only.

Changed Signature

• LensCalc.einstein_radius_jit_from(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) — init_guess was required; two new trailing keyword arguments.

Changed Behaviour

• einstein_radius_jit_from() with no init_guess runs the in-trace seed search instead of raising TypeError. A model with several distinct tangential critical curves is seeded on only one of them; pass init_guess explicitly for those.

Migration

• Before: lens_calc.einstein_radius_jit_from(init_guess=jnp.array([[1.0, 0.0], [0.0, 1.0], [-1.0, 0.0], [0.0, -1.0]]))
• After: lens_calc.einstein_radius_jit_from() (or seed_grid_extent=<half-width> for off-centre / cluster-scale lenses). The old call still works.

Generated by the PyAutoLabs agent workflow.

🤖 Generated with Claude Code

https://claude.ai/code/session_014JqpPWfk6nUAqnyv4P5oNM