Test across different iter values

Author: matt-graham

tests/test_spherical_custom_grads.py

@@ -161,13 +161,15 @@ def func(flm):
 @pytest.mark.parametrize("L_lower", L_lower_to_test)
 @pytest.mark.parametrize("spin", spin_to_test)
 @pytest.mark.parametrize("reality", reality_to_test)
+@pytest.mark.parametrize("iter", [0, 1, 2, 3])
 @pytest.mark.filterwarnings("ignore::RuntimeWarning")
 def test_healpix_forward_custom_gradients(
     flm_generator,
     nside: int,
     L_lower: int,
     spin: int,
     reality: bool,
+    iter: int,
 ):
     sampling = "healpix"
     L = 2 * nside
@@ -191,15 +193,17 @@ def test_healpix_forward_custom_gradients(
     )
 
     def func(f):
-        flm = spherical.forward_jax(
+        flm = spherical.forward(
             f,
             L,
+            method="jax",
             spin=spin,
             nside=nside,
             L_lower=L_lower,
             reality=reality,
             precomps=precomps,
             sampling=sampling,
+            iter=iter,
         )
         return jnp.sum(jnp.abs(flm - flm_target) ** 2)
 

tests/test_spherical_transform.py

@@ -150,12 +150,14 @@ def test_transform_forward(
 @pytest.mark.parametrize("nside", nside_to_test)
 @pytest.mark.parametrize("method", method_to_test)
 @pytest.mark.parametrize("spmd", multiple_gpus)
+@pytest.mark.parametrize("iter", [0, 1, 2, 3])
 @pytest.mark.filterwarnings("ignore::RuntimeWarning")
 def test_transform_forward_healpix(
     flm_generator,
     nside: int,
     method: str,
     spmd: bool,
+    iter: int,
 ):
     sampling = "healpix"
     L = 2 * nside
@@ -174,10 +176,11 @@ def test_transform_forward_healpix(
         reality=True,
         precomps=precomps,
         spmd=spmd,
+        iter=iter,
     )
     flm_check = samples.flm_2d_to_hp(flm_check, L)
 
-    flm = hp.sphtfunc.map2alm(f, lmax=L - 1, iter=0)
+    flm = hp.sphtfunc.map2alm(f, lmax=L - 1, iter=iter)
 
     np.testing.assert_allclose(flm, flm_check, atol=1e-14)