fix: use analytic potential for truncated NFW

autogalaxy/profiles/mass/dark/nfw_truncated.py

@@ -31,6 +31,108 @@ def coord_func_m_from(grid_radius, tau, xp=np):
     )
 
 
+def potential_func_sph_from(grid_radius, tau, xp=np):
+    r"""Return the dimensionless analytic tNFW lensing potential.
+
+    This is equation (18) of Baltz, Marshall & Oguri (2009) for their
+    :math:`n=1` smoothly truncated NFW profile.  The returned function is the
+    paper's dimensionless radial term; ``NFWTruncatedSph.potential_2d_from``
+    supplies the PyAutoGalaxy normalization :math:`2\kappa_s r_s^2`.
+
+    The potential is defined only up to an additive constant.  We retain the
+    paper's convention, in which the central potential is zero, but use its
+    small-radius series through :math:`x=10^{-1}`.  Direct evaluation there
+    subtracts large, nearly equal terms and loses precision in JAX float32.
+    """
+    grid_radius = xp.real(grid_radius)
+    series_radius = xp.maximum(grid_radius, 1.0e-12)
+    use_small_radius_series = grid_radius <= 1.0e-1
+
+    # Keep every branch of the closed form finite when JAX traces ``where``.
+    # Values below the switch are replaced by the series before returning.
+    grid_radius = xp.where(use_small_radius_series, 1.0e-1, grid_radius)
+
+    u = xp.square(grid_radius)
+    tau_squared = xp.square(tau)
+    root = xp.sqrt(tau_squared + u)
+
+    radius_lt = xp.sqrt(xp.where(grid_radius < 1.0, 1.0 - u, 0.25))
+    radius_gt = xp.sqrt(xp.where(grid_radius > 1.0, u - 1.0, 0.25))
+    f_lt = xp.arctanh(radius_lt) / radius_lt
+    f_gt = xp.arctan(radius_gt) / radius_gt
+    f_r = xp.where(grid_radius < 1.0, f_lt, xp.where(grid_radius > 1.0, f_gt, 1.0))
+
+    l_r = coord_func_k_from(grid_radius=grid_radius, tau=tau, xp=xp)
+
+    inverse_radius = 1.0 / grid_radius
+    inverse_radius_lt = xp.where(grid_radius < 1.0, inverse_radius, 1.5)
+    inverse_radius_gt = xp.where(grid_radius > 1.0, inverse_radius, 0.5)
+    cos_lt = -xp.square(xp.arccosh(inverse_radius_lt))
+    cos_gt = xp.square(xp.arccos(inverse_radius_gt))
+    cos_term = xp.where(
+        grid_radius < 1.0,
+        cos_lt,
+        xp.where(grid_radius > 1.0, cos_gt, 0.0),
+    )
+
+    potential = (
+        2.0 * tau_squared * xp.pi * (tau - root + tau * xp.log(tau + root))
+        + 2.0 * (tau_squared - 1.0) * tau * root * l_r
+        + tau_squared * (tau_squared - 1.0) * xp.square(l_r)
+        + 4.0 * tau_squared * (u - 1.0) * f_r
+        + tau_squared * (tau_squared - 1.0) * cos_term
+        + tau_squared
+        * ((tau_squared - 1.0) * xp.log(tau) - tau_squared - 1.0)
+        * xp.log(u)
+        - tau_squared
+        * (
+            (tau_squared - 1.0) * xp.log(tau) * xp.log(4.0 * tau)
+            + 2.0 * xp.log(tau / 2.0)
+            - 2.0 * tau * (tau - xp.pi) * xp.log(2.0 * tau)
+        )
+    )
+
+    potential = potential / xp.square(tau_squared + 1.0)
+
+    small_radius_log = xp.log(2.0 / series_radius)
+    small_radius_coefficient_2 = (
+        small_radius_log * xp.square(tau_squared + 1.0)
+        - tau_squared * xp.log(tau)
+        + tau_squared
+        - xp.pi * tau
+        + xp.log(tau)
+        + 1.0
+    )
+    small_radius_coefficient_2 /= 2.0 * xp.square(tau_squared + 1.0)
+    small_radius_coefficient_4 = (
+        small_radius_log * (3.0 * tau**6.0 + 5.0 * tau**4.0 + tau_squared - 1.0)
+        - tau**6.0
+        - tau**4.0
+        + tau_squared * xp.log(tau)
+        + xp.pi * tau
+        - xp.log(tau)
+    )
+    small_radius_coefficient_4 /= 16.0 * tau_squared * xp.square(tau_squared + 1.0)
+    small_radius_coefficient_6 = (
+        small_radius_log * (20.0 * tau**8.0 + 28.0 * tau**6.0 - 4.0 * tau_squared + 4.0)
+        - 9.0 * tau**8.0
+        - 11.0 * tau**6.0
+        - 4.0 * tau_squared * xp.log(tau)
+        + tau_squared
+        - 4.0 * xp.pi * tau
+        + 4.0 * xp.log(tau)
+        - 1.0
+    )
+    small_radius_coefficient_6 /= 192.0 * tau**4.0 * xp.square(tau_squared + 1.0)
+    potential_small_radius = (
+        xp.square(series_radius) * small_radius_coefficient_2
+        + series_radius**4.0 * small_radius_coefficient_4
+        + series_radius**6.0 * small_radius_coefficient_6
+    )
+
+    return xp.where(use_small_radius_series, potential_small_radius, potential)
+
+
 class NFWTruncatedSph(AbstractgNFW):
     r"""
     Spherical truncated NFW (tNFW) dark matter halo profile (Baltz, Marshall & Oguri 2009).
@@ -59,7 +161,7 @@ class NFWTruncatedSph(AbstractgNFW):
 
     References
     ----------
-    - Baltz, Marshall & Oguri 2009, JCAP, 2009, 015  (arXiv:0705.0735)
+    - Baltz, Marshall & Oguri 2009, JCAP, 2009, 015  (arXiv:0705.0682)
     - Navarro, Frenk & White 1997, ApJ, 490, 493
     """
 
@@ -135,15 +237,21 @@ 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.scale_radius / 1000.0
-        radii_max = self.truncation_radius * 5.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="major", three_D=True,
+        """Calculate the analytic lensing potential of the spherical tNFW profile."""
+        eta = xp.multiply(
+            1.0 / self.scale_radius,
+            self.radial_grid_from(grid=grid, xp=xp, **kwargs).array,
+        )
+
+        return (
+            2.0
+            * self.kappa_s
+            * self.scale_radius**2.0
+            * potential_func_sph_from(
+                grid_radius=eta,
+                tau=self.tau,
+                xp=xp,
+            )
         )
 
     def coord_func_k(self, grid_radius, xp=np):
@@ -197,10 +305,7 @@ def _delta_c_from_concentration(concentration: float) -> float:
             / 3.0
             * (
                 concentration**3
-                / (
-                    np.log(1.0 + concentration)
-                    - concentration / (1.0 + concentration)
-                )
+                / (np.log(1.0 + concentration) - concentration / (1.0 + concentration))
             )
         )
 
@@ -374,12 +479,7 @@ def m200_concentration_from(
         delta_c = rho_s / critical_density
 
         def equation(c):
-            return (
-                200.0
-                / 3.0
-                * (c**3 / (np.log(1.0 + c) - c / (1.0 + c)))
-                - delta_c
-            )
+            return 200.0 / 3.0 * (c**3 / (np.log(1.0 + c) - c / (1.0 + c))) - delta_c
 
         concentration = float(fsolve(equation, 10.0)[0])
         r200_kpc = concentration * rs_kpc

test_autogalaxy/profiles/mass/dark/test_nfw_truncated.py

@@ -114,6 +114,116 @@ def test__convergence_2d_from__scale_radius_5():
     assert convergence == pytest.approx(1.51047026, abs=1.0e-4)
 
 
+def test__potential_2d_from__matches_independent_tnfw_implementation():
+    """Potential differences cross-checked against lenstronomy's analytic TNFW."""
+    mp = ag.mp.NFWTruncatedSph(
+        kappa_s=0.13,
+        scale_radius=1.7,
+        truncation_radius=3.4,
+    )
+
+    potential = mp.potential_2d_from(
+        grid=ag.Grid2DIrregular([[0.0, 0.085], [0.0, 0.51], [0.0, 5.1]])
+    )
+    potential -= potential[0]
+
+    assert potential == pytest.approx(
+        [0.0, 0.05804607, 0.88448954],
+        abs=1.0e-8,
+    )
+
+
+def test__potential_2d_from__centre_and_scale_radius_are_finite():
+    mp = ag.mp.NFWTruncatedSph(
+        kappa_s=0.13,
+        scale_radius=1.7,
+        truncation_radius=3.4,
+    )
+
+    potential = mp.potential_2d_from(grid=ag.Grid2DIrregular([[0.0, 0.0], [0.0, 1.7]]))
+
+    assert np.all(np.isfinite(potential))
+    assert potential[0] == pytest.approx(0.0, abs=1.0e-20)
+
+
+@pytest.mark.parametrize(
+    "tau,radius_over_scale",
+    [
+        (0.5, 0.3),
+        (0.5, 2.0),
+        (1.0, 0.3),
+        (1.0, 2.0),
+        (2.0, 0.3),
+        (2.0, 2.0),
+        (10.0, 0.3),
+        (10.0, 2.0),
+    ],
+)
+def test__potential_2d_from__gradient_matches_analytic_deflections(
+    tau, radius_over_scale
+):
+    scale_radius = 1.7
+    radius = radius_over_scale * scale_radius
+    step = 1.0e-4 * scale_radius
+
+    mp = ag.mp.NFWTruncatedSph(
+        kappa_s=0.13,
+        scale_radius=scale_radius,
+        truncation_radius=tau * scale_radius,
+    )
+
+    radial_grid = ag.Grid2DIrregular([[0.0, radius - step], [0.0, radius + step]])
+    potential = mp.potential_2d_from(grid=radial_grid)
+    potential_gradient = (potential[1] - potential[0]) / (2.0 * step)
+
+    deflection = mp.deflections_yx_2d_from(grid=ag.Grid2DIrregular([[0.0, radius]]))[
+        0, 1
+    ]
+
+    assert potential_gradient == pytest.approx(deflection, rel=1.0e-6)
+
+
+@pytest.mark.parametrize(
+    "tau,radius_over_scale",
+    [
+        (0.5, 0.3),
+        (0.5, 2.0),
+        (1.0, 0.3),
+        (1.0, 2.0),
+        (2.0, 0.3),
+        (2.0, 2.0),
+        (10.0, 0.3),
+        (10.0, 2.0),
+    ],
+)
+def test__potential_2d_from__laplacian_matches_analytic_convergence(
+    tau, radius_over_scale
+):
+    scale_radius = 1.7
+    radius = radius_over_scale * scale_radius
+    step = 1.0e-3 * scale_radius
+
+    mp = ag.mp.NFWTruncatedSph(
+        kappa_s=0.13,
+        scale_radius=scale_radius,
+        truncation_radius=tau * scale_radius,
+    )
+
+    radial_grid = ag.Grid2DIrregular(
+        [[0.0, radius - step], [0.0, radius], [0.0, radius + step]]
+    )
+    potential = mp.potential_2d_from(grid=radial_grid)
+    potential_gradient = (potential[2] - potential[0]) / (2.0 * step)
+    potential_second_derivative = (
+        potential[2] - 2.0 * potential[1] + potential[0]
+    ) / step**2.0
+    potential_laplacian = potential_second_derivative + potential_gradient / radius
+
+    convergence = mp.convergence_2d_from(grid=ag.Grid2DIrregular([[0.0, radius]]))[0]
+
+    assert potential_laplacian == pytest.approx(2.0 * convergence, rel=1.0e-5)
+
+
 def test__mass_at_truncation_radius():
     mp = ag.mp.NFWTruncatedSph(
         centre=(0.0, 0.0), kappa_s=1.0, scale_radius=1.0, truncation_radius=1.0
@@ -158,6 +268,19 @@ def test__compare_nfw_and_truncated_nfw_with_large_truncation_radius():
 
     assert truncated_nfw_deflections == pytest.approx(nfw_deflections.array, abs=1.0e-4)
 
+    truncated_nfw_potential = truncated_nfw.potential_2d_from(
+        grid=ag.Grid2DIrregular([[2.0, 2.0], [3.0, 1.0], [-1.0, -9.0]])
+    )
+    nfw_potential = nfw.potential_2d_from(
+        grid=ag.Grid2DIrregular([[2.0, 2.0], [3.0, 1.0], [-1.0, -9.0]])
+    )
+
+    # The lensing potential is defined only up to an additive constant.
+    truncated_nfw_potential -= truncated_nfw_potential[0]
+    nfw_potential -= nfw_potential[0]
+
+    assert truncated_nfw_potential == pytest.approx(nfw_potential, abs=1.0e-4)
+
 
 # ---------------------------------------------------------------------------
 # Helpers: reference implementation of the los_pipes unit-conversion formulas
@@ -166,6 +289,7 @@ def test__compare_nfw_and_truncated_nfw_with_large_truncation_radius():
 # regression tests below.
 # ---------------------------------------------------------------------------
 
+
 def _los_pipes_reference_delta_c(concentration):
     """NFW characteristic overdensity as computed by los_pipes."""
     c = concentration
@@ -197,8 +321,10 @@ def _los_pipes_reference_convert_to_lens_unit(
 
     critical_density = cosmo.critical_density(z_halo)
     kpc_per_arcsec = cosmo.kpc_per_arcsec_from(z_halo)
-    sigma_crit = cosmo.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from(
-        z_halo, z_source
+    sigma_crit = (
+        cosmo.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from(
+            z_halo, z_source
+        )
     )
 
     r200_kpc = (m200 / (200.0 * critical_density * (4.0 * np.pi / 3.0))) ** (1.0 / 3.0)
@@ -483,11 +609,15 @@ def test__mass_ratio_from_concentration_and_truncation_factor__matches_los_pipes
 def test__mass_ratio_from_concentration_and_truncation_factor__larger_for_smaller_factor():
     """A smaller truncation factor (larger truncation radius) gives a larger mass ratio."""
     c = 10.0
-    ratio_100 = ag.mp.NFWTruncatedSph.mass_ratio_from_concentration_and_truncation_factor(
-        c, 100.0
+    ratio_100 = (
+        ag.mp.NFWTruncatedSph.mass_ratio_from_concentration_and_truncation_factor(
+            c, 100.0
+        )
     )
-    ratio_50 = ag.mp.NFWTruncatedSph.mass_ratio_from_concentration_and_truncation_factor(
-        c, 50.0
+    ratio_50 = (
+        ag.mp.NFWTruncatedSph.mass_ratio_from_concentration_and_truncation_factor(
+            c, 50.0
+        )
     )
     assert ratio_50 > ratio_100
 
@@ -496,8 +626,10 @@ def test__mass_ratio_from_concentration_and_truncation_factor__various_concentra
     """Spot-check mass ratios at several concentrations against los_pipes reference."""
     for c in [5.0, 10.0, 20.0]:
         expected = _los_pipes_reference_mass_ratio(c, 100.0)
-        result = ag.mp.NFWTruncatedSph.mass_ratio_from_concentration_and_truncation_factor(
-            c, 100.0
+        result = (
+            ag.mp.NFWTruncatedSph.mass_ratio_from_concentration_and_truncation_factor(
+                c, 100.0
+            )
         )
         assert result == pytest.approx(expected, rel=1.0e-6), f"failed for c={c}"