From a7b6bdcdb3d5bda97687c066557e34bcda9d19a4 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 30 Oct 2026 09:32:48 +0000 Subject: [PATCH 1/2] fix: analytic Chameleon potential consistent with its deflections (#682) The MGE potential used the circularised ellipticity convention while the convergence uses the major-axis elliptical radius, so grad(psi) was ~1/q times the deflections. Replace it with the exact closed-form potential of the two cored isothermals, built from the same psi_from / normalisation as the analytic deflections, and pin grad(psi) = alpha and lap(psi) = 2 kappa. Co-Authored-By: Claude Opus 5.5 --- autogalaxy/profiles/mass/stellar/chameleon.py | 63 +++++++++++--- .../profiles/mass/stellar/test_chameleon.py | 83 +++++++++++++++++++ 2 files changed, 136 insertions(+), 10 deletions(-) diff --git a/autogalaxy/profiles/mass/stellar/chameleon.py b/autogalaxy/profiles/mass/stellar/chameleon.py index 7065dea3a..e29947a4f 100644 --- a/autogalaxy/profiles/mass/stellar/chameleon.py +++ b/autogalaxy/profiles/mass/stellar/chameleon.py @@ -178,16 +178,59 @@ def convergence_func(self, grid_radius: float, xp=np) -> float: @aa.decorators.to_array @aa.decorators.transform def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): - from autogalaxy.profiles.mass.abstract.mge import MGEDecomposer - - radii_min = self.core_radius_0 / 10.0 - radii_max = self.core_radius_1 * 200.0 - sigmas = xp.exp(xp.linspace(xp.log(radii_min), xp.log(radii_max), 30)) - mge_decomp = MGEDecomposer(mass_profile=self) - return mge_decomp.potential_2d_via_mge_from( - grid=grid, xp=xp, sigma_log_list=sigmas, - ellipticity_convention="circularised", three_D=False, - ) + r""" + Calculate the lensing potential at a given set of arc-second gridded coordinates. + + The potential is analytic: each cored isothermal component has the closed form (Keeton 2001) + + .. math:: + + \psi_s = x \alpha_x + y \alpha_y + - b q s \ln \sqrt{(\Psi + s)^2 + (1 - q^2) x^2} + b q s \ln \left[ (1 + q) s \right] + + where :math:`\Psi` is given by ``psi_from`` and :math:`\alpha_x, \alpha_y` are the component's + deflections, with the same normalisation and core radii as ``deflections_2d_via_analytic_from``. + The Chameleon potential is the difference of the two components. + + Parameters + ---------- + grid + The grid of (y,x) arc-second coordinates the potential is computed on. + """ + axis_ratio = self.axis_ratio(xp) + b = 2.0 * self.mass_to_light_ratio * self.intensity / (1.0 + axis_ratio) + sqrt_1_minus_q2 = xp.sqrt(1.0 - axis_ratio**2.0) + + y = grid.array[:, 0] + x = grid.array[:, 1] + + def potential_cored_isothermal_from(core_radius): + psi = psi_from( + grid=grid, axis_ratio=axis_ratio, core_radius=core_radius, xp=xp + ) + + deflection_y = xp.arctanh( + sqrt_1_minus_q2 * y / (psi + axis_ratio**2.0 * core_radius) + ) + deflection_x = xp.arctan(sqrt_1_minus_q2 * x / (psi + core_radius)) + + return ( + b * axis_ratio / sqrt_1_minus_q2 * (y * deflection_y + x * deflection_x) + - b + * axis_ratio + * core_radius + * xp.log( + xp.sqrt((psi + core_radius) ** 2.0 + (1.0 - axis_ratio**2.0) * x**2.0) + ) + + b * axis_ratio * core_radius * xp.log((1.0 + axis_ratio) * core_radius) + ) + + core_radius_0 = 2.0 * self.core_radius_0 / (1.0 + axis_ratio) + core_radius_1 = 2.0 * self.core_radius_1 / (1.0 + axis_ratio) + + return potential_cored_isothermal_from( + core_radius_0 + ) - potential_cored_isothermal_from(core_radius_1) def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np): """Calculate the intensity of the Chamelon light profile on a grid of radial coordinates. diff --git a/test_autogalaxy/profiles/mass/stellar/test_chameleon.py b/test_autogalaxy/profiles/mass/stellar/test_chameleon.py index d0bec5ffd..0e34edf45 100644 --- a/test_autogalaxy/profiles/mass/stellar/test_chameleon.py +++ b/test_autogalaxy/profiles/mass/stellar/test_chameleon.py @@ -107,3 +107,86 @@ def test__convergence_2d_from__elliptical_vs_spherical(): assert elliptical.convergence_2d_from(grid=grid) == pytest.approx( spherical.convergence_2d_from(grid=grid).array, 1.0e-4 ) + + +def _potential_gradient_from(mp, points, h=1.0e-4): + def potential(shift): + return np.asarray( + mp.potential_2d_from(grid=ag.Grid2DIrregular(points + shift)) + ) + + gradient_y = (potential([h, 0.0]) - potential([-h, 0.0])) / (2.0 * h) + gradient_x = (potential([0.0, h]) - potential([0.0, -h])) / (2.0 * h) + + return np.stack([gradient_y, gradient_x], axis=-1) + + +_potential_points = np.array([[0.7, 0.3], [-0.5, 0.9], [1.2, -1.0], [2.5, 1.5]]) + +_potential_profiles = [ + ag.mp.Chameleon(ell_comps=(0.2, -0.1), intensity=1.0), + ag.mp.ChameleonSph(intensity=1.0), + ag.mp.Chameleon( + centre=(0.1, -0.2), + ell_comps=(0.2, -0.1), + intensity=1.0, + core_radius_0=0.1, + core_radius_1=0.5, + mass_to_light_ratio=2.0, + ), +] + + +@pytest.mark.parametrize("mp", _potential_profiles) +def test__potential_2d_from__gradient_matches_deflections(mp): + deflections = np.asarray( + mp.deflections_yx_2d_from(grid=ag.Grid2DIrregular(_potential_points)) + ) + + gradient = _potential_gradient_from(mp=mp, points=_potential_points) + + assert np.max(np.abs(gradient - deflections)) < 1.0e-4 * np.max( + np.abs(deflections) + ) + + +@pytest.mark.parametrize("mp", _potential_profiles) +def test__potential_2d_from__laplacian_is_twice_convergence(mp): + h = 1.0e-3 + + def potential(shift): + return np.asarray( + mp.potential_2d_from(grid=ag.Grid2DIrregular(_potential_points + shift)) + ) + + laplacian = ( + potential([h, 0.0]) + + potential([-h, 0.0]) + + potential([0.0, h]) + + potential([0.0, -h]) + - 4.0 * potential([0.0, 0.0]) + ) / h**2 + + convergence = np.asarray( + mp.convergence_2d_from(grid=ag.Grid2DIrregular(_potential_points)) + ) + + assert laplacian == pytest.approx(2.0 * convergence, rel=1.0e-3) + + +def test__potential_2d_from(): + mp = ag.mp.Chameleon( + centre=(-0.4, -0.2), + ell_comps=(-0.07142, -0.085116), + intensity=5.0, + core_radius_0=0.2, + core_radius_1=0.4, + mass_to_light_ratio=3.0, + ) + + potential = mp.potential_2d_from( + grid=ag.Grid2DIrregular([[-0.4, -0.2], [0.1625, 0.1625]]) + ) + + assert potential[0] == pytest.approx(0.0, abs=1.0e-8) + assert potential[1] == pytest.approx(1.68326, 1e-4) From 0b62e62c95df85c52e732268b73aa89756ba98d9 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 30 Oct 2026 10:06:26 +0000 Subject: [PATCH 2/2] fix: Chameleon jit-safe axis ratio, light xp pass-through, docstring matches code - Chameleon.axis_ratio (mass + light) capped q with a Python if, which raised TracerBoolConversionError under jax.jit; use xp.where (as Gaussian does). - lp.Chameleon.image_2d_from dropped xp when calling image_2d_via_radii_from. - Mass Chameleon / ChameleonSph docstrings claimed the second core is s0 + s1; the code uses core_radius_1 directly with 2s/(1+q) rescaling and a 1/(1+q) prefactor. Docstrings now state the implemented convergence. - Tests: jit over ell_comps matches numpy for deflections, potential and image. Co-Authored-By: Claude Opus 5.5 --- .../profiles/light/standard/chameleon.py | 4 +-- autogalaxy/profiles/mass/stellar/chameleon.py | 33 ++++++++++--------- .../profiles/light/standard/test_chameleon.py | 22 +++++++++++++ .../profiles/mass/stellar/test_chameleon.py | 32 ++++++++++++++++++ 4 files changed, 73 insertions(+), 18 deletions(-) diff --git a/autogalaxy/profiles/light/standard/chameleon.py b/autogalaxy/profiles/light/standard/chameleon.py index 5edd28ce8..a4eae44dc 100644 --- a/autogalaxy/profiles/light/standard/chameleon.py +++ b/autogalaxy/profiles/light/standard/chameleon.py @@ -50,7 +50,7 @@ def axis_ratio(self, xp=np) -> float: `axis_ratio=1.0`, thus we remove these solutions. """ axis_ratio = super().axis_ratio(xp=xp) - return axis_ratio if axis_ratio < 0.99999 else 0.99999 + return xp.where(axis_ratio < 0.99999, axis_ratio, 0.99999) def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray: """ @@ -117,7 +117,7 @@ def image_2d_from( The image of the Chameleon evaluated at every (y,x) coordinate on the transformed grid. """ return self.image_2d_via_radii_from( - self.elliptical_radii_grid_from(grid=grid, xp=xp, **kwargs) + self.elliptical_radii_grid_from(grid=grid, xp=xp, **kwargs), xp=xp ) diff --git a/autogalaxy/profiles/mass/stellar/chameleon.py b/autogalaxy/profiles/mass/stellar/chameleon.py index e29947a4f..a18ffbb08 100644 --- a/autogalaxy/profiles/mass/stellar/chameleon.py +++ b/autogalaxy/profiles/mass/stellar/chameleon.py @@ -14,20 +14,22 @@ class Chameleon(MassProfile, StellarProfile): Elliptical Chameleon stellar mass profile (Dutton et al. 2011). The Chameleon profile is the difference of two cored isothermal (pseudo-Jaffe) profiles - with core radii :math:`s_0` and :math:`s_0 + s_1`, providing a flexible approximation + with core radii :math:`s_0` and :math:`s_1`, providing a flexible approximation to a variety of stellar light profiles: .. math:: - \kappa(\xi) = \Upsilon \, I \left( - \frac{1}{\sqrt{\xi^2 + s_0^2}} - - \frac{1}{\sqrt{\xi^2 + (s_0 + s_1)^2}} + \kappa(\xi) = \frac{\Upsilon \, I}{1 + q} \left( + \frac{1}{\sqrt{\xi^2 + 4 s_0^2 / (1 + q)^2}} + - \frac{1}{\sqrt{\xi^2 + 4 s_1^2 / (1 + q)^2}} \right) - where :math:`\xi^2 = x^2 + (y/q)^2` is the elliptical radius, :math:`\Upsilon` is - the mass-to-light ratio (``mass_to_light_ratio``), :math:`I` is the intensity - normalisation (``intensity``), :math:`s_0` = ``core_radius_0``, and - :math:`s_1` = ``core_radius_1``. + where :math:`\xi^2 = x^2 + (y/q)^2` is the elliptical radius, :math:`q` is the axis + ratio, :math:`\Upsilon` is the mass-to-light ratio (``mass_to_light_ratio``), + :math:`I` is the intensity normalisation (``intensity``), :math:`s_0` = ``core_radius_0`` + and :math:`s_1` = ``core_radius_1``. The convergence is positive when + :math:`s_1 > s_0`. For :math:`q = 1` the rescaled cores reduce to :math:`s_0` and + :math:`s_1` and the prefactor to :math:`\Upsilon I / 2`. Deflection angles are computed analytically via the cored isothermal deflection formula (Eq. 15–16 of Dutton et al. 2011). @@ -58,9 +60,8 @@ def __init__( core_radius_0 Core radius :math:`s_0` of the first cored isothermal component (arcseconds). core_radius_1 - Additional core size increment :math:`s_1` such that the second component has - core radius :math:`s_0 + s_1` (arcseconds). Using an increment avoids - negative parameter values. + Core radius :math:`s_1` of the second cored isothermal component (arcseconds), + which is subtracted from the first. mass_to_light_ratio The mass-to-light ratio :math:`\Upsilon` in solar units. """ @@ -269,7 +270,7 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np): def axis_ratio(self, xp=np): axis_ratio = super().axis_ratio(xp=xp) - return axis_ratio if axis_ratio < 0.99999 else 0.99999 + return xp.where(axis_ratio < 0.99999, axis_ratio, 0.99999) class ChameleonSph(Chameleon): @@ -281,9 +282,9 @@ class ChameleonSph(Chameleon): .. math:: - \kappa(r) = \Upsilon \, I \left( + \kappa(r) = \frac{\Upsilon \, I}{2} \left( \frac{1}{\sqrt{r^2 + s_0^2}} - - \frac{1}{\sqrt{r^2 + (s_0 + s_1)^2}} + - \frac{1}{\sqrt{r^2 + s_1^2}} \right) References @@ -309,8 +310,8 @@ def __init__( core_radius_0 Core radius :math:`s_0` of the first cored isothermal component (arcseconds). core_radius_1 - Additional core size increment :math:`s_1` such that the second component has - core radius :math:`s_0 + s_1` (arcseconds). + Core radius :math:`s_1` of the second cored isothermal component (arcseconds), + which is subtracted from the first. mass_to_light_ratio The mass-to-light ratio :math:`\Upsilon` in solar units. """ diff --git a/test_autogalaxy/profiles/light/standard/test_chameleon.py b/test_autogalaxy/profiles/light/standard/test_chameleon.py index bd1be5b79..0191bb942 100644 --- a/test_autogalaxy/profiles/light/standard/test_chameleon.py +++ b/test_autogalaxy/profiles/light/standard/test_chameleon.py @@ -1,4 +1,5 @@ from __future__ import division, print_function +import numpy as np import pytest import autogalaxy as ag @@ -59,3 +60,24 @@ def test__image_2d_from__spherical_profile__matches_elliptical_with_zero_ellipti image_spherical = spherical.image_2d_from(grid=grid) assert image_elliptical.array == pytest.approx(image_spherical.array, 1.0e-4) + + +def test__image_2d_from__jit_over_ell_comps_matches_numpy(): + jax = pytest.importorskip("jax") + import jax.numpy as jnp + + with jax.enable_x64(True): + grid = ag.Grid2DIrregular([[0.7, 0.3], [-0.5, 0.9], [1.2, -1.0]]) + + @jax.jit + def f(e): + lp = ag.lp.Chameleon(ell_comps=(e[0], e[1]), intensity=1.0) + return lp.image_2d_from(grid=grid, xp=jnp).array + + image = f(jnp.asarray([0.2, -0.1])) + + lp = ag.lp.Chameleon(ell_comps=(0.2, -0.1), intensity=1.0) + + assert np.asarray(image) == pytest.approx( + np.asarray(lp.image_2d_from(grid=grid)), abs=1.0e-12 + ) diff --git a/test_autogalaxy/profiles/mass/stellar/test_chameleon.py b/test_autogalaxy/profiles/mass/stellar/test_chameleon.py index 0e34edf45..00ab0c6ec 100644 --- a/test_autogalaxy/profiles/mass/stellar/test_chameleon.py +++ b/test_autogalaxy/profiles/mass/stellar/test_chameleon.py @@ -190,3 +190,35 @@ def test__potential_2d_from(): assert potential[0] == pytest.approx(0.0, abs=1.0e-8) assert potential[1] == pytest.approx(1.68326, 1e-4) + + +def test__jit_over_ell_comps_matches_numpy(): + jax = pytest.importorskip("jax") + import jax.numpy as jnp + + with jax.enable_x64(True): + grid = ag.Grid2DIrregular(_potential_points) + + def mp_from(ell_comps): + return ag.mp.Chameleon( + centre=(0.1, -0.2), ell_comps=ell_comps, intensity=1.0 + ) + + @jax.jit + def f(e): + mp = mp_from(ell_comps=(e[0], e[1])) + return ( + mp.deflections_yx_2d_from(grid=grid, xp=jnp).array, + mp.potential_2d_from(grid=grid, xp=jnp).array, + ) + + deflections, potential = f(jnp.asarray([0.2, -0.1])) + + mp = mp_from(ell_comps=(0.2, -0.1)) + + assert np.asarray(deflections) == pytest.approx( + np.asarray(mp.deflections_yx_2d_from(grid=grid)), abs=1.0e-12 + ) + assert np.asarray(potential) == pytest.approx( + np.asarray(mp.potential_2d_from(grid=grid)), abs=1.0e-12 + )