From 0e09f146cbadfc005117603b58b9efe840dd94a9 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 9 Aug 2026 13:55:46 -0400 Subject: [PATCH] fix: use analytic truncated NFW potential --- .../profiles/mass/dark/nfw_truncated.py | 140 ++++++++++++++--- .../profiles/mass/dark/test_nfw_truncated.py | 148 +++++++++++++++++- 2 files changed, 260 insertions(+), 28 deletions(-) diff --git a/autogalaxy/profiles/mass/dark/nfw_truncated.py b/autogalaxy/profiles/mass/dark/nfw_truncated.py index 5c834fd4..dee7bcef 100644 --- a/autogalaxy/profiles/mass/dark/nfw_truncated.py +++ b/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 diff --git a/test_autogalaxy/profiles/mass/dark/test_nfw_truncated.py b/test_autogalaxy/profiles/mass/dark/test_nfw_truncated.py index a7ed878d..e57fa888 100644 --- a/test_autogalaxy/profiles/mass/dark/test_nfw_truncated.py +++ b/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}"