Source code for caustics.lenses.func.epl

from ...utils import translate_rotate, derotate
from ...backend_obj import backend


def _r_omega(z, t, q, n_iter, chunk_size=None):
    """
    Iteratively compute the omega term given in Tessore et al. 2015 equation 23.
    This is done using the approximation given in Tessore et al. 2015 equation
    29. Note that in Tessore et al. 2015 the omega term is independent of
    radius, our "r_omega" term includes an r-factor for numerical reasons.
    """
    # constants
    f = (1.0 - q) / (1.0 + q)
    phi = z / backend.conj(z) * -f

    # defaults to do all the iterations in parallel
    chunk_size = n_iter if chunk_size is None else chunk_size

    if chunk_size == 1:
        # first term in series
        omega_i = z
        part_sum = omega_i
        for i in range(1, n_iter):
            factor = (2.0 * i - (2.0 - t)) / (2.0 * i + (2.0 - t))  # fmt: skip
            omega_i = factor * phi * omega_i  # fmt: skip
            part_sum = part_sum + omega_i  # fmt: skip

        return part_sum
    else:
        i = backend.arange(1, n_iter, device=backend.device(z), dtype=z.real.dtype)
        factor = (2.0 * i - (2.0 - t)) / (2.0 * i + (2.0 - t))
        factor = backend.view(factor, (-1, *([1] * z.ndim)))
        phi_expanded = phi[None]
        cumprod_res = None
        n_chunks = n_iter // chunk_size
        for i in range(n_chunks):
            rec = (
                factor[i * chunk_size : max((i + 1) * chunk_size, n_iter)]
                * phi_expanded
            )
            if cumprod_res is None:
                cumprod = backend.cumprod(rec, dim=0)
                cumprod_res = backend.sum(cumprod, dim=0)
            else:
                cumprod = backend.cumprod(
                    backend.concatenate((cumprod[-1][None], rec)), dim=0
                )[1:]
                cumprod_res = cumprod_res + backend.sum(cumprod, dim=0)
        return z * (1 + cumprod_res)


[docs] def reduced_deflection_angle_epl( x0, y0, q, phi, Rein, t, x, y, n_iter, chunk_size=None ): """ Calculate the reduced deflection angle. Given in Tessore et al. 2015 equation 13. Parameters ---------- x0: ArrayLike The x-coordinate of the lens center. *Unit: arcsec* y0: ArrayLike The y-coordinate of the lens center. *Unit: arcsec* q: ArrayLike The axis ratio of the lens. Semi-minor over semi-major axis lengths. *Unit: unitless* phi: ArrayLike The orientation angle of the lens (position angle). *Unit: radians* Rein: ArrayLike The Einstein radius of the lens. *Unit: arcsec* t: ArrayLike Power law slope (`gamma-1`) of the lens. If not provided, it is considered as a free parameter. *Unit: unitless* x: ArrayLike The x-coordinate of the lens. *Unit: arcsec* y: ArrayLike The y-coordinate of the lens. *Unit: arcsec* n_iter: int Number of iterations for the iterative solver. *Unit: number* chunk_size: int Number of iterations to do in parallel for the iterative solver. *Unit: number* Returns -------- x_component: ArrayLike The x-component of the deflection angle. *Unit: arcsec* y_component: ArrayLike The y-component of the deflection angle. *Unit: arcsec* """ x, y = translate_rotate(x, y, x0, y0, phi) # follow Tessore et al 2015 (eq. 5) z = q * x + y * 1j r = backend.abs(z) # Tessore et al 2015 (eq. 23) r_omega = _r_omega(z, t, q, n_iter, chunk_size) # Tessore et al 2015 (eq. 13) alpha_c = 2.0 / (1.0 + q) * (Rein * backend.sqrt(q) / r) ** t * r_omega # fmt: skip alpha_real = backend.nan_to_num(alpha_c.real, posinf=10**10, neginf=-(10**10)) alpha_imag = backend.nan_to_num(alpha_c.imag, posinf=10**10, neginf=-(10**10)) return derotate(alpha_real, alpha_imag, phi)
[docs] def potential_epl(x0, y0, q, phi, Rein, t, x, y, n_iter, chunk_size): """ Calculate the potential for the EPL as defined in Tessore et al. 2015 equation 15. Parameters ---------- x0: ArrayLike The x-coordinate of the lens center. *Unit: arcsec* y0: ArrayLike The y-coordinate of the lens center. *Unit: arcsec* q: ArrayLike The axis ratio of the lens. Semi-minor over semi-major axis lengths. *Unit: unitless* phi: ArrayLike The orientation angle of the lens (position angle). *Unit: radians* Rein: ArrayLike The Einstein radius of the lens. *Unit: arcsec* t: ArrayLike Power law slope (`gamma-1`) of the lens. If not provided, it is considered as a free parameter. *Unit: unitless* x: ArrayLike The x-coordinate of the lens. *Unit: arcsec* y: ArrayLike The y-coordinate of the lens. *Unit: arcsec* n_iter: int Number of iterations for the iterative solver. *Unit: number* chunk_size: int Number of iterations to do in parallel for the iterative solver. *Unit: number* Returns -------- x_component: ArrayLike The x-component of the deflection angle. *Unit: arcsec* y_component: ArrayLike The y-component of the deflection angle. *Unit: arcsec* """ ax, ay = reduced_deflection_angle_epl( x0, y0, q, phi, Rein, t, x, y, n_iter, chunk_size ) ax, ay = derotate(ax, ay, -phi) x, y = translate_rotate(x, y, x0, y0, phi) return (x * ax + y * ay) / (2 - t) # fmt: skip
[docs] def convergence_epl(x0, y0, q, phi, Rein, t, x, y, s=0.0): """ Calculate the reduced deflection angle. See Tessore et al. 2015 equation 2. Parameters ---------- x0: ArrayLike The x-coordinate of the lens center. *Unit: arcsec* y0: ArrayLike The y-coordinate of the lens center. *Unit: arcsec* q: ArrayLike The axis ratio of the lens. Semi-minor over semi-major axis lengths. *Unit: unitless* phi: ArrayLike The orientation angle of the lens (position angle). *Unit: radians* Rein: ArrayLike The Einstein radius of the lens. *Unit: arcsec* t: ArrayLike Power law slope (`gamma-1`) of the lens. If not provided, it is considered as a free parameter. *Unit: unitless* x: ArrayLike The x-coordinate of the lens. *Unit: arcsec* y: ArrayLike The y-coordinate of the lens. *Unit: arcsec* s: float The core radius of the lens (defaults to 0.0). *Unit: arcsec* Returns -------- x_component: ArrayLike The x-component of the deflection angle. *Unit: arcsec* y_component: ArrayLike The y-component of the deflection angle. *Unit: arcsec* """ x, y = translate_rotate(x, y, x0, y0, phi) psi = backend.sqrt(q ** 2 * x ** 2 + y ** 2 + s ** 2) # fmt: skip return (2 - t) / 2 * (Rein * backend.sqrt(q) / psi) ** t # fmt: skip