Source code for autogalaxy.profiles.mass.dark.gnfw

import inspect
import numpy as np
from scipy import LowLevelCallable
from scipy import special
from scipy.integrate import quad
from typing import Tuple

import autoarray as aa

from autogalaxy.profiles.mass.dark.abstract import AbstractgNFW


def jit_integrand(integrand_function):
    from numba import cfunc
    from numba.types import intc, CPointer, float64

    jitted_function = aa.util.numba.jit(nopython=True, cache=True)(integrand_function)
    no_args = len(inspect.getfullargspec(integrand_function).args)

    wrapped = None

    if no_args == 4:
        # noinspection PyUnusedLocal
        def wrapped(n, xx):
            return jitted_function(xx[0], xx[1], xx[2], xx[3])

    elif no_args == 5:
        # noinspection PyUnusedLocal
        def wrapped(n, xx):
            return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4])

    elif no_args == 6:
        # noinspection PyUnusedLocal
        def wrapped(n, xx):
            return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4], xx[5])

    elif no_args == 7:
        # noinspection PyUnusedLocal
        def wrapped(n, xx):
            return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4], xx[5], xx[6])

    elif no_args == 8:
        # noinspection PyUnusedLocal
        def wrapped(n, xx):
            return jitted_function(
                xx[0], xx[1], xx[2], xx[3], xx[4], xx[5], xx[6], xx[7]
            )

    elif no_args == 9:
        # noinspection PyUnusedLocal
        def wrapped(n, xx):
            return jitted_function(
                xx[0], xx[1], xx[2], xx[3], xx[4], xx[5], xx[6], xx[7], xx[8]
            )

    elif no_args == 10:
        # noinspection PyUnusedLocal
        def wrapped(n, xx):
            return jitted_function(
                xx[0], xx[1], xx[2], xx[3], xx[4], xx[5], xx[6], xx[7], xx[8], xx[9]
            )

    elif no_args == 11:
        # noinspection PyUnusedLocal
        def wrapped(n, xx):
            return jitted_function(
                xx[0],
                xx[1],
                xx[2],
                xx[3],
                xx[4],
                xx[5],
                xx[6],
                xx[7],
                xx[8],
                xx[9],
                xx[10],
            )

    cf = cfunc(float64(intc, CPointer(float64)))

    return LowLevelCallable(cf(wrapped).ctypes)


[docs] class gNFW(AbstractgNFW): def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): return self.deflections_2d_via_mge_from(grid=grid, **kwargs) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform @aa.grid_dec.relocate_to_radial_minimum def deflections_2d_via_mge_from(self, grid: aa.type.Grid2DLike, **kwargs): return self._deflections_2d_via_mge_from( grid=grid, sigmas_factor=self.axis_ratio )
[docs] @aa.grid_dec.to_vector_yx @aa.grid_dec.transform @aa.grid_dec.relocate_to_radial_minimum def deflections_2d_via_integral_from( self, grid: aa.type.Grid2DLike, tabulate_bins=1000, **kwargs ): """ Calculate the deflection angles at a given set of arc-second gridded coordinates. Parameters ---------- grid The grid of (y,x) arc-second coordinates the deflection angles are computed on. tabulate_bins The number of bins to tabulate the inner integral of this profile. """ @jit_integrand def surface_density_integrand(x, kappa_radius, scale_radius, inner_slope): return ( (3 - inner_slope) * (x + kappa_radius / scale_radius) ** (inner_slope - 4) * (1 - np.sqrt(1 - x * x)) ) def calculate_deflection_component(npow, yx_index): deflection_grid = np.zeros(grid.shape[0]) for i in range(grid.shape[0]): deflection_grid[i] = ( 2.0 * self.kappa_s * self.axis_ratio * grid[i, yx_index] * quad( self.deflection_func, a=0.0, b=1.0, args=( grid[i, 0], grid[i, 1], npow, self.axis_ratio, minimum_log_eta, maximum_log_eta, tabulate_bins, surface_density_integral, ), epsrel=gNFW.epsrel, )[0] ) return deflection_grid ( eta_min, eta_max, minimum_log_eta, maximum_log_eta, bin_size, ) = self.tabulate_integral(grid, tabulate_bins) surface_density_integral = np.zeros((tabulate_bins,)) for i in range(tabulate_bins): eta = 10.0 ** (minimum_log_eta + (i - 1) * bin_size) integral = quad( surface_density_integrand, a=0.0, b=1.0, args=(eta, self.scale_radius, self.inner_slope), epsrel=gNFW.epsrel, )[0] surface_density_integral[i] = ( (eta / self.scale_radius) ** (1 - self.inner_slope) ) * (((1 + eta / self.scale_radius) ** (self.inner_slope - 3)) + integral) deflection_y = calculate_deflection_component(npow=1.0, yx_index=0) deflection_x = calculate_deflection_component(npow=0.0, yx_index=1) return self.rotated_grid_from_reference_frame_from( np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T) )
@staticmethod def deflection_func( u, y, x, npow, axis_ratio, minimum_log_eta, maximum_log_eta, tabulate_bins, surface_density_integral, ): _eta_u = np.sqrt((u * ((x**2) + (y**2 / (1 - (1 - axis_ratio**2) * u))))) bin_size = (maximum_log_eta - minimum_log_eta) / (tabulate_bins - 1) i = 1 + int((np.log10(_eta_u) - minimum_log_eta) / bin_size) r1 = 10.0 ** (minimum_log_eta + (i - 1) * bin_size) r2 = r1 * 10.0**bin_size kap = surface_density_integral[i] + ( surface_density_integral[i + 1] - surface_density_integral[i] ) * (_eta_u - r1) / (r2 - r1) return kap / (1.0 - (1.0 - axis_ratio**2) * u) ** (npow + 0.5) def convergence_func(self, grid_radius: float) -> float: def integral_y(y, eta): return (y + eta) ** (self.inner_slope - 4) * (1 - np.sqrt(1 - y**2)) grid_radius = (1.0 / self.scale_radius) * grid_radius for index in range(grid_radius.shape[0]): integral_y_value = quad( integral_y, a=0.0, b=1.0, args=grid_radius[index], epsrel=gNFW.epsrel, )[0] grid_radius[index] = ( 2.0 * self.kappa_s * (grid_radius[index] ** (1 - self.inner_slope)) * ( (1 + grid_radius[index]) ** (self.inner_slope - 3) + ((3 - self.inner_slope) * integral_y_value) ) ) return grid_radius
[docs] @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform @aa.grid_dec.relocate_to_radial_minimum def potential_2d_from(self, grid: aa.type.Grid2DLike, tabulate_bins=1000, **kwargs): """ Calculate the potential at a given set of arc-second gridded coordinates. Parameters ---------- grid The grid of (y,x) arc-second coordinates the deflection angles are computed on. tabulate_bins The number of bins to tabulate the inner integral of this profile. """ @jit_integrand def deflection_integrand(x, kappa_radius, scale_radius, inner_slope): return (x + kappa_radius / scale_radius) ** (inner_slope - 3) * ( (1 - np.sqrt(1 - x**2)) / x ) ( eta_min, eta_max, minimum_log_eta, maximum_log_eta, bin_size, ) = self.tabulate_integral(grid, tabulate_bins) potential_grid = np.zeros(grid.shape[0]) deflection_integral = np.zeros((tabulate_bins,)) for i in range(tabulate_bins): eta = 10.0 ** (minimum_log_eta + (i - 1) * bin_size) integral = quad( deflection_integrand, a=0.0, b=1.0, args=(eta, self.scale_radius, self.inner_slope), epsrel=gNFW.epsrel, )[0] deflection_integral[i] = ( (eta / self.scale_radius) ** (2 - self.inner_slope) ) * ( (1.0 / (3 - self.inner_slope)) * special.hyp2f1( 3 - self.inner_slope, 3 - self.inner_slope, 4 - self.inner_slope, -(eta / self.scale_radius), ) + integral ) for i in range(grid.shape[0]): potential_grid[i] = (2.0 * self.kappa_s * self.axis_ratio) * quad( self.potential_func, a=0.0, b=1.0, args=( grid[i, 0], grid[i, 1], self.axis_ratio, minimum_log_eta, maximum_log_eta, tabulate_bins, deflection_integral, ), epsrel=gNFW.epsrel, )[0] return potential_grid
@staticmethod def potential_func( u, y, x, axis_ratio, minimum_log_eta, maximum_log_eta, tabulate_bins, potential_integral, ): _eta_u = np.sqrt((u * ((x**2) + (y**2 / (1 - (1 - axis_ratio**2) * u))))) bin_size = (maximum_log_eta - minimum_log_eta) / (tabulate_bins - 1) i = 1 + int((np.log10(_eta_u) - minimum_log_eta) / bin_size) r1 = 10.0 ** (minimum_log_eta + (i - 1) * bin_size) r2 = r1 * 10.0**bin_size angle = potential_integral[i] + ( potential_integral[i + 1] - potential_integral[i] ) * (_eta_u - r1) / (r2 - r1) return _eta_u * (angle / u) / (1.0 - (1.0 - axis_ratio**2) * u) ** 0.5
[docs] class gNFWSph(gNFW): def __init__( self, centre: Tuple[float, float] = (0.0, 0.0), kappa_s: float = 0.05, inner_slope: float = 1.0, scale_radius: float = 1.0, ): """ The spherical NFW profiles, used to fit the dark matter halo of the lens. Parameters ---------- centre The (y,x) arc-second coordinates of the profile centre. kappa_s The overall normalization of the dark matter halo \ (kappa_s = (rho_s * scale_radius)/lensing_critical_density) inner_slope The inner slope of the dark matter halo. scale_radius The arc-second radius where the average density within this radius is 200 times the critical density of \ the Universe.. """ super().__init__( centre=centre, ell_comps=(0.0, 0.0), kappa_s=kappa_s, inner_slope=inner_slope, scale_radius=scale_radius, )
[docs] @aa.grid_dec.to_vector_yx @aa.grid_dec.transform @aa.grid_dec.relocate_to_radial_minimum def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs): """ Calculate the deflection angles at a given set of arc-second gridded coordinates. Parameters ---------- grid The grid of (y,x) arc-second coordinates the deflection angles are computed on. """ eta = np.multiply( 1.0 / self.scale_radius, self.radial_grid_from(grid, **kwargs) ) deflection_grid = np.zeros(grid.shape[0]) for i in range(grid.shape[0]): deflection_grid[i] = np.multiply( 4.0 * self.kappa_s * self.scale_radius, self.deflection_func_sph(eta[i]) ) return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_grid)
@staticmethod def deflection_integrand(y, eta, inner_slope): return (y + eta) ** (inner_slope - 3) * ((1 - np.sqrt(1 - y**2)) / y) def deflection_func_sph(self, eta): integral_y_2 = quad( self.deflection_integrand, a=0.0, b=1.0, args=(eta, self.inner_slope), epsrel=1.49e-6, )[0] return eta ** (2 - self.inner_slope) * ( (1.0 / (3 - self.inner_slope)) * special.hyp2f1( 3 - self.inner_slope, 3 - self.inner_slope, 4 - self.inner_slope, -eta ) + integral_y_2 )