Feature/jax e nfw

autogalaxy/profiles/mass/dark/nfw.py

@@ -42,8 +42,63 @@ def __init__(
         )
         super(MassProfileCSE, self).__init__()
 
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike):
-        return self.deflections_2d_via_cse_from(grid=grid)
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        return self.deflections_2d_via_analytic_from(grid=grid, xp=xp, **kwargs)
+
+    @aa.grid_dec.to_vector_yx
+    @aa.grid_dec.transform
+    def deflections_2d_via_analytic_from(
+            self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ):
+        """
+        Analytic calculation deflection angles from Heyrovský & Karamazov 2024 via Eq. 30 & 31
+
+        Parameters
+        ----------
+        grid
+            The grid of (y,x) arc-second coordinates the deflection angles are computed on.
+        """
+
+        # Convert e definitions:
+        # from q = (1-e)/(1+e) to q = sqrt(1-e**2)
+
+        e_autolens = xp.sqrt(self.ell_comps[1] ** 2 + self.ell_comps[0] ** 2)
+        e_hk24 = 2 * xp.sqrt(e_autolens) / (1 + e_autolens)
+
+        # Define dimensionless length coords
+
+        x1 = grid.array[:, 1] / self.scale_radius
+        x2 = grid.array[:, 0] / self.scale_radius
+
+        r2 = x1 ** 2 + x2 ** 2
+
+        # Avoid nans
+
+        mask = r2 > 1e-24
+
+        prefactor = xp.where(mask, 4 * self.kappa_s * xp.sqrt(1 - e_hk24 ** 2) / (
+                ((x1 - e_hk24) ** 2 + x2 ** 2) * ((x1 + e_hk24) ** 2 + x2 ** 2)
+        ), 0.0)
+
+        f1 = xp.where(mask, nfw_hk24_util.small_f_1(x1, x2, e_hk24, xp=xp), 0.0)
+        f2 = xp.where(mask, nfw_hk24_util.small_f_2(x1, x2, e_hk24, xp=xp), 0.0)
+        f3 = xp.where(mask, nfw_hk24_util.small_f_3(x1, x2, e_hk24, xp=xp), 0.0)
+
+        deflection_x = (x1 * ((x1**2 - e_hk24**2) * (1 - e_hk24**2) + x2**2 * (1 + e_hk24**2)) * f1
+                        + x1 * (x1**2 + x2**2 - e_hk24**2) * f2
+                        - x2 * (x1**2 + x2**2 + e_hk24**2) * f3)
+        deflection_x *= prefactor
+
+        deflection_y = (x2 * (x1**2 * (1 - 2 * e_hk24**2) + x2**2 + e_hk24**2) * f1
+                        + x2 * (x1**2 + x2**2 + e_hk24**2) * f2
+                        + x1 * (x1**2 + x2**2 - e_hk24**2) * f3)
+        deflection_y *= prefactor
+
+        return self.rotated_grid_from_reference_frame_from(
+            xp.multiply(self.scale_radius, xp.vstack((deflection_y, deflection_x)).T),
+            xp=xp,
+            **kwargs,
+        )
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
@@ -146,7 +201,7 @@ def convergence_func(self, grid_radius: float) -> float:
         return np.real(
             2.0
             * self.kappa_s
-            * np.array(self.coord_func_g(grid_radius=grid_radius, xp=xp))
+            * np.array(self.coord_func_g(grid_radius=grid_radius))
         )
 
     @aa.over_sample
@@ -271,26 +326,26 @@ def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         # Convert e definitions:
         # from q = (1-e)/(1+e) to q = sqrt(1-e**2)
 
-        e_autolens = np.sqrt(self.ell_comps[1] ** 2 + self.ell_comps[0] ** 2)
-        e_hk24 = 2 * np.sqrt(e_autolens) / np.sqrt(1 + 2 * e_autolens + e_autolens**2)
+        e_autolens = xp.sqrt(self.ell_comps[1] ** 2 + self.ell_comps[0] ** 2)
+        e_hk24 = 2 * xp.sqrt(e_autolens) / (1 + e_autolens)
 
         # Define dimensionless length coords
 
         x1 = grid.array[:, 1] / self.scale_radius
         x2 = grid.array[:, 0] / self.scale_radius
 
         # Avoid nans due to x=0
-        x1 = np.where(np.abs(x1) < 1e-6, 1e-6, x1)
-        x2 = np.where(np.abs(x2) < 1e-6, 1e-6, x2)
+        x1 = xp.where(xp.abs(x1) < 1e-6, 1e-6, x1)
+        x2 = xp.where(xp.abs(x2) < 1e-6, 1e-6, x2)
 
         # Calculate shear from nfw_HK24.py
 
-        g1, g2 = nfw_hk24_util.g1_g2_from(x1=x1, x2=x2, e=e_hk24, k_s=self.kappa_s)
+        g1, g2 = nfw_hk24_util.g1_g2_from(x1=x1, x2=x2, e=e_hk24, k_s=self.kappa_s, xp=xp)
 
         # Rotation for shear
 
         shear_field = self.rotated_grid_from_reference_frame_from(
-            grid=np.vstack((g2, g1)).T, angle=self.angle(xp) * 2
+            grid=xp.vstack((g2, g1)).T, angle=self.angle(xp) * 2, xp=xp
         )
 
         return aa.VectorYX2DIrregular(values=shear_field, grid=grid)
@@ -315,8 +370,8 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         # Convert e definitions:
         # from q = (1-e)/(1+e) to q = sqrt(1-e**2)
 
-        e_autolens = np.sqrt(self.ell_comps[1] ** 2 + self.ell_comps[0] ** 2)
-        e_hk24 = 2 * np.sqrt(e_autolens) / np.sqrt(1 + 2 * e_autolens + e_autolens**2)
+        e_autolens = xp.sqrt(self.ell_comps[1] ** 2 + self.ell_comps[0] ** 2)
+        e_hk24 = 2 * xp.sqrt(e_autolens) / (1 + e_autolens)
 
         # Define dimensionless length coords
 
@@ -325,13 +380,13 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
 
         # Avoid nans due to x=0
 
-        x1 = np.where(np.abs(x1) < 1e-6, 1e-6, x1)
-        x2 = np.where(np.abs(x2) < 1e-6, 1e-6, x2)
+        x1 = xp.where(xp.abs(x1) < 1e-6, 1e-6, x1)
+        x2 = xp.where(xp.abs(x2) < 1e-6, 1e-6, x2)
 
         # Calculate convergence from nfw_HK24.py
-        a = nfw_hk24_util.semi_major_axis_from(x1, x2, e_hk24)
+        a = nfw_hk24_util.semi_major_axis_from(x1, x2, e_hk24, xp=xp)
 
-        return nfw_hk24_util.kappa_from(k_s=self.kappa_s, a=a)
+        return nfw_hk24_util.kappa_from(k_s=self.kappa_s, a=a, xp=xp)
 
 
 class NFWSph(NFW):

autogalaxy/profiles/mass/dark/nfw_hk24_util.py

@@ -7,7 +7,7 @@
 import numpy as np
 
 
-def semi_major_axis_from(x1: np.ndarray, x2: np.ndarray, e: np.ndarray) -> np.ndarray:
+def semi_major_axis_from(x1, x2, e, xp=np):
     """
     Returns the semi-major axis of the ellipse at a given point.
 
@@ -20,10 +20,10 @@ def semi_major_axis_from(x1: np.ndarray, x2: np.ndarray, e: np.ndarray) -> np.nd
     e
         Eccentricity.
     """
-    return np.sqrt(x1**2 + x2**2 / (1 - e**2))
+    return xp.sqrt(x1**2 + x2**2 / (1 - e**2))
 
 
-def capital_F_from(chi: np.ndarray) -> np.ndarray:
+def capital_F_from(chi, xp=np):
     """
     Equation 16 from Heyrovský & Karamazov.
 
@@ -36,19 +36,30 @@ def capital_F_from(chi: np.ndarray) -> np.ndarray:
     -------
     F(chi)
     """
-    F = np.zeros(chi.shape)
+    chi = xp.asarray(chi)
+    eps = 1e-12
 
-    root_min = np.sqrt(1 - chi[chi < 1] ** 2)
-    F[chi < 1] = np.arctanh(root_min) / root_min
-    F[chi == 1] = 1
-    root_plus = np.sqrt(chi[chi > 1] ** 2 - 1)
-    F[chi > 1] = np.arctan(root_plus) / root_plus
-    return F
+    less = chi < 1 - eps
+    greater = chi > 1 + eps
 
+    root_min_arg = xp.where(less, 1 - chi ** 2, 0.0)
+    root_min = xp.sqrt(root_min_arg)
+    root_min_safe = xp.where(less, root_min, 1.0)
+    F_less = xp.where(less, xp.arctanh(root_min) / root_min_safe, 0.0)
 
-def kappa_from(k_s: float, a: np.ndarray) -> np.ndarray:
+    root_plus_arg = xp.where(greater, chi ** 2 - 1, 0.0)
+    root_plus = xp.sqrt(root_plus_arg)
+    root_plus_safe = xp.where(greater, root_plus, 1.0)
+    F_greater = xp.where(greater, xp.arctan(root_plus) / root_plus_safe, 0.0)
+
+    F_equal = xp.ones_like(chi)
+
+    return xp.where(less, F_less, xp.where(greater, F_greater, F_equal))
+
+
+def kappa_from(k_s, a, xp=np):
     """
-    Equation 16 from Heyrovský & Karamazov.
+    Equation 21 from Heyrovský & Karamazov.
 
     Parameters
     ----------
@@ -61,13 +72,12 @@ def kappa_from(k_s: float, a: np.ndarray) -> np.ndarray:
     -------
     Convergence as a function of a
     """
-    F = capital_F_from(a)
-    kappa = 2 * k_s * (1 - F) / (a**2 - 1)
-    kappa[a == 1] = 2 / 3 * k_s
+    F = capital_F_from(a, xp=xp)
+    kappa = xp.where(xp.abs(a - 1) < 1e-12, 2/3 * k_s, 2 * k_s * (1 - F) / (a**2 - 1))
     return kappa
 
 
-def small_f_1(x1: np.ndarray, x2: np.ndarray, e: float) -> np.ndarray:
+def small_f_1(x1, x2, e, xp=np):
     """
     Equation 32 HK+24
 
@@ -84,13 +94,13 @@ def small_f_1(x1: np.ndarray, x2: np.ndarray, e: float) -> np.ndarray:
     -------
     f_1
     """
-    a = semi_major_axis_from(x1, x2, e)
-    F = capital_F_from(a)
+    a = semi_major_axis_from(x1, x2, e, xp=xp)
+    F = capital_F_from(a, xp=xp)
     f1 = (1 - e**2) ** (-1 / 2) * F
     return f1
 
 
-def small_f_2(x1: np.ndarray, x2: np.ndarray, e: float) -> np.ndarray:
+def small_f_2(x1, x2, e, xp=np):
     """
     Equation 32 HK+24
 
@@ -108,12 +118,12 @@ def small_f_2(x1: np.ndarray, x2: np.ndarray, e: float) -> np.ndarray:
     f_3
 
     """
-    norm = np.sqrt(x1**2 + x2**2)
-    f2 = np.log(norm / (1 + np.sqrt(1 - e**2)))
+    norm = xp.sqrt(x1**2 + x2**2)
+    f2 = xp.log(norm / (1 + xp.sqrt(1 - e**2)))
     return f2
 
 
-def small_f_3(x1: np.ndarray, x2: np.ndarray, e: float) -> np.ndarray:
+def small_f_3(x1, x2, e, xp=np):
     """
     Equation 32 HK+24
 
@@ -131,12 +141,12 @@ def small_f_3(x1: np.ndarray, x2: np.ndarray, e: float) -> np.ndarray:
     f_3
 
     """
-    root = np.sqrt(1 - e**2)
-    f3 = np.arctan(x1 * x2 * (1 - root) / (x1**2 * root + x2**2))
+    root = xp.sqrt(1 - e**2)
+    f3 = xp.arctan(x1 * x2 * (1 - root) / (x1**2 * root + x2**2))
     return f3
 
 
-def small_f_0(x1: np.ndarray, x2: np.ndarray, e: float) -> np.ndarray:
+def small_f_0(x1, x2, e, xp=np):
     """
     Equation 37 HK+24
 
@@ -154,9 +164,9 @@ def small_f_0(x1: np.ndarray, x2: np.ndarray, e: float) -> np.ndarray:
     f_0
 
     """
-    a = semi_major_axis_from(x1, x2, e)
-    F = capital_F_from(a)
-    pre_factor = 1 / (2 * np.sqrt(1 - e**2))
+    a = semi_major_axis_from(x1, x2, e, xp=xp)
+    F = capital_F_from(a, xp=xp)
+    pre_factor = 1 / (2 * xp.sqrt(1 - e**2))
     nominator = x1**2 + x2**2 + e**2 - 2 + (1 - e**2 * x1**2) * F
     denominator = 1 - x1**2 - x2**2 / (1 - e**2)
 
@@ -165,7 +175,7 @@ def small_f_0(x1: np.ndarray, x2: np.ndarray, e: float) -> np.ndarray:
     return f0
 
 
-def g1_g2_from(x1, x2, e, k_s):
+def g1_g2_from(x1, x2, e, k_s, xp=np):
     """
     Both components of the shear
 
@@ -189,16 +199,16 @@ def g1_g2_from(x1, x2, e, k_s):
     """
 
     # Factorized functions g1 and g2
-    f0 = small_f_0(x1, x2, e)
-    f1 = small_f_1(x1, x2, e)
-    f2 = small_f_2(x1, x2, e)
-    f3 = small_f_3(x1, x2, e)
+    f0 = small_f_0(x1, x2, e, xp=xp)
+    f1 = small_f_1(x1, x2, e, xp=xp)
+    f2 = small_f_2(x1, x2, e, xp=xp)
+    f3 = small_f_3(x1, x2, e, xp=xp)
 
     # Prefactor for both g1 and g2
     full_pre_factor = (
         4
         * k_s
-        * np.sqrt(1 - e**2)
+        * xp.sqrt(1 - e**2)
         / (((x1 - e) ** 2 + x2**2) ** 2 * ((x1 + e) ** 2 + x2**2) ** 2)
     )
 

test_autogalaxy/profiles/mass/dark/test_nfw.py

@@ -5,6 +5,68 @@
 
 grid = ag.Grid2DIrregular([[1.0, 1.0], [2.0, 2.0], [3.0, 3.0], [2.0, 4.0]])
 
+def test__deflections_via_analytical_from():
+
+    nfw = ag.mp.NFW(
+        centre=(0.0, 0.0),
+        ell_comps=(0.0, 0.0),
+        kappa_s=1.0,
+        scale_radius=1.0,
+    )
+
+    deflections = nfw.deflections_2d_via_analytic_from(
+        grid=ag.Grid2DIrregular([[0.1625, 0.1625]])
+    )
+
+    assert deflections[0, 0] == pytest.approx(0.56194, 1e-3)
+    assert deflections[0, 1] == pytest.approx(0.56194, 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,
+    )
+
+    deflections = nfw.deflections_2d_via_analytic_from(
+        grid=ag.Grid2DIrregular([(0.1625, 0.1625)])
+    )
+
+    assert deflections[0, 0] == pytest.approx(-2.59480, 1e-3)
+    assert deflections[0, 1] == pytest.approx(-0.44204, 1e-3)
+
+    nfw = ag.mp.NFW(
+        centre=(0.0, 0.0),
+        ell_comps=(0.0, 0.0),
+        kappa_s=1.0,
+        scale_radius=1.0,
+    )
+
+    deflections_via_analytic = nfw.deflections_2d_via_analytic_from(
+        grid=ag.Grid2DIrregular([[0.0, 0.0]])
+    )
+    deflections_via_cse = nfw.deflections_2d_via_cse_from(
+        grid=ag.Grid2DIrregular([[0.0, 0.0]])
+    )
+
+    assert deflections_via_analytic == pytest.approx(deflections_via_cse.array, 1.0e-4)
+
+    nfw = ag.mp.NFW(
+        centre=(0.3, -0.2),
+        ell_comps=(0.09, 0.172614),
+        kappa_s=0.05,
+        scale_radius=4.0,
+    )
+
+    deflections_via_analytic = nfw.deflections_2d_via_analytic_from(
+        grid=ag.Grid2DIrregular([[0.1625, -0.1625]])
+    )
+    deflections_via_cse = nfw.deflections_2d_via_cse_from(
+        grid=ag.Grid2DIrregular([[0.1625, -0.1625]])
+    )
+
+    assert deflections_via_analytic == pytest.approx(deflections_via_cse.array, 1.0e-4)
+
 
 def test__deflections_via_integral_from():
     nfw = ag.mp.NFWSph(centre=(0.0, 0.0), kappa_s=1.0, scale_radius=1.0)