diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py index ccf9a9eaf..01e8c8c30 100644 --- a/autogalaxy/profiles/mass/dark/nfw.py +++ b/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,8 +326,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) / (1 + e_autolens) # Define dimensionless length coords @@ -280,17 +335,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 +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): diff --git a/autogalaxy/profiles/mass/dark/nfw_hk24_util.py b/autogalaxy/profiles/mass/dark/nfw_hk24_util.py index 0d24b679d..667240baa 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,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) ) diff --git a/test_autogalaxy/profiles/mass/dark/test_nfw.py b/test_autogalaxy/profiles/mass/dark/test_nfw.py index 9260afcbd..2fc4d4a06 100644 --- a/test_autogalaxy/profiles/mass/dark/test_nfw.py +++ b/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)