From eaff99e13292a0b2c1f90823a169df8c55f9226b Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 11 Sep 2026 01:37:42 +0000 Subject: [PATCH] refactor: effective_einstein_radius latent drops its hardcoded seed fan MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `LensCalc.einstein_radius_jit_from` (PyAutoGalaxy#614) now finds its own Newton seed inside the JIT trace, so the JAX path of the `effective_einstein_radius` latent calls it with no `init_guess`. The fixed 4-seed fan at ±1 arcsec from the origin could not converge for off-centre lenses or Einstein radii outside its basin. The test stubs mirror the new optional-argument signature. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014JqpPWfk6nUAqnyv4P5oNM --- autolens/analysis/latent.py | 21 ++++++++++----------- test_autolens/analysis/test_latent.py | 4 ++-- 2 files changed, 12 insertions(+), 13 deletions(-) diff --git a/autolens/analysis/latent.py b/autolens/analysis/latent.py index d0024aafc..fe53c611a 100644 --- a/autolens/analysis/latent.py +++ b/autolens/analysis/latent.py @@ -415,12 +415,15 @@ def effective_einstein_radius(fit, magzero, xp=np): """ Effective Einstein radius via the tangential critical curve. - JAX path: ``LensCalc.einstein_radius_jit_from(init_guess=fan)``, where - ``fan`` is a fixed 4-seed fan at ±1 arcsec from the lens centre — the - JIT-compatible variant required because ``ZeroSolver`` (line 1520 of - ``autogalaxy/operate/lens_calc.py``) uses ``lax.cond`` / - ``lax.while_loop`` early termination that is incompatible with - ``jax.vmap`` but fine under ``jax.jit``. + JAX path: ``LensCalc.einstein_radius_jit_from()`` with no seed — the + JIT-compatible variant, which since PyAutoGalaxy#614 finds its own Newton + seed inside the trace (argmin of the tangential eigen value on a coarse + grid, ±3 arcsec by default) instead of needing a caller-supplied + ``init_guess``. The earlier fixed 4-seed fan at ±1 arcsec from the origin + could not converge for off-centre lenses or Einstein radii outside its + basin. ``ZeroSolver`` uses ``lax.cond`` / ``lax.while_loop`` early + termination that is incompatible with ``jax.vmap`` but fine under + ``jax.jit``. NumPy path: ``LensCalc.einstein_radius_from(grid=fit.dataset.grids.lp)``. """ @@ -429,11 +432,7 @@ def effective_einstein_radius(fit, magzero, xp=np): try: lens_calc = LensCalc.from_mass_obj(fit.tracer) if xp is not np and _jax_zero_contour_available(): - import jax.numpy as jnp - init_guess = jnp.array( - [[1.0, 0.0], [0.0, 1.0], [-1.0, 0.0], [0.0, -1.0]] - ) - return lens_calc.einstein_radius_jit_from(init_guess=init_guess) + return lens_calc.einstein_radius_jit_from() return lens_calc.einstein_radius_from(grid=fit.dataset.grids.lp) except (ValueError, AttributeError): return xp.nan diff --git a/test_autolens/analysis/test_latent.py b/test_autolens/analysis/test_latent.py index 3efe41976..f353153cb 100644 --- a/test_autolens/analysis/test_latent.py +++ b/test_autolens/analysis/test_latent.py @@ -272,7 +272,7 @@ def einstein_radius_from(self, grid): calls["grid"] = grid return 1.234 - def einstein_radius_jit_from(self, init_guess): + def einstein_radius_jit_from(self, init_guess=None): calls["init_guess"] = init_guess raise AssertionError("numpy path must not use jit_from") @@ -318,7 +318,7 @@ def einstein_radius_from(self, grid): calls["grid"] = grid return 5.678 - def einstein_radius_jit_from(self, init_guess): + def einstein_radius_jit_from(self, init_guess=None): raise AssertionError( "jit path must not run when jax_zero_contour is missing" )