feat: New parametrizations for external shear and point mass

src/caustics/lenses/external_shear.py

@@ -1,5 +1,5 @@
 # mypy: disable-error-code="dict-item"
-from typing import Optional, Union, Annotated
+from typing import Optional, Union, Annotated, Literal
 
 from torch import Tensor
 import torch
@@ -71,6 +71,7 @@
         gamma_2: Annotated[
             Optional[Union[Tensor, float]], "Shear component in the y-direction", True
         ] = None,
+        parametrization: Literal["cartesian", "angular"] = "cartesian",
         s: Annotated[
             float, "Softening length for the elliptical power-law profile"
         ] = 0.0,
@@ -82,8 +83,44 @@
         self.y0 = Param("y0", y0, units="arcsec")
         self.gamma_1 = Param("gamma_1", gamma_1, units="unitless")
         self.gamma_2 = Param("gamma_2", gamma_2, units="unitless")
+        self._parametrization = "cartesian"
+        self.parametrization = parametrization
         self.s = s
 
+    @property
+    def parametrization(self) -> str:
+        return self._parametrization
+
+    @parametrization.setter
+    def parametrization(self, value: str):
+        if value not in ["cartesian", "angular"]:
+            raise ValueError(
+                f"Invalid parametrization: {value}. Must be 'cartesian' or 'angular'."
+            )
+        if value == "angular" and self._parametrization != "angular":
+            self.gamma = Param("gamma", shape=self.gamma_1.shape, units="unitless")
+            self.theta = Param("theta", shape=self.gamma_1.shape, units="radians")
+            self.gamma_1.value = lambda p: func.gamma_theta_to_gamma1(
+                p["gamma"].value, p["theta"].value
+            )
+            self.gamma_2.value = lambda p: func.gamma_theta_to_gamma2(
+                p["gamma"].value, p["theta"].value
+            )
+            self.gamma_1.link(self.gamma)
+            self.gamma_1.link(self.theta)
+            self.gamma_2.link(self.gamma)
+            self.gamma_2.link(self.theta)
+        if value == "cartesian" and self._parametrization != "cartesian":
+            try:
+                self.gamma_1 = None
+                self.gamma_2 = None
+                del self.gamma
+                del self.theta
+            except AttributeError:
+                pass
+
+        self._parametrization = value
+
     @forward
     def reduced_deflection_angle(
         self,

src/caustics/lenses/func/__init__.py

@@ -28,6 +28,8 @@
 from .external_shear import (
     reduced_deflection_angle_external_shear,
     potential_external_shear,
+    gamma_theta_to_gamma1,
+    gamma_theta_to_gamma2,
 )
 from .nfw import (
     physical_deflection_angle_nfw,
@@ -107,6 +109,8 @@
     "convergence_epl",
     "reduced_deflection_angle_external_shear",
     "potential_external_shear",
+    "gamma_theta_to_gamma1",
+    "gamma_theta_to_gamma2",
     "physical_deflection_angle_nfw",
     "potential_nfw",
     "convergence_nfw",

src/caustics/lenses/func/external_shear.py

@@ -1,3 +1,5 @@
+import torch
+
 from ...utils import translate_rotate
 
 
@@ -105,3 +107,57 @@ def potential_external_shear(x0, y0, gamma_1, gamma_2, x, y):
     """
     x, y = translate_rotate(x, y, x0, y0)
     return 0.5 * gamma_1 * (x**2 - y**2) + gamma_2 * x * y
+
+
+def gamma_theta_to_gamma1(gamma, theta):
+    """
+    Convert the shear magnitude and angle to the gamma_1 component.
+
+    Parameters
+    ----------
+    gamma: Tensor
+        The shear magnitude.
+
+        *Unit: unitless*
+
+    theta: Tensor
+        The shear angle.
+
+        *Unit: radians*
+
+    Returns
+    -------
+    Tensor
+        The gamma_1 component of the shear.
+
+        *Unit: unitless*
+
+    """
+    return gamma * torch.cos(2 * theta)
+
+
+def gamma_theta_to_gamma2(gamma, theta):
+    """
+    Convert the shear magnitude and angle to the gamma_2 component.
+
+    Parameters
+    ----------
+    gamma: Tensor
+        The shear magnitude.
+
+        *Unit: unitless*
+
+    theta: Tensor
+        The shear angle.
+
+        *Unit: radians*
+
+    Returns
+    -------
+    Tensor
+        The gamma_2 component of the shear.
+
+        *Unit: unitless*
+
+    """
+    return gamma * torch.sin(2 * theta)

src/caustics/lenses/point.py

@@ -1,5 +1,5 @@
 # mypy: disable-error-code="operator,dict-item"
-from typing import Optional, Union, Annotated
+from typing import Optional, Union, Annotated, Literal
 
 from torch import Tensor
 from caskade import forward, Param
@@ -72,6 +72,7 @@
         th_ein: Annotated[
             Optional[Union[Tensor, float]], "Einstein radius of the lens", True
         ] = None,
+        parametrization: Literal["Rein", "mass"] = "Rein",
         s: Annotated[
             float, "Softening parameter to prevent numerical instabilities"
         ] = 0.0,
@@ -119,8 +120,47 @@
         self.x0 = Param("x0", x0, units="arcsec")
         self.y0 = Param("y0", y0, units="arcsec")
         self.th_ein = Param("th_ein", th_ein, units="arcsec", valid=(0, None))
+        self._parametrization = "Rein"
+        self.parametrization = parametrization
         self.s = s
 
+    @property
+    def parametrization(self):
+        return self._parametrization
+
+    @parametrization.setter
+    def parametrization(self, value):
+        if value not in ["Rein", "mass"]:
+            raise ValueError(
+                f"Invalid parametrization {value}. Choose from ['Rein', 'mass']"
+            )
+        if value == "mass" and self.parametrization != "mass":
+            self.z_s = Param("z_s", value=None, shape=self.z_l.shape, units="unitless")
+            self.mass = Param("mass", shape=self.th_ein.shape, units="Msol")
+
+            def mass_to_rein(p):
+                Dls = p["cosmology"].angular_diameter_distance_z1z2(
+                    p["z_l"].value, p["z_s"].value
+                )
+                Dl = p["cosmology"].angular_diameter_distance(p["z_l"].value)
+                Ds = p["cosmology"].angular_diameter_distance(p["z_s"].value)
+                return func.mass_to_rein_point(p["mass"].value, Dls, Dl, Ds)
+
+            self.th_ein.value = mass_to_rein
+            self.th_ein.link(self.mass)
+            self.th_ein.link(self.z_s)
+            self.th_ein.link("cosmology", self.cosmology)
+            self.th_ein.link(self.z_l)
+        if value == "Rein" and self.parametrization != "Rein":
+            try:
+                del self.mass
+                del self.z_s
+                self.th_ein = None
+            except AttributeError:
+                pass
+
+        self._parametrization = value
+
     @forward
     def mass_to_rein(
         self, mass: Tensor, z_s: Tensor, z_l: Annotated[Tensor, "Param"]

tests/test_external_shear.py

@@ -1,13 +1,16 @@
 from io import StringIO
 
 import torch
+import numpy as np
 from lenstronomy.LensModel.lens_model import LensModel
 
 from utils import lens_test_helper
 from caustics.cosmology import FlatLambdaCDM
 from caustics.lenses import ExternalShear
 from caustics.sims import build_simulator
 
+import pytest
+
 
 def test(sim_source, device):
     atol = 1e-5
@@ -50,3 +53,32 @@ def test(sim_source, device):
     lens_test_helper(
         lens, lens_ls, z_s, x, kwargs_ls, rtol, atol, test_kappa=False, device=device
     )
+
+
+def test_external_shear_parametrization():
+
+    cosmology = FlatLambdaCDM(name="cosmo")
+    lens = ExternalShear(name="shear", cosmology=cosmology)
+
+    # Check default
+    assert lens.parametrization == "cartesian"
+
+    # Check set to angular
+    lens.parametrization = "angular"
+    assert lens.parametrization == "angular"
+    # Check setting gamma theta to get gamma1 and gamma2
+    lens.gamma = 1.0
+    lens.theta = np.pi / 4
+    assert np.allclose(lens.gamma_1.value.item(), 0.0, atol=1e-5)
+    assert np.allclose(lens.gamma_2.value.item(), 1.0, atol=1e-5)
+
+    # Check reset to cartesian
+    lens.parametrization = "cartesian"
+    assert lens.parametrization == "cartesian"
+    assert lens.gamma_1.value is None
+    assert lens.gamma_2.value is None
+    assert not hasattr(lens, "gamma")
+    assert not hasattr(lens, "theta")
+
+    with pytest.raises(ValueError):
+        lens.parametrization = "weird"

tests/test_point.py

@@ -1,6 +1,7 @@
 from io import StringIO
 
 import torch
+import numpy as np
 from lenstronomy.LensModel.lens_model import LensModel
 from utils import lens_test_helper
 
@@ -44,3 +45,30 @@ def test_point_lens(sim_source, device, th_ein):
     kwargs_ls = [{"center_x": x[0].item(), "center_y": x[1].item(), "theta_E": th_ein}]
 
     lens_test_helper(lens, lens_ls, z_s, x, kwargs_ls, rtol, atol, device=device)
+
+
+def test_point_parametrization():
+
+    cosmology = FlatLambdaCDM(name="cosmo")
+    lens = Point(name="point", cosmology=cosmology, z_l=0.5)
+
+    # Check default
+    assert lens.parametrization == "Rein"
+
+    # Check set to mass
+    lens.parametrization = "mass"
+    assert lens.parametrization == "mass"
+    # Check setting mass and z_l, and z_s to get mass
+    lens.mass = 1e10
+    lens.z_s = 1.0
+    assert np.allclose(lens.th_ein.value.item(), 0.1637, atol=1e-3)
+
+    # Check reset to cartesian
+    lens.parametrization = "Rein"
+    assert lens.parametrization == "Rein"
+    assert lens.th_ein.value is None
+    assert not hasattr(lens, "mass")
+    assert not hasattr(lens, "z_s")
+
+    with pytest.raises(ValueError):
+        lens.parametrization = "weird"