# mypy: disable-error-code="index,dict-item"
from typing import Optional, Annotated, Union
import numpy as np
from caskade import forward, Param
from ..backend_obj import backend, ArrayLike
from ..utils import interp_bicubic
from .base import ThinLens, CosmologyType, NameType, ZType
__all__ = ("PixelatedPotential",)
[docs]
class PixelatedPotential(ThinLens):
_null_params = {
"x0": 0.0,
"y0": 0.0,
"potential_map": np.logspace(0, 1, 100, dtype=np.float32).reshape(10, 10),
}
def __init__(
self,
pixelscale: Annotated[float, "pixelscale", True],
cosmology: CosmologyType,
z_l: ZType = None,
z_s: ZType = None,
x0: Annotated[
Optional[Union[ArrayLike, float]],
"The x-coordinate of the center of the grid",
True,
] = backend.make_array(0.0),
y0: Annotated[
Optional[Union[ArrayLike, float]],
"The y-coordinate of the center of the grid",
True,
] = backend.make_array(0.0),
potential_map: Annotated[
Optional[ArrayLike],
"A 2D tensor representing the potential map",
True,
] = None,
shape: Annotated[
tuple[Optional[int], ...], "The shape of the potential map"
] = (
None,
None,
),
name: NameType = None,
):
"""Strong lensing with user provided kappa map
PixelatedConvergence is a class for strong gravitational lensing with a
user-provided kappa map. It inherits from the ThinLens class.
This class enables the computation of deflection angles and
lensing potential by applying the user-provided kappa map to a
grid using either Fast Fourier Transform (FFT) or a 2D
convolution.
Attributes
----------
name: string
The name of the PixelatedConvergence object.
fov: float
The field of view in arcseconds.
*Unit: arcsec*
cosmology: Cosmology
An instance of the cosmological parameters.
z_l: Optional[ArrayLike]
The redshift of the lens.
*Unit: unitless*
x0: Optional[ArrayLike]
The x-coordinate of the center of the grid.
*Unit: arcsec*
y0: Optional[ArrayLike]
The y-coordinate of the center of the grid.
*Unit: arcsec*
potential_map: Optional[ArrayLike]
A 2D tensor representing the potential map.
*Unit: unitless*
shape: Optional[tuple[int, ...]]
The shape of the potential map.
"""
super().__init__(cosmology, z_l, name=name, z_s=z_s)
if potential_map is not None and potential_map.ndim != 2:
raise ValueError(
f"potential_map must be 2D. Received a {potential_map.ndim}D tensor)"
)
elif shape is not None and len(shape) != 2:
raise ValueError(f"shape must specify a 2D tensor. Received shape={shape}")
self.x0 = Param("x0", x0, shape=(), units="arcsec")
self.y0 = Param("y0", y0, shape=(), units="arcsec")
self.potential_map = Param(
"potential_map", potential_map, shape, units="unitless"
)
self.pixelscale = Param(
"pixelscale", pixelscale, shape=(), units="arcsec/pixel"
)
[docs]
@forward
def reduced_deflection_angle(
self,
x: ArrayLike,
y: ArrayLike,
x0: Annotated[ArrayLike, "Param"],
y0: Annotated[ArrayLike, "Param"],
potential_map: Annotated[ArrayLike, "Param"],
pixelscale: Annotated[ArrayLike, "Param"],
) -> tuple[ArrayLike, ArrayLike]:
"""
Compute the deflection angles at the specified positions using the given convergence map.
Parameters
----------
x: ArrayLike
The x-coordinates of the positions to compute the deflection angles for.
*Unit: arcsec*
y: ArrayLike
The y-coordinates of the positions to compute the deflection angles for.
*Unit: arcsec*
Returns
-------
x_component: ArrayLike
Deflection Angle in the x-direction.
*Unit: arcsec*
y_component: ArrayLike
Deflection Angle in the y-direction.
*Unit: arcsec*
"""
fov_x = potential_map.shape[1] * pixelscale
fov_y = potential_map.shape[0] * pixelscale
return tuple(
alpha.reshape(x.shape) / pixelscale
for alpha in interp_bicubic(
backend.view(x - x0, -1) / fov_x * 2,
backend.view(y - y0, -1) / fov_y * 2,
potential_map,
get_Y=False,
get_dY=True,
get_ddY=False,
)
)
[docs]
@forward
def potential(
self,
x: ArrayLike,
y: ArrayLike,
x0: Annotated[ArrayLike, "Param"],
y0: Annotated[ArrayLike, "Param"],
potential_map: Annotated[ArrayLike, "Param"],
pixelscale: Annotated[ArrayLike, "Param"],
) -> ArrayLike:
"""
Compute the lensing potential at the specified positions using the given convergence map.
Parameters
----------
x: ArrayLike
The x-coordinates of the positions to compute the lensing potential for.
*Unit: arcsec*
y: ArrayLike
The y-coordinates of the positions to compute the lensing potential for.
*Unit: arcsec*
Returns
-------
ArrayLike
The lensing potential at the specified positions.
*Unit: arcsec^2*
"""
fov_x = potential_map.shape[1] * pixelscale
fov_y = potential_map.shape[0] * pixelscale
return interp_bicubic(
backend.view(x - x0, -1) / fov_x * 2,
backend.view(y - y0, -1) / fov_y * 2,
potential_map,
get_Y=True,
get_dY=False,
get_ddY=False,
)[0].reshape(x.shape)
[docs]
@forward
def convergence(
self,
x: ArrayLike,
y: ArrayLike,
x0: Annotated[ArrayLike, "Param"],
y0: Annotated[ArrayLike, "Param"],
potential_map: Annotated[ArrayLike, "Param"],
pixelscale: Annotated[ArrayLike, "Param"],
) -> ArrayLike:
"""
Compute the convergence at the specified positions.
Parameters
----------
x: ArrayLike
The x-coordinates of the positions to compute the convergence for.
*Unit: arcsec*
y: ArrayLike
The y-coordinates of the positions to compute the convergence for.
*Unit: arcsec*
Returns
-------
ArrayLike
The convergence at the specified positions.
*Unit: unitless*
"""
fov_x = potential_map.shape[1] * pixelscale
fov_y = potential_map.shape[0] * pixelscale
_, dY11, dY22 = interp_bicubic(
backend.view(x - x0, -1) / fov_x * 2,
backend.view(y - y0, -1) / fov_y * 2,
potential_map,
get_Y=False,
get_dY=False,
get_ddY=True,
)
return (dY11 + dY22).reshape(x.shape) / (2 * pixelscale**2)
