diff --git a/src/caustics/__init__.py b/src/caustics/__init__.py index 1e1e7788..76c4380e 100644 --- a/src/caustics/__init__.py +++ b/src/caustics/__init__.py @@ -43,6 +43,7 @@ "ThickLens", "EPL", "ExternalShear", + "ExternalShear_angle", "PixelatedConvergence", "Multiplane", "NFW", diff --git a/src/caustics/lenses/external_shear.py b/src/caustics/lenses/external_shear.py index f10995bb..2d5db61e 100644 --- a/src/caustics/lenses/external_shear.py +++ b/src/caustics/lenses/external_shear.py @@ -1,17 +1,21 @@ -# mypy: disable-error-code="dict-item" from typing import Optional, Union, Annotated from torch import Tensor +import torch -from ..utils import translate_rotate from .base import ThinLens, CosmologyType, NameType, ZLType from ..parametrized import unpack +from ..parametrization import convert_params from ..packed import Packed __all__ = ("ExternalShear",) class ExternalShear(ThinLens): + _null_params = { + "gamma_1": 0.01, + "gamma_2": 0.0, + } """ Represents an external shear effect in a gravitational lensing system. @@ -28,64 +32,97 @@ class ExternalShear(ThinLens): *Unit: unitless* - x0, y0: Optional[Union[Tensor, float]] - Coordinates of the shear center in the lens plane. + gamma_1, gamma_2: Optional[Union[Tensor, float]] + Shear components (cartesian parametrization). - *Unit: arcsec* + *Unit: unitless* - gamma_1, gamma_2: Optional[Union[Tensor, float]] - Shear components. + gamma: Optional[Union[Tensor, float]] + Shear magnitude (polar parametrization). *Unit: unitless* + phi: Optional[Union[Tensor, float]] + Shear angle (polar parametrization). + + *Unit: radians* + Notes ------ The shear components gamma_1 and gamma_2 represent an external shear, a gravitational distortion that can be caused by nearby structures outside of the main lens galaxy. """ - _null_params = { - "x0": 0.0, - "y0": 0.0, - "gamma_1": 0.1, - "gamma_2": 0.1, - } - def __init__( self, cosmology: CosmologyType, z_l: ZLType = None, - x0: Annotated[ + gamma_1: Annotated[ Optional[Union[Tensor, float]], - "x-coordinate of the shear center in the lens plane", + "Shear component in the vertical direction", True, ] = None, - y0: Annotated[ + gamma_2: Annotated[ Optional[Union[Tensor, float]], - "y-coordinate of the shear center in the lens plane", + "Shear component in the y=-x direction", True, ] = None, - gamma_1: Annotated[ - Optional[Union[Tensor, float]], "Shear component in the x-direction", True - ] = None, - gamma_2: Annotated[ - Optional[Union[Tensor, float]], "Shear component in the y-direction", True + gamma: Annotated[ + Optional[Union[Tensor, float]], "Shear magnitude", True ] = None, + phi: Annotated[Optional[Union[Tensor, float]], "Shear angle", True] = None, + parametrization: Annotated[ + str, "Parametrization of the shear field", {"cartesian", "polar"} + ] = "cartesian", s: Annotated[ float, "Softening length for the elliptical power-law profile" ] = 0.0, name: NameType = None, ): super().__init__(cosmology, z_l, name=name) - - self.add_param("x0", x0) - self.add_param("y0", y0) - self.add_param("gamma_1", gamma_1) - self.add_param("gamma_2", gamma_2) + if parametrization.lower() == "cartesian": + self.add_param("gamma_1", gamma_1) + self.add_param("gamma_2", gamma_2) + elif parametrization.lower() == "polar": + self.add_param("gamma", gamma) + self.add_param("phi", phi) + else: + raise ValueError( + f"parametrization must be either 'cartesian' or 'polar', got {parametrization}" + ) self.s = s + # We take cartesian as the fiducial parametrization + if parametrization.lower() == "cartesian": + self._convert_params_to_fiducial = ( + lambda self, *args, **kwargs: kwargs + ) # do nothing + elif parametrization.lower() == "polar": + self._convert_params_to_fiducial = self._convert_polar_to_cartesian + + def _convert_polar_to_cartesian(self, *args, **kwargs): + gamma = kwargs.get("gamma") + phi = kwargs.get("phi") + # This breaks if gamma or phi are not provided (or are None) + gamma_1, gamma_2 = self._polar_to_cartesian(gamma, phi) + kwargs["gamma_1"] = gamma_1 + kwargs["gamma_2"] = gamma_2 + return kwargs + + @staticmethod + def _polar_to_cartesian(gamma: Tensor, phi: Tensor) -> tuple[Tensor, Tensor]: + gamma_1 = gamma * torch.cos(2 * phi) + gamma_2 = gamma * torch.sin(2 * phi) + return gamma_1, gamma_2 + + @staticmethod + def _cartesian_to_polar(gamma_1: Tensor, gamma_2: Tensor) -> tuple[Tensor, Tensor]: + gamma = torch.sqrt(gamma_1**2 + gamma_2**2) + phi = 0.5 * torch.atan2(gamma_2, gamma_1) + return gamma, phi @unpack - def reduced_deflection_angle( + @convert_params + def potential( self, x: Tensor, y: Tensor, @@ -93,14 +130,14 @@ def reduced_deflection_angle( *args, params: Optional["Packed"] = None, z_l: Optional[Tensor] = None, - x0: Optional[Tensor] = None, - y0: Optional[Tensor] = None, gamma_1: Optional[Tensor] = None, gamma_2: Optional[Tensor] = None, + gamma: Optional[Tensor] = None, + phi: Optional[Tensor] = None, **kwargs, - ) -> tuple[Tensor, Tensor]: + ) -> Tensor: """ - Calculates the reduced deflection angle. + Calculates the lensing potential. Parameters ---------- @@ -124,30 +161,18 @@ def reduced_deflection_angle( Returns ------- - x_component: Tensor - Deflection Angle in x-direction. - - *Unit: arcsec* - - y_component: Tensor - Deflection Angle in y-direction. + Tensor + The lensing potential. - *Unit: arcsec* + *Unit: arcsec^2* """ - x, y = translate_rotate(x, y, x0, y0) - # Meneghetti eq 3.83 - # TODO, why is it not: - # th = (x**2 + y**2).sqrt() + self.s - # a1 = x/th + x * gamma_1 + y * gamma_2 - # a2 = y/th + x * gamma_2 - y * gamma_1 - a1 = x * gamma_1 + y * gamma_2 - a2 = x * gamma_2 - y * gamma_1 - - return a1, a2 # I'm not sure but I think no derotation necessary + # Equation 5.127 of Meneghetti et al. 2019 + return 0.5 * gamma_1 * (x**2 - y**2) + gamma_2 * x * y @unpack - def potential( + @convert_params + def reduced_deflection_angle( self, x: Tensor, y: Tensor, @@ -155,14 +180,14 @@ def potential( *args, params: Optional["Packed"] = None, z_l: Optional[Tensor] = None, - x0: Optional[Tensor] = None, - y0: Optional[Tensor] = None, gamma_1: Optional[Tensor] = None, gamma_2: Optional[Tensor] = None, + gamma: Optional[Tensor] = None, + phi: Optional[Tensor] = None, **kwargs, - ) -> Tensor: + ) -> tuple[Tensor, Tensor]: """ - Calculates the lensing potential. + Calculates the reduced deflection angle. Parameters ---------- @@ -186,17 +211,24 @@ def potential( Returns ------- - Tensor - The lensing potential. + x_component: Tensor + Deflection Angle in x-direction. - *Unit: arcsec^2* + *Unit: arcsec* + + y_component: Tensor + Deflection Angle in y-direction. + + *Unit: arcsec* """ - ax, ay = self.reduced_deflection_angle(x, y, z_s, params) - x, y = translate_rotate(x, y, x0, y0) - return 0.5 * (x * ax + y * ay) + # Derivative of the potential + a1 = x * gamma_1 + y * gamma_2 + a2 = x * gamma_2 - y * gamma_1 + return a1, a2 @unpack + @convert_params def convergence( self, x: Tensor, @@ -205,14 +237,14 @@ def convergence( *args, params: Optional["Packed"] = None, z_l: Optional[Tensor] = None, - x0: Optional[Tensor] = None, - y0: Optional[Tensor] = None, gamma_1: Optional[Tensor] = None, gamma_2: Optional[Tensor] = None, + gamma: Optional[Tensor] = None, + phi: Optional[Tensor] = None, **kwargs, ) -> Tensor: """ - The convergence is undefined for an external shear. + The convergence is zero by definition for an external shear. Parameters ---------- @@ -241,10 +273,5 @@ def convergence( *Unit: unitless* - Raises - ------ - NotImplementedError - This method is not implemented as the convergence is not defined - for an external shear. """ - raise NotImplementedError("convergence undefined for external shear") + return torch.zeros_like(x) diff --git a/src/caustics/parametrization.py b/src/caustics/parametrization.py new file mode 100644 index 00000000..3ec62718 --- /dev/null +++ b/src/caustics/parametrization.py @@ -0,0 +1,12 @@ +from functools import wraps +from typing import Callable + + +def convert_params(method: Callable): + @wraps(method) + def wrapper(self, *args, **kwargs): + if hasattr(self, "_convert_params_to_fiducial"): + kwargs = self._convert_params_to_fiducial(*args, **kwargs) + return method(self, *args, **kwargs) + + return wrapper diff --git a/tests/test_external_shear.py b/tests/test_external_shear.py index 8af65ee4..a5b2bae6 100644 --- a/tests/test_external_shear.py +++ b/tests/test_external_shear.py @@ -36,13 +36,11 @@ def test(sim_source, device, lens_models): # Parameters z_s = torch.tensor(2.0, device=device) - x = torch.tensor([0.7, 0.12, -0.52, -0.1, 0.1], device=device) + x = torch.tensor([-0.52, -0.1, 0.1], device=device) kwargs_ls = [ { - "ra_0": x[1].item(), - "dec_0": x[2].item(), - "gamma1": x[3].item(), - "gamma2": x[4].item(), + "gamma1": x[1].item(), + "gamma2": x[2].item(), } ] @@ -51,5 +49,26 @@ def test(sim_source, device, lens_models): ) +def test_parametrization(): + cosmo = FlatLambdaCDM(name="cosmo") + lens = ExternalShear(cosmology=cosmo, z_l=0.5, parametrization="cartesian") + lens_polar = ExternalShear(cosmology=cosmo, z_l=0.5, parametrization="polar") + + gamma_1 = torch.tensor(0.1) + gamma_2 = torch.tensor(0.2) + gamma = torch.sqrt(gamma_1**2 + gamma_2**2) + phi = 0.5 * torch.atan2(gamma_2, gamma_1) + + # Check that the conversion yields the same results in the deflection angle + x = torch.tensor([0.1, 0.2]) + y = torch.tensor([0.2, 0.1]) + z_s = torch.tensor(2.0) + + a1, a2 = lens.reduced_deflection_angle(x, y, z_s, gamma_1=gamma_1, gamma_2=gamma_2) + a1_p, a2_p = lens_polar.reduced_deflection_angle(x, y, z_s, gamma=gamma, phi=phi) + assert torch.allclose(a1, a1_p) + assert torch.allclose(a2, a2_p) + + if __name__ == "__main__": test(None)