# mypy: disable-error-code="operator,union-attr,dict-item"
from typing import Optional, Union, Annotated, Callable
from caskade import forward, Param
from ..backend_obj import backend, ArrayLike
from .base import ThinLens, CosmologyType, NameType, ZType
from .func import (
physical_deflection_angle_enclosed_mass,
convergence_enclosed_mass,
reduced_from_physical_deflection_angle,
)
__all__ = ("EnclosedMass",)
[docs]
class EnclosedMass(ThinLens):
"""
A class for representing a lens with an enclosed mass profile. This generic
lens profile can represent any lens with a mass distribution that can be
described by a function that returns the enclosed mass as a function of
radius.
"""
_null_params = {
"x0": 0.0,
"y0": 0.0,
"q": 1.0,
"phi": 0.0,
"p": 1.0,
}
def __init__(
self,
cosmology: CosmologyType,
enclosed_mass: Callable,
z_l: ZType = None,
z_s: ZType = None,
x0: Annotated[
Optional[Union[ArrayLike, float]],
"The x-coordinate of the lens center",
True,
] = None,
y0: Annotated[
Optional[Union[ArrayLike, float]],
"The y-coordinate of the lens center",
True,
] = None,
q: Annotated[
Optional[Union[ArrayLike, float]], "The axis ratio of the lens", True
] = None,
phi: Annotated[
Optional[Union[ArrayLike, float]], "The position angle of the lens", True
] = None,
p: Annotated[
Optional[Union[ArrayLike, list[float]]],
"parameters for the enclosed mass function",
True,
] = None,
s: Annotated[
float, "Softening parameter to prevent numerical instabilities"
] = 0.0,
name: NameType = None,
**kwargs,
):
"""
Initialize an enclosed mass lens.
Parameters
----------
name : str
The name of the lens.
cosmology : Cosmology
The cosmology object that describes the Universe.
enclosed_mass : callable
A function that takes a radius and a set of parameters and returns the enclosed mass. Should be of the form ``enclosed_mass(r, p) -> M`` and returns units of Msun.
*Unit: Msun*
z_l : float
The redshift of the lens.
*Unit: unitless*
z_s : float
The redshift of the source.
*Unit: unitless*
x0 : float or ArrayLike, optional
The x-coordinate of the lens center.
*Unit: arcsec*
y0 : float or ArrayLike, optional
The y-coordinate of the lens center.
*Unit: arcsec*
q : float or ArrayLike, optional
The axis ratio of the lens. ratio of semi-minor to semi-major axis (b/a).
*Unit: unitless*
phi : float or ArrayLike, optional
The position angle of the lens.
*Unit: radians*
p : list[float] or ArrayLike, optional
The parameters for the enclosed mass function.
*Unit: user-defined*
s : float, optional
Softening parameter to prevent numerical instabilities.
*Unit: arcsec*
"""
super().__init__(cosmology, z_l, name=name, z_s=z_s, **kwargs)
self.enclosed_mass = enclosed_mass
self.x0 = Param("x0", x0, shape=(), units="arcsec")
self.y0 = Param("y0", y0, shape=(), units="arcsec")
self.q = Param("q", q, shape=(), units="unitless", valid=(0, 1))
self.phi = Param(
"phi", phi, shape=(), units="radians", valid=(0, backend.pi), cyclic=True
)
self.p = Param("p", p, units="user-defined")
self.s = s
[docs]
@forward
def physical_deflection_angle(
self,
x: ArrayLike,
y: ArrayLike,
x0: Annotated[ArrayLike, "Param"],
y0: Annotated[ArrayLike, "Param"],
q: Annotated[ArrayLike, "Param"],
phi: Annotated[ArrayLike, "Param"],
p: Annotated[ArrayLike, "Param"],
) -> tuple[ArrayLike, ArrayLike]:
"""
Calculate the physical deflection angle of the lens at a given position.
Parameters
----------
x: ArrayLike
The x-coordinate on the lens plane.
*Unit: arcsec*
y: ArrayLike
The y-coordinate on the lens plane.
*Unit: arcsec*
Returns
-------
The physical deflection angle at the given position. [ArrayLike, ArrayLike]
*Unit: arcsec*
"""
return physical_deflection_angle_enclosed_mass(
x0, y0, q, phi, lambda r: self.enclosed_mass(r, p), x, y, self.s
)
[docs]
@forward
def reduced_deflection_angle(self, x, y, z_s, z_l):
d_s = self.cosmology.angular_diameter_distance(z_s)
d_ls = self.cosmology.angular_diameter_distance_z1z2(z_l, z_s)
deflection_angle_x, deflection_angle_y = self.physical_deflection_angle(x, y)
return reduced_from_physical_deflection_angle(
deflection_angle_x, deflection_angle_y, d_s, d_ls
)
[docs]
@forward
def potential(
self,
x: ArrayLike,
y: ArrayLike,
*args,
**kwargs,
) -> ArrayLike:
raise NotImplementedError(
"Potential is not implemented for enclosed mass profiles."
)
[docs]
@forward
def convergence(
self,
x: ArrayLike,
y: ArrayLike,
z_s: Annotated[ArrayLike, "Param"],
z_l: Annotated[ArrayLike, "Param"],
x0: Annotated[ArrayLike, "Param"],
y0: Annotated[ArrayLike, "Param"],
q: Annotated[ArrayLike, "Param"],
phi: Annotated[ArrayLike, "Param"],
p: Annotated[ArrayLike, "Param"],
) -> ArrayLike:
"""
Calculate the dimensionless convergence of the lens at a given position.
Parameters
----------
x: ArrayLike
The x-coordinate on the lens plane.
*Unit: arcsec*
y: ArrayLike
The y-coordinate on the lens plane.
*Unit: arcsec*
Returns
-------
The dimensionless convergence at the given position. [ArrayLike]
*Unit: unitless*
"""
csd = self.cosmology.critical_surface_density(z_l, z_s)
return convergence_enclosed_mass(
x0,
y0,
q,
phi,
lambda r: self.enclosed_mass(r, p),
x,
y,
csd,
self.s,
)
