Merge pull request #470 from PyAutoLabs/feature/dark-matter-potentials

Lensing potential for elliptical/spherical dark-matter profiles (NFW/gNFW) + NFWSph fix

Author: Jammy2211

autogalaxy/profiles/mass/abstract/abstract.py

@@ -173,7 +173,7 @@ def convergence_func(self, grid_radius: float, xp=np) -> float:
         """
         raise NotImplementedError
 
-    def potential_2d_from(self, grid):
+    def potential_2d_from(self, grid, xp=np, **kwargs):
         """
         Returns the 2D lensing potential of the mass profile from a 2D grid of Cartesian (y,x) coordinates.
 

autogalaxy/profiles/mass/dark/gnfw.py

@@ -60,6 +60,33 @@ def deflections_2d_via_mge_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs)
         )
         return deflections_via_mge
 
+    @aa.over_sample
+    @aa.decorators.to_array
+    @aa.decorators.transform
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        """
+        Returns the 2D lensing potential via the same MGE decomposition used for the
+        deflection angles, so that ``grad(psi)`` is self-consistent with ``deflections_yx_2d_from``.
+
+        This implementation is inherited by every elliptical and spherical gNFW/NFW variant
+        (``NFW``, ``gNFWSph``, and the MCR / Virial-mass subclasses), none of which have a
+        closed-form elliptical potential. The MGE decomposer expands the 3D density into a
+        Gaussian sum and integrates each Gaussian's analytic potential.
+        """
+        radii_min = self.scale_radius / 20000.0
+        radii_max = self.scale_radius * 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="major",
+            three_D=True,
+        )
+
     def convergence_func(self, grid_radius: float, xp=np) -> float:
 
         from scipy.integrate import quad

autogalaxy/profiles/mass/dark/nfw.py

@@ -413,11 +413,19 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
             grid=grid, **kwargs
         ) + 0j
         return np.real(
-            2.0 * self.scale_radius * self.kappa_s * self.potential_func_sph(eta, xp=xp)
+            2.0
+            * self.scale_radius**2
+            * self.kappa_s
+            * self.potential_func_sph(eta, xp=xp)
         )
 
     @staticmethod
     def potential_func_sph(eta, xp=np):
+        # Spherical NFW lensing potential psi(x) = 2 kappa_s r_s^2 g(x), with
+        # g(x) = ln^2(x/2) - arctanh^2(sqrt(1 - x^2)), obtained by integrating
+        # the deflection field radially. The previous prefactor used a single
+        # factor of r_s instead of r_s^2, so grad(psi) came out a factor 1/r_s
+        # too small relative to deflections_yx_2d_from.
         return ((xp.log(eta.array / 2.0)) ** 2) - (
             xp.arctanh(xp.sqrt(1 - eta.array**2))
         ) ** 2

test_autogalaxy/profiles/mass/dark/test_gnfw.py

@@ -67,8 +67,8 @@ def test__potential_2d_from__elliptical_vs_spherical():
         centre=(0.1, 0.2), kappa_s=2.0, inner_slope=1.5, scale_radius=3.0
     )
 
-    assert elliptical.convergence_2d_from(grid) == pytest.approx(
-        spherical.convergence_2d_from(grid).array, 1e-4
+    assert elliptical.potential_2d_from(grid) == pytest.approx(
+        spherical.potential_2d_from(grid).array, 1e-4
     )
 
 

test_autogalaxy/profiles/mass/dark/test_nfw.py

@@ -287,46 +287,28 @@ def test__convergence_2d_from__nfw_ell():
     assert convergence == pytest.approx(1.388511, 1e-3)
 
 
-# def test__potential_2d_from():
-#     nfw = ag.mp.NFWSph(centre=(0.3, 0.2), kappa_s=2.5, scale_radius=4.0)
-#
-#     potential = nfw.potential_2d_from(grid=ag.Grid2DIrregular([[0.1875, 0.1625]]))
-#
-#     assert potential == pytest.approx(0.03702, 1e-3)
-#
-#     nfw = ag.mp.NFWSph(centre=(0.3, 0.2), kappa_s=2.5, scale_radius=4.0)
-#
-#     potential = nfw.potential_2d_from(grid=ag.Grid2DIrregular([[0.1875, 0.1625]]))
-#
-#     assert potential == pytest.approx(0.03702, 1e-3)
-#
-#     nfw = ag.mp.NFW(
-#         centre=(0.3, 0.2),
-#         ell_comps=(0.03669, 0.172614),
-#         kappa_s=2.5,
-#         scale_radius=4.0,
-#     )
-#
-#     potential = nfw.potential_2d_from(grid=ag.Grid2DIrregular([[0.1625, 0.1625]]))
-#
-#     assert potential == pytest.approx(0.05380, 1e-3)
-#
-#     nfw_spherical = ag.mp.NFWSph(centre=(0.3, 0.2), kappa_s=2.5, scale_radius=4.0)
-#     nfw_elliptical = ag.mp.NFW(
-#         centre=(0.3, 0.2),
-#         ell_comps=(0.0, 0.0),
-#         kappa_s=2.5,
-#         scale_radius=4.0,
-#     )
-#
-#     potential_spherical = nfw_spherical.potential_2d_from(
-#         grid=ag.Grid2DIrregular([[0.1875, 0.1625]])
-#     )
-#     potential_elliptical = nfw_elliptical.potential_2d_from(
-#         grid=ag.Grid2DIrregular([[0.1875, 0.1625]])
-#     )
-#
-#     assert potential_spherical == pytest.approx(potential_elliptical, 1e-3)
+def test__potential_2d_from__nfw_sph():
+    nfw = ag.mp.NFWSph(centre=(0.3, 0.2), kappa_s=2.5, scale_radius=4.0)
+
+    potential = nfw.potential_2d_from(grid=ag.Grid2DIrregular([[0.1875, 0.1625]]))
+
+    assert potential == pytest.approx(0.148108, 1e-3)
+
+
+def test__potential_2d_from__elliptical_reduces_to_spherical():
+    # At zero ellipticity the elliptical NFW (MGE potential) must match the
+    # spherical NFW (analytic potential) — a cross-check of two independent
+    # implementations.
+    spherical = ag.mp.NFWSph(centre=(0.3, 0.2), kappa_s=2.5, scale_radius=4.0)
+    elliptical = ag.mp.NFW(
+        centre=(0.3, 0.2), ell_comps=(0.0, 0.0), kappa_s=2.5, scale_radius=4.0
+    )
+
+    grid = ag.Grid2DIrregular([[0.1875, 0.1625]])
+
+    assert elliptical.potential_2d_from(grid=grid) == pytest.approx(
+        spherical.potential_2d_from(grid=grid).array, 1e-3
+    )
 
 
 def test__shear_yx_2d_from__nfw_sph_config_1():