From e154ff8c416b5cf1cf054bae32bbc096cf79978d Mon Sep 17 00:00:00 2001 From: "Connor Stone, PhD" Date: Thu, 12 Sep 2024 09:27:49 -0700 Subject: [PATCH] use pytest mark parametrize to massively increase the number of tests (#255) --- tests/test_epl.py | 26 ++++++++++++++------------ tests/test_masssheet.py | 18 +++++++++++++----- tests/test_nfw.py | 14 +++++++------- tests/test_point.py | 15 ++++++--------- tests/test_pseudo_jaffe.py | 33 ++++++++------------------------- tests/test_sersic.py | 17 +++++++++-------- tests/test_sie.py | 19 ++++++++++--------- tests/test_sis.py | 7 +++++-- tests/test_tnfw.py | 18 ++++++++---------- tests/utils/__init__.py | 19 +++++++++++-------- 10 files changed, 91 insertions(+), 95 deletions(-) diff --git a/tests/test_epl.py b/tests/test_epl.py index f0d79491..1c854e13 100644 --- a/tests/test_epl.py +++ b/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,10 +43,10 @@ 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, @@ -47,19 +54,19 @@ def test_lenstronomy(sim_source, device, lens_models): "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) diff --git a/tests/test_masssheet.py b/tests/test_masssheet.py index a1f79ef2..b455d05b 100644 --- a/tests/test_masssheet.py +++ b/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)) diff --git a/tests/test_nfw.py b/tests/test_nfw.py index a9c30191..cd6c2984 100644 --- a/tests/test_nfw.py +++ b/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() diff --git a/tests/test_point.py b/tests/test_point.py index 2584c944..022e3fcc 100644 --- a/tests/test_point.py +++ b/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) diff --git a/tests/test_pseudo_jaffe.py b/tests/test_pseudo_jaffe.py index 84d75fe7..6e45b84e 100644 --- a/tests/test_pseudo_jaffe.py +++ b/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) diff --git a/tests/test_sersic.py b/tests/test_sersic.py index c215bbb3..0c8089e9 100644 --- a/tests/test_sersic.py +++ b/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) diff --git a/tests/test_sie.py b/tests/test_sie.py index f4d06e3a..8b379936 100644 --- a/tests/test_sie.py +++ b/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) diff --git a/tests/test_sis.py b/tests/test_sis.py index a4594e55..cf170620 100644 --- a/tests/test_sis.py +++ b/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()} ] diff --git a/tests/test_tnfw.py b/tests/test_tnfw.py index 33a2b5eb..46e1f6e1 100644 --- a/tests/test_tnfw.py +++ b/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) diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py index d1121b0d..b18b7605 100644 --- a/tests/utils/__init__.py +++ b/tests/utils/__init__.py @@ -1,6 +1,7 @@ """ Utilities for testing """ + from typing import Any, Dict, List, Union import torch @@ -16,11 +17,12 @@ from caustics.cosmology import FlatLambdaCDM from .models import mock_from_file -__all__ = ( - "mock_from_file", -) +__all__ = ("mock_from_file",) + -def setup_simulator(cosmo_static=False, use_nfw=True, simulator_static=False, batched_params=False, device=None): +def setup_simulator( + cosmo_static=False, use_nfw=True, simulator_static=False, batched_params=False, device=None +): n_pix = 20 class Sim(Simulator): @@ -88,7 +90,7 @@ def forward(self, params): cosmo_params = [_x[0] for _x in cosmo_params] lens_params = [_x[0] for _x in lens_params] source_params = [_x[0] for _x in source_params] - + sim = Sim() # Set device when not None if device is not None: @@ -97,7 +99,7 @@ def forward(self, params): cosmo_params = [_p.to(device=device) for _p in cosmo_params] lens_params = [_p.to(device=device) for _p in lens_params] source_params = [_p.to(device=device) for _p in source_params] - + return sim, (sim_params, cosmo_params, lens_params, source_params) @@ -161,7 +163,7 @@ def forward(self, params): lens_params = [_x[0] for _x in lens_params] kappa = kappa[0] source = source[0] - + sim = Sim() # Set device when not None if device is not None: @@ -177,7 +179,7 @@ def forward(self, params): def get_default_cosmologies(device=None): cosmology = FlatLambdaCDM("cosmo") cosmology_ap = FlatLambdaCDM_AP(100 * cosmology.h0.value, cosmology.Om0.value, Tcmb0=0) - + if device is not None: cosmology = cosmology.to(device=device) return cosmology, cosmology_ap @@ -210,6 +212,7 @@ def alpha_test_helper(lens, lens_ls, z_s, x, kwargs_ls, atol, rtol, device=None) thx, thy, thx_ls, thy_ls = setup_grids(device=device) alpha_x, alpha_y = lens.reduced_deflection_angle(thx, thy, z_s, x) alpha_x_ls, alpha_y_ls = lens_ls.alpha(thx_ls, thy_ls, kwargs_ls) + print(np.sum(np.abs(1 - alpha_x.cpu().numpy() / alpha_x_ls) > 1e-3)) assert np.allclose(alpha_x.cpu().numpy(), alpha_x_ls, rtol, atol) assert np.allclose(alpha_y.cpu().numpy(), alpha_y_ls, rtol, atol)