Add tests for new parametrizations

src/caustics/lenses/external_shear.py

@@ -100,7 +100,10 @@ def parametrization(self, value: str):
         if value == "angular" and self._parametrization != "angular":
             try:
                 gamma = torch.sqrt(self.gamma_1.value**2 + self.gamma_2.value**2)
-                theta = 0.5 * torch.acos(self.gamma_1.value / gamma)
+                if gamma.item() == 0:
+                    theta = 0.0
+                else:
+                    theta = 0.5 * torch.acos(self.gamma_1.value / gamma)
             except TypeError:
                 gamma = None
                 theta = None

src/caustics/lenses/point.py

@@ -142,8 +142,8 @@ def mass_to_rein(p):
                 Dls = p["cosmology"].angular_diameter_distance_z1z2(
                     p["z_l"].value, p["z_s"].value
                 )
-                Dl = p["cosmology"].cosmology.angular_diameter_distance(p["z_l"].value)
-                Ds = p["cosmology"].cosmology.angular_diameter_distance(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

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,49 @@ 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")
+
+    # Check set to angular when gamma1 and gamma2 have values
+    lens.gamma_1 = 0.0
+    lens.gamma_2 = 1.0
+    lens.parametrization = "angular"
+    assert np.allclose(lens.gamma.value.item(), 1.0)
+    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 case where gamma = 0
+    lens.parametrization = "cartesian"
+    lens.gamma_1 = 0.0
+    lens.gamma_2 = 0.0
+    lens.parametrization = "angular"
+    assert np.allclose(lens.gamma.value.item(), 0.0, atol=1e-5)
+    assert np.allclose(lens.gamma_1.value.item(), 0.0, atol=1e-5)
+    assert np.allclose(lens.gamma_2.value.item(), 0.0, atol=1e-5)
+
+    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"