Skip to content

refactor: effective_einstein_radius latent drops its hardcoded seed fan #735

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jammy2211 merged 1 commit into from
Sep 11, 2026
Merged

refactor: effective_einstein_radius latent drops its hardcoded seed fan #735

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 10 additions & 11 deletions autolens/analysis/latent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)``.
"""
Expand All @@ -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
Expand Down
4 changes: 2 additions & 2 deletions test_autolens/analysis/test_latent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down Expand Up @@ -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"
)
Expand Down
Loading