From 3a1b4ad4bf07a0ae3c408201a8828e0069339980 Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Wed, 11 Feb 2026 11:48:46 +0000
Subject: [PATCH 01/19] jaxified the helper functions

---
 .../profiles/mass/dark/nfw_hk24_util.py       | 73 ++++++++++---------
 1 file changed, 38 insertions(+), 35 deletions(-)

diff --git a/autogalaxy/profiles/mass/dark/nfw_hk24_util.py b/autogalaxy/profiles/mass/dark/nfw_hk24_util.py
index 0d24b679d..3b26c6af7 100644
--- a/autogalaxy/profiles/mass/dark/nfw_hk24_util.py
+++ b/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,9 @@ 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 +35,24 @@ def capital_F_from(chi: np.ndarray) -> np.ndarray:
     -------
     F(chi)
     """
-    F = np.zeros(chi.shape)
+    chi = xp.asarray(chi)
+    eps = 1e-12
+
+    less = chi < 1 - eps
+    greater = chi > 1 + eps
 
-    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
+    root_min = xp.sqrt(1 - chi**2)
+    root_plus = xp.sqrt(chi**2 - 1)
 
+    F_less = xp.arctanh(root_min) / root_min
+    F_greater = xp.arctan(root_plus) / root_plus
+    F_equal = xp.ones_like(chi)
 
-def kappa_from(k_s: float, a: np.ndarray) -> np.ndarray:
+    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 +65,12 @@ def kappa_from(k_s: float, a: np.ndarray) -> np.ndarray:
     -------
     Convergence as a function of a
     """
-    F = capital_F_from(a)
+    F = capital_F_from(a, xp=xp)
     kappa = 2 * k_s * (1 - F) / (a**2 - 1)
-    kappa[a == 1] = 2 / 3 * k_s
-    return kappa
+    return xp.where(xp.abs(a - 1) < 1e-12, 2/3 * k_s, kappa)
 
 
-def small_f_1(x1: np.ndarray, x2: np.ndarray, e: float) -> np.ndarray:
+def small_f_1(x1: np.ndarray, x2: np.ndarray, e: float, xp=np) -> np.ndarray:
     """
     Equation 32 HK+24
 
@@ -84,13 +87,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: np.ndarray, x2: np.ndarray, e: float, xp=np) -> np.ndarray:
     """
     Equation 32 HK+24
 
@@ -108,12 +111,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: np.ndarray, x2: np.ndarray, e: float, xp=np) -> np.ndarray:
     """
     Equation 32 HK+24
 
@@ -131,12 +134,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: np.ndarray, x2: np.ndarray, e: float, xp=np) -> np.ndarray:
     """
     Equation 37 HK+24
 
@@ -154,9 +157,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 +168,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 +192,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)
     )
 

From db506e11dd6a1c26eaccb2df0d1e19fa5e04dafd Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Wed, 11 Feb 2026 15:13:25 +0000
Subject: [PATCH 02/19] added analytical deflection calc for eNFW and jaxified
 some functions

---
 autogalaxy/profiles/mass/dark/nfw.py | 83 +++++++++++++++++++++++-----
 1 file changed, 68 insertions(+), 15 deletions(-)

diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py
index ccf9a9eaf..cae8d377c 100644
--- a/autogalaxy/profiles/mass/dark/nfw.py
+++ b/autogalaxy/profiles/mass/dark/nfw.py
@@ -42,8 +42,61 @@ 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):
+        return self.deflections_2d_via_analytic_from(grid=grid, xp=xp)
+
+    @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) / xp.sqrt(1 + 2 * e_autolens + e_autolens ** 2)
+
+        # 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 = xp.where(np.abs(x1) < 1e-6, 1e-6, x1)
+        x2 = xp.where(np.abs(x2) < 1e-6, 1e-6, x2)
+
+        prefactor = 4 * self.kappa_s * xp.sqrt(1 - e_hk24**2) / (
+            ((x1 - e_hk24)**2 + x2**2) * ((x1 + e_hk24)**2 + x2**2)
+        )
+
+        deflection_x = ((x1 * ((x1**2 - e_hk24**2) * (1 - e_hk24**2) + x2**2 * (1 + e_hk24**2)) * nfw_hk24_util.small_f_1(x1, x2, e_hk24, xp=xp)
+                        + x1 * (x1**2 + x2**2 - e_hk24**2) * nfw_hk24_util.small_f_2(x1, x2, e_hk24, xp=xp))
+                        - x2 * (x1**2 + x2**2 + e_hk24**2) * nfw_hk24_util.small_f_3(x1, x2, e_hk24, xp=xp))
+
+        deflection_y = ((x2 * (x1**2 * (1 - 2 * e_hk24**2) + x2**2 + e_hk24**2) * nfw_hk24_util.small_f_1(x1, x2, e_hk24, xp=xp)
+                        + x2 * (x1**2 + x2**2 + e_hk24**2) * nfw_hk24_util.small_f_2(x1, x2, e_hk24, xp=xp))
+                        + x1 * (x1**2 + x2**2 - e_hk24**2) * nfw_hk24_util.small_f_3(x1, x2, e_hk24, xp=xp))
+
+        return self.rotated_grid_from_reference_frame_from(
+            grid=xp.multiply(prefactor, xp.vstack((deflection_y, deflection_x)).T),
+            xp=xp,
+            **kwargs,
+        )
+        # return self.transformed_from_reference_frame_grid_from(
+        #     grid=xp.multiply(prefactor, xp.vstack((deflection_y, deflection_x)).T),
+        #     centre=self.centre,
+        #     angle=self.angle(xp),
+        #     xp=xp
+        # )
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
@@ -146,7 +199,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,8 +324,8 @@ 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) / xp.sqrt(1 + 2 * e_autolens + e_autolens**2)
 
         # Define dimensionless length coords
 
@@ -280,17 +333,17 @@ def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         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 +368,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) / xp.sqrt(1 + 2 * e_autolens + e_autolens**2)
 
         # Define dimensionless length coords
 
@@ -325,13 +378,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):

From fba1d1035348d139fd79cdac54ca7650bb8ad46c Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Wed, 11 Feb 2026 15:16:12 +0000
Subject: [PATCH 03/19] added analytical deflection tests

---
 .../profiles/mass/dark/test_nfw.py            | 34 +++++++++++++++++--
 1 file changed, 32 insertions(+), 2 deletions(-)

diff --git a/test_autogalaxy/profiles/mass/dark/test_nfw.py b/test_autogalaxy/profiles/mass/dark/test_nfw.py
index 9260afcbd..097ba57ce 100644
--- a/test_autogalaxy/profiles/mass/dark/test_nfw.py
+++ b/test_autogalaxy/profiles/mass/dark/test_nfw.py
@@ -5,6 +5,36 @@
 
 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)
+
 
 def test__deflections_via_integral_from():
     nfw = ag.mp.NFWSph(centre=(0.0, 0.0), kappa_s=1.0, scale_radius=1.0)
@@ -157,11 +187,11 @@ def test__deflections_yx_2d_from():
     nfw = ag.mp.NFW(centre=(0.0, 0.0), kappa_s=1.0, scale_radius=1.0)
 
     deflections = nfw.deflections_yx_2d_from(grid=ag.Grid2DIrregular([[1.0, 0.0]]))
-    deflections_via_integral = nfw.deflections_2d_via_cse_from(
+    deflections_via_analytical = nfw.deflections_2d_via_analytic_from(
         grid=ag.Grid2DIrregular([[1.0, 0.0]])
     )
 
-    assert deflections == pytest.approx(deflections_via_integral.array, 1.0e-4)
+    assert deflections == pytest.approx(deflections_via_analytical.array, 1.0e-4)
 
 
 def test__convergence_2d_via_mge_from():

From 99644fc46acdb0328b7310ac1170f3c922532533 Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Wed, 11 Feb 2026 15:47:52 +0000
Subject: [PATCH 04/19] mathematical simplification of e_hk24

---
 autogalaxy/profiles/mass/dark/nfw.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py
index cae8d377c..6aa13764a 100644
--- a/autogalaxy/profiles/mass/dark/nfw.py
+++ b/autogalaxy/profiles/mass/dark/nfw.py
@@ -63,7 +63,7 @@ def deflections_2d_via_analytic_from(
         # 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) / xp.sqrt(1 + 2 * e_autolens + e_autolens ** 2)
+        e_hk24 = 2 * xp.sqrt(e_autolens) / (1 + e_autolens)
 
         # Define dimensionless length coords
 
@@ -325,7 +325,7 @@ def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         # 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) / xp.sqrt(1 + 2 * e_autolens + e_autolens**2)
+        e_hk24 = 2 * xp.sqrt(e_autolens) / (1 + e_autolens)
 
         # Define dimensionless length coords
 
@@ -369,7 +369,7 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         # 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) / xp.sqrt(1 + 2 * e_autolens + e_autolens**2)
+        e_hk24 = 2 * xp.sqrt(e_autolens) / (1 + e_autolens)
 
         # Define dimensionless length coords
 

From a1f360425cca6643864bbd67cc5b7d29153c5c82 Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Wed, 11 Feb 2026 16:46:50 +0000
Subject: [PATCH 05/19] multiply resulting deflection angles by scale_radius to
 go to physical deflection angles

---
 autogalaxy/profiles/mass/dark/nfw.py | 16 +++++-----------
 1 file changed, 5 insertions(+), 11 deletions(-)

diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py
index 6aa13764a..8c06a7079 100644
--- a/autogalaxy/profiles/mass/dark/nfw.py
+++ b/autogalaxy/profiles/mass/dark/nfw.py
@@ -78,25 +78,19 @@ def deflections_2d_via_analytic_from(
             ((x1 - e_hk24)**2 + x2**2) * ((x1 + e_hk24)**2 + x2**2)
         )
 
-        deflection_x = ((x1 * ((x1**2 - e_hk24**2) * (1 - e_hk24**2) + x2**2 * (1 + e_hk24**2)) * nfw_hk24_util.small_f_1(x1, x2, e_hk24, xp=xp)
-                        + x1 * (x1**2 + x2**2 - e_hk24**2) * nfw_hk24_util.small_f_2(x1, x2, e_hk24, xp=xp))
+        deflection_x = (x1 * ((x1**2 - e_hk24**2) * (1 - e_hk24**2) + x2**2 * (1 + e_hk24**2)) * nfw_hk24_util.small_f_1(x1, x2, e_hk24, xp=xp)
+                        + x1 * (x1**2 + x2**2 - e_hk24**2) * nfw_hk24_util.small_f_2(x1, x2, e_hk24, xp=xp)
                         - x2 * (x1**2 + x2**2 + e_hk24**2) * nfw_hk24_util.small_f_3(x1, x2, e_hk24, xp=xp))
 
-        deflection_y = ((x2 * (x1**2 * (1 - 2 * e_hk24**2) + x2**2 + e_hk24**2) * nfw_hk24_util.small_f_1(x1, x2, e_hk24, xp=xp)
-                        + x2 * (x1**2 + x2**2 + e_hk24**2) * nfw_hk24_util.small_f_2(x1, x2, e_hk24, xp=xp))
+        deflection_y = (x2 * (x1**2 * (1 - 2 * e_hk24**2) + x2**2 + e_hk24**2) * nfw_hk24_util.small_f_1(x1, x2, e_hk24, xp=xp)
+                        + x2 * (x1**2 + x2**2 + e_hk24**2) * nfw_hk24_util.small_f_2(x1, x2, e_hk24, xp=xp)
                         + x1 * (x1**2 + x2**2 - e_hk24**2) * nfw_hk24_util.small_f_3(x1, x2, e_hk24, xp=xp))
 
         return self.rotated_grid_from_reference_frame_from(
-            grid=xp.multiply(prefactor, xp.vstack((deflection_y, deflection_x)).T),
+            grid=xp.multiply(self.scale_radius, xp.multiply(prefactor, xp.vstack((deflection_y, deflection_x)).T)),
             xp=xp,
             **kwargs,
         )
-        # return self.transformed_from_reference_frame_grid_from(
-        #     grid=xp.multiply(prefactor, xp.vstack((deflection_y, deflection_x)).T),
-        #     centre=self.centre,
-        #     angle=self.angle(xp),
-        #     xp=xp
-        # )
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform

From 38fe2c10600dd013fa3b2cff673d5d2b357ba640 Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Wed, 11 Feb 2026 17:21:27 +0000
Subject: [PATCH 06/19] multiplying with prefactor before vstack

---
 autogalaxy/profiles/mass/dark/nfw.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py
index 8c06a7079..9dedc66b0 100644
--- a/autogalaxy/profiles/mass/dark/nfw.py
+++ b/autogalaxy/profiles/mass/dark/nfw.py
@@ -81,13 +81,15 @@ def deflections_2d_via_analytic_from(
         deflection_x = (x1 * ((x1**2 - e_hk24**2) * (1 - e_hk24**2) + x2**2 * (1 + e_hk24**2)) * nfw_hk24_util.small_f_1(x1, x2, e_hk24, xp=xp)
                         + x1 * (x1**2 + x2**2 - e_hk24**2) * nfw_hk24_util.small_f_2(x1, x2, e_hk24, xp=xp)
                         - x2 * (x1**2 + x2**2 + e_hk24**2) * nfw_hk24_util.small_f_3(x1, x2, e_hk24, xp=xp))
+        deflection_x *= prefactor
 
         deflection_y = (x2 * (x1**2 * (1 - 2 * e_hk24**2) + x2**2 + e_hk24**2) * nfw_hk24_util.small_f_1(x1, x2, e_hk24, xp=xp)
                         + x2 * (x1**2 + x2**2 + e_hk24**2) * nfw_hk24_util.small_f_2(x1, x2, e_hk24, xp=xp)
                         + x1 * (x1**2 + x2**2 - e_hk24**2) * nfw_hk24_util.small_f_3(x1, x2, e_hk24, xp=xp))
+        deflection_y *= prefactor
 
         return self.rotated_grid_from_reference_frame_from(
-            grid=xp.multiply(self.scale_radius, xp.multiply(prefactor, xp.vstack((deflection_y, deflection_x)).T)),
+            xp.multiply(self.scale_radius, xp.vstack((deflection_y, deflection_x)).T),
             xp=xp,
             **kwargs,
         )

From bd2daa609dd3362c3b2de58154125b8c3f3789f8 Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Thu, 12 Feb 2026 14:19:45 +0000
Subject: [PATCH 07/19] changed the handling near 0,0

---
 autogalaxy/profiles/mass/dark/nfw.py | 21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py
index 9dedc66b0..7fd6d920a 100644
--- a/autogalaxy/profiles/mass/dark/nfw.py
+++ b/autogalaxy/profiles/mass/dark/nfw.py
@@ -71,21 +71,26 @@ def deflections_2d_via_analytic_from(
         x2 = grid.array[:, 0] / self.scale_radius
 
         # Avoid nans due to x=0
-        x1 = xp.where(np.abs(x1) < 1e-6, 1e-6, x1)
-        x2 = xp.where(np.abs(x2) < 1e-6, 1e-6, x2)
+        eps = xp.array(1e-12)
+        x1 = xp.where(xp.abs(x1) < eps, xp.sign(x1) * eps, x1)
+        x2 = xp.where(xp.abs(x2) < eps, xp.sign(x2) * eps, x2)
 
         prefactor = 4 * self.kappa_s * xp.sqrt(1 - e_hk24**2) / (
             ((x1 - e_hk24)**2 + x2**2) * ((x1 + e_hk24)**2 + x2**2)
         )
 
-        deflection_x = (x1 * ((x1**2 - e_hk24**2) * (1 - e_hk24**2) + x2**2 * (1 + e_hk24**2)) * nfw_hk24_util.small_f_1(x1, x2, e_hk24, xp=xp)
-                        + x1 * (x1**2 + x2**2 - e_hk24**2) * nfw_hk24_util.small_f_2(x1, x2, e_hk24, xp=xp)
-                        - x2 * (x1**2 + x2**2 + e_hk24**2) * nfw_hk24_util.small_f_3(x1, x2, e_hk24, xp=xp))
+        f1 = nfw_hk24_util.small_f_1(x1, x2, e_hk24, xp=xp)
+        f2 = nfw_hk24_util.small_f_2(x1, x2, e_hk24, xp=xp)
+        f3 = nfw_hk24_util.small_f_3(x1, x2, e_hk24, xp=xp)
+
+        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) * nfw_hk24_util.small_f_1(x1, x2, e_hk24, xp=xp)
-                        + x2 * (x1**2 + x2**2 + e_hk24**2) * nfw_hk24_util.small_f_2(x1, x2, e_hk24, xp=xp)
-                        + x1 * (x1**2 + x2**2 - e_hk24**2) * nfw_hk24_util.small_f_3(x1, x2, e_hk24, xp=xp))
+        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(

From 7a940d7474a355d3de0409b0ae07d93553e0811d Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Thu, 12 Feb 2026 14:26:37 +0000
Subject: [PATCH 08/19] prevent nans at the centre of the mass distribution

---
 autogalaxy/profiles/mass/dark/nfw.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py
index 7fd6d920a..7baa4dd50 100644
--- a/autogalaxy/profiles/mass/dark/nfw.py
+++ b/autogalaxy/profiles/mass/dark/nfw.py
@@ -93,6 +93,12 @@ def deflections_2d_via_analytic_from(
                         + x1 * (x1**2 + x2**2 - e_hk24**2) * f3)
         deflection_y *= prefactor
 
+        # prevent nans at the centre
+        r2 = x1 ** 2 + x2 ** 2
+        mask0 = r2 < 1e-24
+        deflection_x = xp.where(mask0, 0.0, deflection_x)
+        deflection_y = xp.where(mask0, 0.0, deflection_y)
+
         return self.rotated_grid_from_reference_frame_from(
             xp.multiply(self.scale_radius, xp.vstack((deflection_y, deflection_x)).T),
             xp=xp,

From 3916fc853be229164e9ac642b9f4c3d3cca2af6d Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Thu, 12 Feb 2026 14:28:20 +0000
Subject: [PATCH 09/19] remove expected input types

---
 autogalaxy/profiles/mass/dark/nfw_hk24_util.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/autogalaxy/profiles/mass/dark/nfw_hk24_util.py b/autogalaxy/profiles/mass/dark/nfw_hk24_util.py
index 3b26c6af7..e8f214262 100644
--- a/autogalaxy/profiles/mass/dark/nfw_hk24_util.py
+++ b/autogalaxy/profiles/mass/dark/nfw_hk24_util.py
@@ -70,7 +70,7 @@ def kappa_from(k_s, a, xp=np):
     return xp.where(xp.abs(a - 1) < 1e-12, 2/3 * k_s, kappa)
 
 
-def small_f_1(x1: np.ndarray, x2: np.ndarray, e: float, xp=np) -> np.ndarray:
+def small_f_1(x1, x2, e, xp=np):
     """
     Equation 32 HK+24
 
@@ -93,7 +93,7 @@ def small_f_1(x1: np.ndarray, x2: np.ndarray, e: float, xp=np) -> np.ndarray:
     return f1
 
 
-def small_f_2(x1: np.ndarray, x2: np.ndarray, e: float, xp=np) -> np.ndarray:
+def small_f_2(x1, x2, e, xp=np):
     """
     Equation 32 HK+24
 
@@ -116,7 +116,7 @@ def small_f_2(x1: np.ndarray, x2: np.ndarray, e: float, xp=np) -> np.ndarray:
     return f2
 
 
-def small_f_3(x1: np.ndarray, x2: np.ndarray, e: float, xp=np) -> np.ndarray:
+def small_f_3(x1, x2, e, xp=np):
     """
     Equation 32 HK+24
 
@@ -139,7 +139,7 @@ def small_f_3(x1: np.ndarray, x2: np.ndarray, e: float, xp=np) -> np.ndarray:
     return f3
 
 
-def small_f_0(x1: np.ndarray, x2: np.ndarray, e: float, xp=np) -> np.ndarray:
+def small_f_0(x1, x2, e, xp=np):
     """
     Equation 37 HK+24
 

From 0ea9abc92a3803e714e1e699cd33e8cdb92565b7 Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Thu, 12 Feb 2026 14:29:59 +0000
Subject: [PATCH 10/19] changed the test to a less symmetric case

---
 autogalaxy/profiles/mass/dark/nfw_test.ipynb  | 314 ++++++++++++++++++
 .../profiles/mass/dark/test_nfw.py            |   2 +-
 2 files changed, 315 insertions(+), 1 deletion(-)
 create mode 100644 autogalaxy/profiles/mass/dark/nfw_test.ipynb

diff --git a/autogalaxy/profiles/mass/dark/nfw_test.ipynb b/autogalaxy/profiles/mass/dark/nfw_test.ipynb
new file mode 100644
index 000000000..226a7ed59
--- /dev/null
+++ b/autogalaxy/profiles/mass/dark/nfw_test.ipynb
@@ -0,0 +1,314 @@
+{
+ "cells": [
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-02-12T14:24:56.097874Z",
+     "start_time": "2026-02-12T14:24:51.586923Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "import jax.numpy as jnp\n",
+    "import numpy as np\n",
+    "import jax\n",
+    "jax.config.update(\"jax_enable_x64\", True)\n",
+    "import autogalaxy as ag\n",
+    "import autolens as al\n",
+    "\n",
+    "grid = ag.Grid2DIrregular([[1.0, 1.0], [2.0, 2.0], [3.0, 3.0], [2.0, 4.0]])\n",
+    "\n",
+    "grid = ag.Grid2DIrregular([[1.0, 1.0], [2.0, 2.0], [3.0, 3.0], [2.0, 4.0]])\n",
+    "nfw = ag.mp.NFW(\n",
+    "        centre=(0.0, 0.0),\n",
+    "        ell_comps=(0.0, 0.111111),\n",
+    "        kappa_s=1.0,\n",
+    "        scale_radius=20.0,\n",
+    "    )\n",
+    "gnfw = ag.mp.gNFW(\n",
+    "        centre=(0.0, 0.0),\n",
+    "        ell_comps=(0.0, 0.111111),\n",
+    "        kappa_s=1.0,\n",
+    "        inner_slope=1.0,\n",
+    "        scale_radius=20.0,\n",
+    ")\n"
+   ],
+   "id": "2e4b1f6e3163e4e3",
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "C:\\Users\\c4072114\\PycharmProjects\\PyAuto\\PyAutoGalaxy\\autogalaxy\\profiles\\mass\\dark\\nfw.py:20: SyntaxWarning: invalid escape sequence '\\|'\n",
+      "  \"\"\"\n"
+     ]
+    }
+   ],
+   "execution_count": 1
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-02-12T14:14:08.874004Z",
+     "start_time": "2026-02-12T14:14:07.935866Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "print(nfw.deflections_yx_2d_from(grid,xp=np))\n",
+    "print(nfw.deflections_yx_2d_from(grid, xp=jnp))"
+   ],
+   "id": "f6c1af783d7268a7",
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "VectorYX2DIrregular([[ 6.03220218,  4.92448729],\n",
+      "       [ 9.09858068,  7.47079345],\n",
+      "       [11.15895124,  9.20530545],\n",
+      "       [ 7.36924937, 12.14404003]])\n",
+      "2026-02-12 14:14:08,002 - jax._src.xla_bridge - INFO - Unable to initialize backend 'cuda': \n",
+      "2026-02-12 14:14:08,003 - jax._src.xla_bridge - INFO - Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
+      "2026-02-12 14:14:08,007 - jax._src.xla_bridge - INFO - Unable to initialize backend 'tpu': UNIMPLEMENTED: LoadPjrtPlugin is not implemented on windows yet.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "C:\\Users\\c4072114\\PycharmProjects\\PyAuto\\PyAutoGalaxy\\autogalaxy\\profiles\\mass\\dark\\nfw_hk24_util.py:45: RuntimeWarning: invalid value encountered in sqrt\n",
+      "  root_plus = xp.sqrt(chi**2 - 1)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Array([[ 6.03220218,  4.92448729],\n",
+      "       [ 9.09858068,  7.47079345],\n",
+      "       [11.15895124,  9.20530545],\n",
+      "       [ 7.36924937, 12.14404003]], dtype=float64)\n"
+     ]
+    }
+   ],
+   "execution_count": 2
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-02-12T14:14:16.941162Z",
+     "start_time": "2026-02-12T14:14:16.924363Z"
+    }
+   },
+   "cell_type": "code",
+   "source": "gnfw.deflections_yx_2d_from(grid)",
+   "id": "4993cd02ab7e8c86",
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "VectorYX2DIrregular([[ 6.03074243,  4.92332473],\n",
+       "       [ 9.09582031,  7.4685857 ],\n",
+       "       [11.15476335,  9.20198906],\n",
+       "       [ 7.3664635 , 12.1396311 ]])"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 3
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-02-12T14:24:59.961666Z",
+     "start_time": "2026-02-12T14:24:58.834078Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "grid = al.Grid2D.uniform(shape_native=(11, 11), pixel_scales=0.2,\n",
+    "                             over_sample_size=8)\n",
+    "\n",
+    "lens_galaxy = al.Galaxy(\n",
+    "    redshift=0.5,\n",
+    "    mass_0=ag.mp.NFW(\n",
+    "        centre=(0.0, 0.00),\n",
+    "        ell_comps=(0.2, 0.111111),\n",
+    "        kappa_s=0.1,\n",
+    "        scale_radius=20.0,\n",
+    "    ),\n",
+    ")\n",
+    "print(lens_galaxy.mass_0.deflections_yx_2d_from(grid=grid, xp=jnp))\n",
+    "print(np.max(lens_galaxy.mass_0.deflections_yx_2d_from(grid=grid, xp=jnp)))"
+   ],
+   "id": "59cf372ee6ae833c",
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2026-02-12 14:24:58,908 - jax._src.xla_bridge - INFO - Unable to initialize backend 'cuda': \n",
+      "2026-02-12 14:24:58,909 - jax._src.xla_bridge - INFO - Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
+      "2026-02-12 14:24:58,913 - jax._src.xla_bridge - INFO - Unable to initialize backend 'tpu': UNIMPLEMENTED: LoadPjrtPlugin is not implemented on windows yet.\n",
+      "Array([[ 0.64421232, -0.54423054],\n",
+      "       [ 0.64868502, -0.46953412],\n",
+      "       [ 0.65154284, -0.3883424 ],\n",
+      "       [ 0.65165847, -0.3002847 ],\n",
+      "       [ 0.6475545 , -0.2055496 ],\n",
+      "       [ 0.63766224, -0.10527277],\n",
+      "       [ 0.62087922, -0.0017213 ],\n",
+      "       [ 0.59716429,  0.10203609],\n",
+      "       [ 0.56769437,  0.20292033],\n",
+      "       [ 0.53442151,  0.29860742],\n",
+      "       [ 0.49938556,  0.38784218],\n",
+      "       [ 0.55699033, -0.54888684],\n",
+      "       [ 0.56089008, -0.47323704],\n",
+      "       [ 0.56365885, -0.39013744],\n",
+      "       [ 0.56403146, -0.29870151],\n",
+      "       [ 0.56003098, -0.19866231],\n",
+      "       [ 0.54923128, -0.09122931],\n",
+      "       [ 0.52983592,  0.02027455],\n",
+      "       [ 0.50203998,  0.13104815],\n",
+      "       [ 0.46819426,  0.23665964],\n",
+      "       [ 0.43149302,  0.33449321],\n",
+      "       [ 0.3945875 ,  0.42382454],\n",
+      "       [ 0.46072474, -0.552298  ],\n",
+      "       [ 0.46312394, -0.47623446],\n",
+      "       [ 0.46508563, -0.39184904],\n",
+      "       [ 0.46543342, -0.29743107],\n",
+      "       [ 0.46169148, -0.19162379],\n",
+      "       [ 0.44991531, -0.07524281],\n",
+      "       [ 0.42680477,  0.04646023],\n",
+      "       [ 0.39339841,  0.16512372],\n",
+      "       [ 0.35472166,  0.27440954],\n",
+      "       [ 0.31566035,  0.37238804],\n",
+      "       [ 0.27884834,  0.45990214],\n",
+      "       [ 0.35433802, -0.55320829],\n",
+      "       [ 0.35356863, -0.47726693],\n",
+      "       [ 0.35302519, -0.39249966],\n",
+      "       [ 0.35212514, -0.29618816],\n",
+      "       [ 0.34852334, -0.18473489],\n",
+      "       [ 0.33578884, -0.05660441],\n",
+      "       [ 0.30671054,  0.07901422],\n",
+      "       [ 0.26584619,  0.20515544],\n",
+      "       [ 0.22400385,  0.3146805 ],\n",
+      "       [ 0.18622396,  0.40959294],\n",
+      "       [ 0.15312727,  0.49332251],\n",
+      "       [ 0.23739015, -0.54971459],\n",
+      "       [ 0.23087867, -0.47399256],\n",
+      "       [ 0.22442588, -0.38950352],\n",
+      "       [ 0.21830314, -0.29302227],\n",
+      "       [ 0.21224924, -0.17815385],\n",
+      "       [ 0.19863225, -0.03383159],\n",
+      "       [ 0.15811807,  0.1216322 ],\n",
+      "       [ 0.11272857,  0.24831262],\n",
+      "       [ 0.07611139,  0.35195516],\n",
+      "       [ 0.0462995 ,  0.44130356],\n",
+      "       [ 0.02114958,  0.52065869],\n",
+      "       [ 0.11101004, -0.53947865],\n",
+      "       [ 0.09584853, -0.46297986],\n",
+      "       [ 0.07872659, -0.37770426],\n",
+      "       [ 0.05893757, -0.2805206 ],\n",
+      "       [ 0.03500201, -0.16486547],\n",
+      "       [ 0.        ,  0.        ],\n",
+      "       [-0.03500201,  0.16486547],\n",
+      "       [-0.05893757,  0.2805206 ],\n",
+      "       [-0.07872659,  0.37770426],\n",
+      "       [-0.09584853,  0.46297986],\n",
+      "       [-0.11101004,  0.53947865],\n",
+      "       [-0.02114958, -0.52065869],\n",
+      "       [-0.0462995 , -0.44130356],\n",
+      "       [-0.07611139, -0.35195516],\n",
+      "       [-0.11272857, -0.24831262],\n",
+      "       [-0.15811807, -0.1216322 ],\n",
+      "       [-0.19863225,  0.03383159],\n",
+      "       [-0.21224924,  0.17815385],\n",
+      "       [-0.21830314,  0.29302227],\n",
+      "       [-0.22442588,  0.38950352],\n",
+      "       [-0.23087867,  0.47399256],\n",
+      "       [-0.23739015,  0.54971459],\n",
+      "       [-0.15312727, -0.49332251],\n",
+      "       [-0.18622396, -0.40959294],\n",
+      "       [-0.22400385, -0.3146805 ],\n",
+      "       [-0.26584619, -0.20515544],\n",
+      "       [-0.30671054, -0.07901422],\n",
+      "       [-0.33578884,  0.05660441],\n",
+      "       [-0.34852334,  0.18473489],\n",
+      "       [-0.35212514,  0.29618816],\n",
+      "       [-0.35302519,  0.39249966],\n",
+      "       [-0.35356863,  0.47726693],\n",
+      "       [-0.35433802,  0.55320829],\n",
+      "       [-0.27884834, -0.45990214],\n",
+      "       [-0.31566035, -0.37238804],\n",
+      "       [-0.35472166, -0.27440954],\n",
+      "       [-0.39339841, -0.16512372],\n",
+      "       [-0.42680477, -0.04646023],\n",
+      "       [-0.44991531,  0.07524281],\n",
+      "       [-0.46169148,  0.19162379],\n",
+      "       [-0.46543342,  0.29743107],\n",
+      "       [-0.46508563,  0.39184904],\n",
+      "       [-0.46312394,  0.47623446],\n",
+      "       [-0.46072474,  0.552298  ],\n",
+      "       [-0.3945875 , -0.42382454],\n",
+      "       [-0.43149302, -0.33449321],\n",
+      "       [-0.46819426, -0.23665964],\n",
+      "       [-0.50203998, -0.13104815],\n",
+      "       [-0.52983592, -0.02027455],\n",
+      "       [-0.54923128,  0.09122931],\n",
+      "       [-0.56003098,  0.19866231],\n",
+      "       [-0.56403146,  0.29870151],\n",
+      "       [-0.56365885,  0.39013744],\n",
+      "       [-0.56089008,  0.47323704],\n",
+      "       [-0.55699033,  0.54888684],\n",
+      "       [-0.49938556, -0.38784218],\n",
+      "       [-0.53442151, -0.29860742],\n",
+      "       [-0.56769437, -0.20292033],\n",
+      "       [-0.59716429, -0.10203609],\n",
+      "       [-0.62087922,  0.0017213 ],\n",
+      "       [-0.63766224,  0.10527277],\n",
+      "       [-0.6475545 ,  0.2055496 ],\n",
+      "       [-0.65165847,  0.3002847 ],\n",
+      "       [-0.65154284,  0.3883424 ],\n",
+      "       [-0.64868502,  0.46953412],\n",
+      "       [-0.64421232,  0.54423054]], dtype=float64)\n",
+      "0.651658467012907\n"
+     ]
+    }
+   ],
+   "execution_count": 2
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": "",
+   "id": "730f63354e47eb46"
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/test_autogalaxy/profiles/mass/dark/test_nfw.py b/test_autogalaxy/profiles/mass/dark/test_nfw.py
index 097ba57ce..6f3ea4396 100644
--- a/test_autogalaxy/profiles/mass/dark/test_nfw.py
+++ b/test_autogalaxy/profiles/mass/dark/test_nfw.py
@@ -184,7 +184,7 @@ def test__deflections_2d__numerical_precision_of_csv_compared_to_integral():
 
 
 def test__deflections_yx_2d_from():
-    nfw = ag.mp.NFW(centre=(0.0, 0.0), kappa_s=1.0, scale_radius=1.0)
+    nfw = ag.mp.NFW(centre=(0.1, -0.1), ell_comps=(0.3, 0.1), kappa_s=1.0, scale_radius=1.0)
 
     deflections = nfw.deflections_yx_2d_from(grid=ag.Grid2DIrregular([[1.0, 0.0]]))
     deflections_via_analytical = nfw.deflections_2d_via_analytic_from(

From 60e1ffb56868e407a2eed7b4c5a9d89764794311 Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Thu, 12 Feb 2026 14:36:24 +0000
Subject: [PATCH 11/19] stop tracking nfw_test.ipynb

---
 autogalaxy/profiles/mass/dark/nfw_test.ipynb | 314 -------------------
 1 file changed, 314 deletions(-)
 delete mode 100644 autogalaxy/profiles/mass/dark/nfw_test.ipynb

diff --git a/autogalaxy/profiles/mass/dark/nfw_test.ipynb b/autogalaxy/profiles/mass/dark/nfw_test.ipynb
deleted file mode 100644
index 226a7ed59..000000000
--- a/autogalaxy/profiles/mass/dark/nfw_test.ipynb
+++ /dev/null
@@ -1,314 +0,0 @@
-{
- "cells": [
-  {
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2026-02-12T14:24:56.097874Z",
-     "start_time": "2026-02-12T14:24:51.586923Z"
-    }
-   },
-   "cell_type": "code",
-   "source": [
-    "import jax.numpy as jnp\n",
-    "import numpy as np\n",
-    "import jax\n",
-    "jax.config.update(\"jax_enable_x64\", True)\n",
-    "import autogalaxy as ag\n",
-    "import autolens as al\n",
-    "\n",
-    "grid = ag.Grid2DIrregular([[1.0, 1.0], [2.0, 2.0], [3.0, 3.0], [2.0, 4.0]])\n",
-    "\n",
-    "grid = ag.Grid2DIrregular([[1.0, 1.0], [2.0, 2.0], [3.0, 3.0], [2.0, 4.0]])\n",
-    "nfw = ag.mp.NFW(\n",
-    "        centre=(0.0, 0.0),\n",
-    "        ell_comps=(0.0, 0.111111),\n",
-    "        kappa_s=1.0,\n",
-    "        scale_radius=20.0,\n",
-    "    )\n",
-    "gnfw = ag.mp.gNFW(\n",
-    "        centre=(0.0, 0.0),\n",
-    "        ell_comps=(0.0, 0.111111),\n",
-    "        kappa_s=1.0,\n",
-    "        inner_slope=1.0,\n",
-    "        scale_radius=20.0,\n",
-    ")\n"
-   ],
-   "id": "2e4b1f6e3163e4e3",
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "C:\\Users\\c4072114\\PycharmProjects\\PyAuto\\PyAutoGalaxy\\autogalaxy\\profiles\\mass\\dark\\nfw.py:20: SyntaxWarning: invalid escape sequence '\\|'\n",
-      "  \"\"\"\n"
-     ]
-    }
-   ],
-   "execution_count": 1
-  },
-  {
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2026-02-12T14:14:08.874004Z",
-     "start_time": "2026-02-12T14:14:07.935866Z"
-    }
-   },
-   "cell_type": "code",
-   "source": [
-    "print(nfw.deflections_yx_2d_from(grid,xp=np))\n",
-    "print(nfw.deflections_yx_2d_from(grid, xp=jnp))"
-   ],
-   "id": "f6c1af783d7268a7",
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "VectorYX2DIrregular([[ 6.03220218,  4.92448729],\n",
-      "       [ 9.09858068,  7.47079345],\n",
-      "       [11.15895124,  9.20530545],\n",
-      "       [ 7.36924937, 12.14404003]])\n",
-      "2026-02-12 14:14:08,002 - jax._src.xla_bridge - INFO - Unable to initialize backend 'cuda': \n",
-      "2026-02-12 14:14:08,003 - jax._src.xla_bridge - INFO - Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
-      "2026-02-12 14:14:08,007 - jax._src.xla_bridge - INFO - Unable to initialize backend 'tpu': UNIMPLEMENTED: LoadPjrtPlugin is not implemented on windows yet.\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "C:\\Users\\c4072114\\PycharmProjects\\PyAuto\\PyAutoGalaxy\\autogalaxy\\profiles\\mass\\dark\\nfw_hk24_util.py:45: RuntimeWarning: invalid value encountered in sqrt\n",
-      "  root_plus = xp.sqrt(chi**2 - 1)\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Array([[ 6.03220218,  4.92448729],\n",
-      "       [ 9.09858068,  7.47079345],\n",
-      "       [11.15895124,  9.20530545],\n",
-      "       [ 7.36924937, 12.14404003]], dtype=float64)\n"
-     ]
-    }
-   ],
-   "execution_count": 2
-  },
-  {
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2026-02-12T14:14:16.941162Z",
-     "start_time": "2026-02-12T14:14:16.924363Z"
-    }
-   },
-   "cell_type": "code",
-   "source": "gnfw.deflections_yx_2d_from(grid)",
-   "id": "4993cd02ab7e8c86",
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "VectorYX2DIrregular([[ 6.03074243,  4.92332473],\n",
-       "       [ 9.09582031,  7.4685857 ],\n",
-       "       [11.15476335,  9.20198906],\n",
-       "       [ 7.3664635 , 12.1396311 ]])"
-      ]
-     },
-     "execution_count": 3,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "execution_count": 3
-  },
-  {
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2026-02-12T14:24:59.961666Z",
-     "start_time": "2026-02-12T14:24:58.834078Z"
-    }
-   },
-   "cell_type": "code",
-   "source": [
-    "grid = al.Grid2D.uniform(shape_native=(11, 11), pixel_scales=0.2,\n",
-    "                             over_sample_size=8)\n",
-    "\n",
-    "lens_galaxy = al.Galaxy(\n",
-    "    redshift=0.5,\n",
-    "    mass_0=ag.mp.NFW(\n",
-    "        centre=(0.0, 0.00),\n",
-    "        ell_comps=(0.2, 0.111111),\n",
-    "        kappa_s=0.1,\n",
-    "        scale_radius=20.0,\n",
-    "    ),\n",
-    ")\n",
-    "print(lens_galaxy.mass_0.deflections_yx_2d_from(grid=grid, xp=jnp))\n",
-    "print(np.max(lens_galaxy.mass_0.deflections_yx_2d_from(grid=grid, xp=jnp)))"
-   ],
-   "id": "59cf372ee6ae833c",
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "2026-02-12 14:24:58,908 - jax._src.xla_bridge - INFO - Unable to initialize backend 'cuda': \n",
-      "2026-02-12 14:24:58,909 - jax._src.xla_bridge - INFO - Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
-      "2026-02-12 14:24:58,913 - jax._src.xla_bridge - INFO - Unable to initialize backend 'tpu': UNIMPLEMENTED: LoadPjrtPlugin is not implemented on windows yet.\n",
-      "Array([[ 0.64421232, -0.54423054],\n",
-      "       [ 0.64868502, -0.46953412],\n",
-      "       [ 0.65154284, -0.3883424 ],\n",
-      "       [ 0.65165847, -0.3002847 ],\n",
-      "       [ 0.6475545 , -0.2055496 ],\n",
-      "       [ 0.63766224, -0.10527277],\n",
-      "       [ 0.62087922, -0.0017213 ],\n",
-      "       [ 0.59716429,  0.10203609],\n",
-      "       [ 0.56769437,  0.20292033],\n",
-      "       [ 0.53442151,  0.29860742],\n",
-      "       [ 0.49938556,  0.38784218],\n",
-      "       [ 0.55699033, -0.54888684],\n",
-      "       [ 0.56089008, -0.47323704],\n",
-      "       [ 0.56365885, -0.39013744],\n",
-      "       [ 0.56403146, -0.29870151],\n",
-      "       [ 0.56003098, -0.19866231],\n",
-      "       [ 0.54923128, -0.09122931],\n",
-      "       [ 0.52983592,  0.02027455],\n",
-      "       [ 0.50203998,  0.13104815],\n",
-      "       [ 0.46819426,  0.23665964],\n",
-      "       [ 0.43149302,  0.33449321],\n",
-      "       [ 0.3945875 ,  0.42382454],\n",
-      "       [ 0.46072474, -0.552298  ],\n",
-      "       [ 0.46312394, -0.47623446],\n",
-      "       [ 0.46508563, -0.39184904],\n",
-      "       [ 0.46543342, -0.29743107],\n",
-      "       [ 0.46169148, -0.19162379],\n",
-      "       [ 0.44991531, -0.07524281],\n",
-      "       [ 0.42680477,  0.04646023],\n",
-      "       [ 0.39339841,  0.16512372],\n",
-      "       [ 0.35472166,  0.27440954],\n",
-      "       [ 0.31566035,  0.37238804],\n",
-      "       [ 0.27884834,  0.45990214],\n",
-      "       [ 0.35433802, -0.55320829],\n",
-      "       [ 0.35356863, -0.47726693],\n",
-      "       [ 0.35302519, -0.39249966],\n",
-      "       [ 0.35212514, -0.29618816],\n",
-      "       [ 0.34852334, -0.18473489],\n",
-      "       [ 0.33578884, -0.05660441],\n",
-      "       [ 0.30671054,  0.07901422],\n",
-      "       [ 0.26584619,  0.20515544],\n",
-      "       [ 0.22400385,  0.3146805 ],\n",
-      "       [ 0.18622396,  0.40959294],\n",
-      "       [ 0.15312727,  0.49332251],\n",
-      "       [ 0.23739015, -0.54971459],\n",
-      "       [ 0.23087867, -0.47399256],\n",
-      "       [ 0.22442588, -0.38950352],\n",
-      "       [ 0.21830314, -0.29302227],\n",
-      "       [ 0.21224924, -0.17815385],\n",
-      "       [ 0.19863225, -0.03383159],\n",
-      "       [ 0.15811807,  0.1216322 ],\n",
-      "       [ 0.11272857,  0.24831262],\n",
-      "       [ 0.07611139,  0.35195516],\n",
-      "       [ 0.0462995 ,  0.44130356],\n",
-      "       [ 0.02114958,  0.52065869],\n",
-      "       [ 0.11101004, -0.53947865],\n",
-      "       [ 0.09584853, -0.46297986],\n",
-      "       [ 0.07872659, -0.37770426],\n",
-      "       [ 0.05893757, -0.2805206 ],\n",
-      "       [ 0.03500201, -0.16486547],\n",
-      "       [ 0.        ,  0.        ],\n",
-      "       [-0.03500201,  0.16486547],\n",
-      "       [-0.05893757,  0.2805206 ],\n",
-      "       [-0.07872659,  0.37770426],\n",
-      "       [-0.09584853,  0.46297986],\n",
-      "       [-0.11101004,  0.53947865],\n",
-      "       [-0.02114958, -0.52065869],\n",
-      "       [-0.0462995 , -0.44130356],\n",
-      "       [-0.07611139, -0.35195516],\n",
-      "       [-0.11272857, -0.24831262],\n",
-      "       [-0.15811807, -0.1216322 ],\n",
-      "       [-0.19863225,  0.03383159],\n",
-      "       [-0.21224924,  0.17815385],\n",
-      "       [-0.21830314,  0.29302227],\n",
-      "       [-0.22442588,  0.38950352],\n",
-      "       [-0.23087867,  0.47399256],\n",
-      "       [-0.23739015,  0.54971459],\n",
-      "       [-0.15312727, -0.49332251],\n",
-      "       [-0.18622396, -0.40959294],\n",
-      "       [-0.22400385, -0.3146805 ],\n",
-      "       [-0.26584619, -0.20515544],\n",
-      "       [-0.30671054, -0.07901422],\n",
-      "       [-0.33578884,  0.05660441],\n",
-      "       [-0.34852334,  0.18473489],\n",
-      "       [-0.35212514,  0.29618816],\n",
-      "       [-0.35302519,  0.39249966],\n",
-      "       [-0.35356863,  0.47726693],\n",
-      "       [-0.35433802,  0.55320829],\n",
-      "       [-0.27884834, -0.45990214],\n",
-      "       [-0.31566035, -0.37238804],\n",
-      "       [-0.35472166, -0.27440954],\n",
-      "       [-0.39339841, -0.16512372],\n",
-      "       [-0.42680477, -0.04646023],\n",
-      "       [-0.44991531,  0.07524281],\n",
-      "       [-0.46169148,  0.19162379],\n",
-      "       [-0.46543342,  0.29743107],\n",
-      "       [-0.46508563,  0.39184904],\n",
-      "       [-0.46312394,  0.47623446],\n",
-      "       [-0.46072474,  0.552298  ],\n",
-      "       [-0.3945875 , -0.42382454],\n",
-      "       [-0.43149302, -0.33449321],\n",
-      "       [-0.46819426, -0.23665964],\n",
-      "       [-0.50203998, -0.13104815],\n",
-      "       [-0.52983592, -0.02027455],\n",
-      "       [-0.54923128,  0.09122931],\n",
-      "       [-0.56003098,  0.19866231],\n",
-      "       [-0.56403146,  0.29870151],\n",
-      "       [-0.56365885,  0.39013744],\n",
-      "       [-0.56089008,  0.47323704],\n",
-      "       [-0.55699033,  0.54888684],\n",
-      "       [-0.49938556, -0.38784218],\n",
-      "       [-0.53442151, -0.29860742],\n",
-      "       [-0.56769437, -0.20292033],\n",
-      "       [-0.59716429, -0.10203609],\n",
-      "       [-0.62087922,  0.0017213 ],\n",
-      "       [-0.63766224,  0.10527277],\n",
-      "       [-0.6475545 ,  0.2055496 ],\n",
-      "       [-0.65165847,  0.3002847 ],\n",
-      "       [-0.65154284,  0.3883424 ],\n",
-      "       [-0.64868502,  0.46953412],\n",
-      "       [-0.64421232,  0.54423054]], dtype=float64)\n",
-      "0.651658467012907\n"
-     ]
-    }
-   ],
-   "execution_count": 2
-  },
-  {
-   "metadata": {},
-   "cell_type": "code",
-   "outputs": [],
-   "execution_count": null,
-   "source": "",
-   "id": "730f63354e47eb46"
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 2
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython2",
-   "version": "2.7.6"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

From 7801ca26db9f3a5cdd3745c288619f9dd96a4663 Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Fri, 13 Feb 2026 10:28:42 +0000
Subject: [PATCH 12/19] changed test__deflection_yx_2d_from to not be
 tautological

---
 test_autogalaxy/profiles/mass/dark/test_nfw.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/test_autogalaxy/profiles/mass/dark/test_nfw.py b/test_autogalaxy/profiles/mass/dark/test_nfw.py
index 6f3ea4396..bed8e5632 100644
--- a/test_autogalaxy/profiles/mass/dark/test_nfw.py
+++ b/test_autogalaxy/profiles/mass/dark/test_nfw.py
@@ -184,14 +184,14 @@ def test__deflections_2d__numerical_precision_of_csv_compared_to_integral():
 
 
 def test__deflections_yx_2d_from():
-    nfw = ag.mp.NFW(centre=(0.1, -0.1), ell_comps=(0.3, 0.1), kappa_s=1.0, scale_radius=1.0)
+    nfw = ag.mp.NFW(centre=(0.0, 0.0), kappa_s=1.0, scale_radius=1.0)
 
     deflections = nfw.deflections_yx_2d_from(grid=ag.Grid2DIrregular([[1.0, 0.0]]))
-    deflections_via_analytical = nfw.deflections_2d_via_analytic_from(
+    deflections_via_integral = nfw.deflections_2d_via_cse_from(
         grid=ag.Grid2DIrregular([[1.0, 0.0]])
     )
 
-    assert deflections == pytest.approx(deflections_via_analytical.array, 1.0e-4)
+    assert deflections == pytest.approx(deflections_via_integral.array, 1.0e-4)
 
 
 def test__convergence_2d_via_mge_from():

From ef54c66338e45caa7607c7b91bbafbd0910e27da Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Fri, 13 Feb 2026 10:40:37 +0000
Subject: [PATCH 13/19] prevented intermediate nans in capital_F_from

---
 autogalaxy/profiles/mass/dark/nfw_hk24_util.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/autogalaxy/profiles/mass/dark/nfw_hk24_util.py b/autogalaxy/profiles/mass/dark/nfw_hk24_util.py
index e8f214262..1b8b82c27 100644
--- a/autogalaxy/profiles/mass/dark/nfw_hk24_util.py
+++ b/autogalaxy/profiles/mass/dark/nfw_hk24_util.py
@@ -41,11 +41,16 @@ def capital_F_from(chi, xp=np):
     less = chi < 1 - eps
     greater = chi > 1 + eps
 
-    root_min = xp.sqrt(1 - chi**2)
-    root_plus = xp.sqrt(chi**2 - 1)
+    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)
+
+    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_less = xp.arctanh(root_min) / root_min
-    F_greater = xp.arctan(root_plus) / root_plus
     F_equal = xp.ones_like(chi)
 
     return xp.where(less, F_less, xp.where(greater, F_greater, F_equal))

From 8552d2b2a42a55b3f0cb7956db9b6441d2d14bad Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Fri, 13 Feb 2026 10:45:26 +0000
Subject: [PATCH 14/19] prevented intermediate nans in kappa_from

---
 autogalaxy/profiles/mass/dark/nfw_hk24_util.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/autogalaxy/profiles/mass/dark/nfw_hk24_util.py b/autogalaxy/profiles/mass/dark/nfw_hk24_util.py
index 1b8b82c27..667240baa 100644
--- a/autogalaxy/profiles/mass/dark/nfw_hk24_util.py
+++ b/autogalaxy/profiles/mass/dark/nfw_hk24_util.py
@@ -22,6 +22,7 @@ def semi_major_axis_from(x1, x2, e, xp=np):
     """
     return xp.sqrt(x1**2 + x2**2 / (1 - e**2))
 
+
 def capital_F_from(chi, xp=np):
     """
     Equation 16 from Heyrovský & Karamazov.
@@ -55,6 +56,7 @@ def capital_F_from(chi, xp=np):
 
     return xp.where(less, F_less, xp.where(greater, F_greater, F_equal))
 
+
 def kappa_from(k_s, a, xp=np):
     """
     Equation 21 from Heyrovský & Karamazov.
@@ -71,8 +73,8 @@ def kappa_from(k_s, a, xp=np):
     Convergence as a function of a
     """
     F = capital_F_from(a, xp=xp)
-    kappa = 2 * k_s * (1 - F) / (a**2 - 1)
-    return xp.where(xp.abs(a - 1) < 1e-12, 2/3 * k_s, kappa)
+    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, x2, e, xp=np):

From cf747b350b64e96f6a72cf304e8d04188c8f839a Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Fri, 13 Feb 2026 11:20:51 +0000
Subject: [PATCH 15/19] avoid nans fixed in deflections_2d_via_analytic_from

---
 autogalaxy/profiles/mass/dark/nfw.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py
index 7baa4dd50..4848be18a 100644
--- a/autogalaxy/profiles/mass/dark/nfw.py
+++ b/autogalaxy/profiles/mass/dark/nfw.py
@@ -70,10 +70,12 @@ def deflections_2d_via_analytic_from(
         x1 = grid.array[:, 1] / self.scale_radius
         x2 = grid.array[:, 0] / self.scale_radius
 
-        # Avoid nans due to x=0
+        # Avoid nans due to x=0 by perturbing very small coordinates away from exactly zero
         eps = xp.array(1e-12)
-        x1 = xp.where(xp.abs(x1) < eps, xp.sign(x1) * eps, x1)
-        x2 = xp.where(xp.abs(x2) < eps, xp.sign(x2) * eps, x2)
+        sign_x1_safe = xp.where(xp.sign(x1) == 0, 1.0, xp.sign(x1))
+        sign_x2_safe = xp.where(xp.sign(x2) == 0, 1.0, xp.sign(x2))
+        x1 = xp.where(xp.abs(x1) < eps, sign_x1_safe * eps, x1)
+        x2 = xp.where(xp.abs(x2) < eps, sign_x2_safe * eps, x2)
 
         prefactor = 4 * self.kappa_s * xp.sqrt(1 - e_hk24**2) / (
             ((x1 - e_hk24)**2 + x2**2) * ((x1 + e_hk24)**2 + x2**2)

From e19067ed298066847cdc1aae75cf88d23912186a Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Fri, 13 Feb 2026 11:22:15 +0000
Subject: [PATCH 16/19] pass **kwargs in deflections_yx_2d_from

---
 autogalaxy/profiles/mass/dark/nfw.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py
index 4848be18a..f93f22f3c 100644
--- a/autogalaxy/profiles/mass/dark/nfw.py
+++ b/autogalaxy/profiles/mass/dark/nfw.py
@@ -42,8 +42,8 @@ def __init__(
         )
         super(MassProfileCSE, self).__init__()
 
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np):
-        return self.deflections_2d_via_analytic_from(grid=grid, xp=xp)
+    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

From 19595b856334d0e1a72280c85ad2e9e17fbbc516 Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Fri, 13 Feb 2026 12:01:53 +0000
Subject: [PATCH 17/19] set mask before calculating prefactor and f_1 to avoid
 nans

---
 autogalaxy/profiles/mass/dark/nfw.py | 31 +++++++++++-----------------
 1 file changed, 12 insertions(+), 19 deletions(-)

diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py
index f93f22f3c..01e8c8c30 100644
--- a/autogalaxy/profiles/mass/dark/nfw.py
+++ b/autogalaxy/profiles/mass/dark/nfw.py
@@ -70,20 +70,19 @@ def deflections_2d_via_analytic_from(
         x1 = grid.array[:, 1] / self.scale_radius
         x2 = grid.array[:, 0] / self.scale_radius
 
-        # Avoid nans due to x=0 by perturbing very small coordinates away from exactly zero
-        eps = xp.array(1e-12)
-        sign_x1_safe = xp.where(xp.sign(x1) == 0, 1.0, xp.sign(x1))
-        sign_x2_safe = xp.where(xp.sign(x2) == 0, 1.0, xp.sign(x2))
-        x1 = xp.where(xp.abs(x1) < eps, sign_x1_safe * eps, x1)
-        x2 = xp.where(xp.abs(x2) < eps, sign_x2_safe * eps, x2)
-
-        prefactor = 4 * self.kappa_s * xp.sqrt(1 - e_hk24**2) / (
-            ((x1 - e_hk24)**2 + x2**2) * ((x1 + e_hk24)**2 + x2**2)
-        )
+        r2 = x1 ** 2 + x2 ** 2
+
+        # Avoid nans
+
+        mask = r2 > 1e-24
 
-        f1 = nfw_hk24_util.small_f_1(x1, x2, e_hk24, xp=xp)
-        f2 = nfw_hk24_util.small_f_2(x1, x2, e_hk24, xp=xp)
-        f3 = nfw_hk24_util.small_f_3(x1, x2, e_hk24, xp=xp)
+        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
@@ -95,12 +94,6 @@ def deflections_2d_via_analytic_from(
                         + x1 * (x1**2 + x2**2 - e_hk24**2) * f3)
         deflection_y *= prefactor
 
-        # prevent nans at the centre
-        r2 = x1 ** 2 + x2 ** 2
-        mask0 = r2 < 1e-24
-        deflection_x = xp.where(mask0, 0.0, deflection_x)
-        deflection_y = xp.where(mask0, 0.0, deflection_y)
-
         return self.rotated_grid_from_reference_frame_from(
             xp.multiply(self.scale_radius, xp.vstack((deflection_y, deflection_x)).T),
             xp=xp,

From 4465e5722bec2b1c9e20e22dc19c46666f349af2 Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Fri, 13 Feb 2026 12:13:24 +0000
Subject: [PATCH 18/19] extended deflections_via_analytical_from test

---
 .../profiles/mass/dark/test_nfw.py            | 32 +++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/test_autogalaxy/profiles/mass/dark/test_nfw.py b/test_autogalaxy/profiles/mass/dark/test_nfw.py
index bed8e5632..128cdb5a9 100644
--- a/test_autogalaxy/profiles/mass/dark/test_nfw.py
+++ b/test_autogalaxy/profiles/mass/dark/test_nfw.py
@@ -35,6 +35,38 @@ def test__deflections_via_analytical_from():
     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_analytical = nfw.deflections_2d_via_analytical_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_analytical == 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_analytical = 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_analytical == 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)

From d1d13bd2102fa0d53624fad9f87cfeff8a3264ed Mon Sep 17 00:00:00 2001
From: Niek Wielders <c4072114@newcastle.ac.uk>
Date: Fri, 13 Feb 2026 12:19:39 +0000
Subject: [PATCH 19/19] fix typo

---
 test_autogalaxy/profiles/mass/dark/test_nfw.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/test_autogalaxy/profiles/mass/dark/test_nfw.py b/test_autogalaxy/profiles/mass/dark/test_nfw.py
index 128cdb5a9..2fc4d4a06 100644
--- a/test_autogalaxy/profiles/mass/dark/test_nfw.py
+++ b/test_autogalaxy/profiles/mass/dark/test_nfw.py
@@ -42,14 +42,14 @@ def test__deflections_via_analytical_from():
         scale_radius=1.0,
     )
 
-    deflections_via_analytical = nfw.deflections_2d_via_analytical_from(
+    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_analytical == pytest.approx(deflections_via_cse.array, 1.0e-4)
+    assert deflections_via_analytic == pytest.approx(deflections_via_cse.array, 1.0e-4)
 
     nfw = ag.mp.NFW(
         centre=(0.3, -0.2),
@@ -58,14 +58,14 @@ def test__deflections_via_analytical_from():
         scale_radius=4.0,
     )
 
-    deflections_via_analytical = nfw.deflections_2d_via_analytic_from(
+    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_analytical == pytest.approx(deflections_via_cse.array, 1.0e-4)
+    assert deflections_via_analytic == pytest.approx(deflections_via_cse.array, 1.0e-4)
 
 
 def test__deflections_via_integral_from():