fix: analytic Chameleon potential + jit-safe Chameleon profiles

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
         )
 
 

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.
         """
@@ -178,16 +179,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.
@@ -226,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):
@@ -238,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
@@ -266,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.
         """

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
+        )

test_autogalaxy/profiles/mass/stellar/test_chameleon.py

@@ -107,3 +107,118 @@ 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)
+
+
+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
+        )