tests: use pytest mark parametrize to massively increase the number of unit tests

tests/test_epl.py

@@ -8,9 +8,16 @@
 
 from caustics.cosmology import FlatLambdaCDM
 from caustics.lenses import EPL
+import numpy as np
 
+import pytest
 
-def test_lenstronomy(sim_source, device, lens_models):
+
+@pytest.mark.parametrize("q", [0.4, 0.7])
+@pytest.mark.parametrize("phi", [pi / 3, -pi / 4])
+@pytest.mark.parametrize("b", [0.1, 1.0])
+@pytest.mark.parametrize("t", [0.1, 1.0, 1.9])
+def test_lenstronomy_epl(sim_source, device, lens_models, q, phi, b, t):
     if sim_source == "yaml":
         yaml_str = """\
         cosmology: &cosmology
@@ -36,30 +43,30 @@ def test_lenstronomy(sim_source, device, lens_models):
 
     # Parameters
     z_s = torch.tensor(1.0, device=device)
-    x = torch.tensor([0.7, 0.912, -0.442, 0.7, pi / 3, 1.4, 1.35], device=device)
+    x = torch.tensor([0.7, 0.912, -0.442, q, phi, b, t], device=device)
 
-    e1, e2 = param_util.phi_q2_ellipticity(phi=x[4].item(), q=x[3].item())
-    theta_E = (x[5] / x[3].sqrt()).item()
+    e1, e2 = param_util.phi_q2_ellipticity(phi=phi, q=q)
+    theta_E = b / np.sqrt(q)  # (x[5] / x[3].sqrt()).item()
     kwargs_ls = [
         {
             "theta_E": theta_E,
             "e1": e1,
             "e2": e2,
             "center_x": x[1].item(),
             "center_y": x[2].item(),
-            "gamma": x[6].item() + 1,  # important: add +1
+            "gamma": t + 1,  # important: add +1
         }
     ]
 
     # Different tolerances for difference quantities
     alpha_test_helper(
-        lens, lens_ls, z_s, x, kwargs_ls, rtol=1e-100, atol=6e-5, device=device
+        lens, lens_ls, z_s, x, kwargs_ls, rtol=1e-100, atol=1e-3, device=device
     )
     kappa_test_helper(
         lens, lens_ls, z_s, x, kwargs_ls, rtol=3e-5, atol=1e-100, device=device
     )
     Psi_test_helper(
-        lens, lens_ls, z_s, x, kwargs_ls, rtol=3e-5, atol=1e-100, device=device
+        lens, lens_ls, z_s, x, kwargs_ls, rtol=1e-3, atol=1e-100, device=device
     )
 
 
@@ -98,8 +105,3 @@ def test_special_case_sie(device):
     Psi_test_helper(
         lens, lens_ls, z_s, x, kwargs_ls, rtol=3e-5, atol=1e-100, device=device
     )
-
-
-if __name__ == "__main__":
-    test_lenstronomy(None)
-    test_special_case_sie(None)

tests/test_masssheet.py

@@ -5,8 +5,11 @@
 from caustics.lenses import MassSheet
 from caustics.utils import meshgrid
 
+import pytest
 
-def test(sim_source, device, lens_models):
+
+@pytest.mark.parametrize("convergence", [-1.0, 0.0, 1.0])
+def test_masssheet(sim_source, device, lens_models, convergence):
     if sim_source == "yaml":
         yaml_str = """\
         cosmology: &cosmology
@@ -30,12 +33,17 @@ def test(sim_source, device, lens_models):
 
     # Parameters
     z_s = torch.tensor(1.2)
-    x = torch.tensor([0.5, 0.0, 0.0, 0.7])
+    x = torch.tensor([0.5, 0.0, 0.0, convergence])
 
     thx, thy = meshgrid(0.01, 10, device=device)
 
-    lens.reduced_deflection_angle(thx, thy, z_s, x)
+    ax, ay = lens.reduced_deflection_angle(thx, thy, z_s, x)
+
+    p = lens.potential(thx, thy, z_s, x)
 
-    lens.potential(thx, thy, z_s, x)
+    c = lens.convergence(thx, thy, z_s, x)
 
-    lens.convergence(thx, thy, z_s, x)
+    assert torch.all(torch.isfinite(ax))
+    assert torch.all(torch.isfinite(ay))
+    assert torch.all(torch.isfinite(p))
+    assert torch.all(torch.isfinite(c))

tests/test_nfw.py

@@ -14,12 +14,16 @@
 from caustics.cosmology import FlatLambdaCDM as CausticFlatLambdaCDM
 from caustics.lenses import NFW
 
+import pytest
+
 h0_default = float(default_cosmology.get().h)
 Om0_default = float(default_cosmology.get().Om0)
 Ob0_default = float(default_cosmology.get().Ob0)
 
 
-def test(sim_source, device, lens_models):
+@pytest.mark.parametrize("m", [1e8, 1e10, 1e12])
+@pytest.mark.parametrize("c", [1.0, 8.0, 20.0])
+def test_nfw(sim_source, device, lens_models, m, c):
     atol = 1e-5
     rtol = 3e-2
     z_l = torch.tensor(0.1)
@@ -54,8 +58,8 @@ def test(sim_source, device, lens_models):
 
     thx0 = 0.457
     thy0 = 0.141
-    m = 1e12
-    c = 8.0
+    # m = 1e12
+    # c = 8.0
     x = torch.tensor([thx0, thy0, m, c])
 
     # Lenstronomy
@@ -113,7 +117,3 @@ def test_runs(sim_source, device, lens_models):
     assert torch.all(torch.isfinite(alpha[1]))
     kappa = lens.convergence(thx, thy, z_s, x)
     assert torch.all(torch.isfinite(kappa))
-
-
-if __name__ == "__main__":
-    test()

tests/test_point.py

@@ -6,8 +6,11 @@
 from caustics.cosmology import FlatLambdaCDM
 from caustics.lenses import Point
 
+import pytest
 
-def test(sim_source, device, lens_models):
+
+@pytest.mark.parametrize("th_ein", [0.1, 1.0, 2.0])
+def test_point_lens(sim_source, device, lens_models, th_ein):
     atol = 1e-5
     rtol = 1e-5
     z_l = torch.tensor(0.9)
@@ -37,13 +40,7 @@ def test(sim_source, device, lens_models):
 
     # Parameters
     z_s = torch.tensor(1.2)
-    x = torch.tensor([0.912, -0.442, 1.1])
-    kwargs_ls = [
-        {"center_x": x[0].item(), "center_y": x[1].item(), "theta_E": x[2].item()}
-    ]
+    x = torch.tensor([0.912, -0.442, 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)
-
-
-if __name__ == "__main__":
-    test(None)

tests/test_pseudo_jaffe.py

@@ -6,8 +6,12 @@
 from caustics.cosmology import FlatLambdaCDM
 from caustics.lenses import PseudoJaffe
 
+import pytest
 
-def test(sim_source, device, lens_models):
+
+@pytest.mark.parametrize("mass", [1e8, 1e10, 1e12])
+@pytest.mark.parametrize("Rc,Rs", [[1.0, 10.0], [1e-2, 1.0], [0.5, 1.0]])
+def test_pseudo_jaffe(sim_source, device, lens_models, mass, Rc, Rs):
     atol = 1e-5
     rtol = 1e-5
 
@@ -35,26 +39,9 @@ def test(sim_source, device, lens_models):
 
     # Parameters, computing kappa_0 with a helper function
     z_s = torch.tensor(2.1)
-    x = torch.tensor([0.5, 0.071, 0.023, -1e100, 0.5, 1.5])
-    d_l = cosmology.angular_diameter_distance(x[0])
-    arcsec_to_rad = 1 / (180 / torch.pi * 60**2)
-    kappa_0 = lens.central_convergence(
-        x[0],
-        z_s,
-        torch.tensor(2e11),
-        x[4] * d_l * arcsec_to_rad,
-        x[5] * d_l * arcsec_to_rad,
-        cosmology.critical_surface_density(x[0], z_s),
-    )
-    x[3] = (
-        2
-        * torch.pi
-        * kappa_0
-        * cosmology.critical_surface_density(x[0], z_s)
-        * x[4]
-        * x[5]
-        * (d_l * arcsec_to_rad) ** 2
-    )
+    x = torch.tensor([0.5, 0.071, 0.023, mass, Rc, Rs])
+    kappa_0 = lens.get_convergence_0(z_s, x)
+
     kwargs_ls = [
         {
             "sigma0": kappa_0.item(),
@@ -97,7 +84,3 @@ def test_massenclosed(device):
     masses = lens.mass_enclosed_2d(xx, z_s, x)
 
     assert torch.all(masses < x[3])
-
-
-if __name__ == "__main__":
-    test(None)

tests/test_sersic.py

@@ -8,8 +8,13 @@
 from caustics.light import Sersic
 from caustics.utils import meshgrid
 
+import pytest
 
-def test(sim_source, device, light_models):
+
+@pytest.mark.parametrize("q", [0.2, 0.7])
+@pytest.mark.parametrize("n", [1.0, 2.0, 3.0])
+@pytest.mark.parametrize("th_e", [1.0, 10.0])
+def test_sersic(sim_source, device, light_models, q, n, th_e):
     # Caustics setup
     res = 0.05
     nx = 200
@@ -48,9 +53,9 @@ def test(sim_source, device, light_models):
     thx0_src = 0.05
     thy0_src = 0.01
     phi_src = 0.0
-    q_src = 0.5
-    index_src = 1.5
-    th_e_src = 0.1
+    q_src = q
+    index_src = n
+    th_e_src = th_e
     I_e_src = 100
     # NOTE: in several places we use np.sqrt(q_src) in order to match
     # the definition used by lenstronomy. This only works when phi = 0.
@@ -86,7 +91,3 @@ def test(sim_source, device, light_models):
     brightness_ls = sersic_ls.surface_brightness(x_ls, y_ls, kwargs_light_source)
 
     assert np.allclose(brightness.cpu().numpy(), brightness_ls)
-
-
-if __name__ == "__main__":
-    test(None)

tests/test_sie.py

@@ -10,10 +10,15 @@
 from caustics.lenses import SIE
 from caustics.utils import meshgrid
 
+import pytest
 
-def test(sim_source, device, lens_models):
+
+@pytest.mark.parametrize("q", [0.5, 0.7, 0.9])
+@pytest.mark.parametrize("phi", [pi / 3, -pi / 4, pi / 6])
+@pytest.mark.parametrize("th_ein", [0.1, 1.0, 2.5])
+def test_sie(sim_source, device, lens_models, q, phi, th_ein):
     atol = 1e-5
-    rtol = 1e-5
+    rtol = 1e-3
 
     if sim_source == "yaml":
         yaml_str = """\
@@ -38,11 +43,11 @@ def test(sim_source, device, lens_models):
 
     # Parameters
     z_s = torch.tensor(1.2)
-    x = torch.tensor([0.5, 0.912, -0.442, 0.7, pi / 3, 1.4])
-    e1, e2 = param_util.phi_q2_ellipticity(phi=x[4].item(), q=x[3].item())
+    x = torch.tensor([0.5, 0.912, -0.442, q, phi, th_ein])
+    e1, e2 = param_util.phi_q2_ellipticity(phi=phi, q=q)
     kwargs_ls = [
         {
-            "theta_E": x[5].item(),
+            "theta_E": th_ein,
             "e1": e1,
             "e2": e2,
             "center_x": x[1].item(),
@@ -97,7 +102,3 @@ def test_sie_time_delay():
             )
         )
     )
-
-
-if __name__ == "__main__":
-    test(None)

tests/test_sis.py

@@ -6,8 +6,11 @@
 from caustics.cosmology import FlatLambdaCDM
 from caustics.lenses import SIS
 
+import pytest
 
-def test(sim_source, device, lens_models):
+
+@pytest.mark.parametrize("th_ein", [0.1, 1.0, 2.0])
+def test(sim_source, device, lens_models, th_ein):
     atol = 1e-5
     rtol = 1e-5
     z_l = torch.tensor(0.5)
@@ -37,7 +40,7 @@ def test(sim_source, device, lens_models):
 
     # Parameters
     z_s = torch.tensor(1.2)
-    x = torch.tensor([-0.342, 0.51, 1.4])
+    x = torch.tensor([-0.342, 0.51, th_ein])
     kwargs_ls = [
         {"center_x": x[0].item(), "center_y": x[1].item(), "theta_E": x[2].item()}
     ]

tests/test_tnfw.py

@@ -13,13 +13,20 @@
 
 from caustics.cosmology import FlatLambdaCDM as CausticFlatLambdaCDM
 from caustics.lenses import TNFW
+import pytest
+
 
 h0_default = float(default_cosmology.get().h)
 Om0_default = float(default_cosmology.get().Om0)
 Ob0_default = float(default_cosmology.get().Ob0)
 
 
-def test(sim_source, device, lens_models):
+@pytest.mark.parametrize(
+    "m", [1e8, 1e10, 1e12]
+)  # Note with m=1e14 the test fails, due to the Rs_angle becoming too large (pytorch is unstable)
+@pytest.mark.parametrize("c", [1.0, 8.0, 40.0])
+@pytest.mark.parametrize("t", [2.0, 5.0, 20.0])
+def test(sim_source, device, lens_models, m, c, t):
     atol = 1e-5
     rtol = 3e-2
     z_l = torch.tensor(0.1)
@@ -51,16 +58,11 @@ def test(sim_source, device, lens_models):
     lens_model_list = ["TNFW"]
     lens_ls = LensModel(lens_model_list=lens_model_list)
 
-    print(lens)
-
     # Parameters
     z_s = torch.tensor(0.5)
 
     thx0 = 0.457
     thy0 = 0.141
-    m = 1e12
-    c = 8.0
-    t = 3.0
     x = torch.tensor([thx0, thy0, m, c, t])
 
     # Lenstronomy
@@ -118,7 +120,3 @@ def test_runs(device):
     Rt = lens.get_truncation_radius(x)
 
     assert Rt == (rs * t)
-
-
-if __name__ == "__main__":
-    test(None)