reduce total number of test cases, reactivate CI

Author: CosmoMatt

.github/workflows/tests.yml

@@ -25,14 +25,14 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
 
-    # - name: Install dependencies
-    #   run: |
-    #     python -m pip install --upgrade pip
-    #     pip install -r requirements/requirements-tests.txt
-    #     pip install -r requirements/requirements-core.txt
-    #     pip install .
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -r requirements/requirements-tests.txt
+        pip install -r requirements/requirements-core.txt
+        pip install .
     
-    # - name: Run tests
-    #   run: |
-    #     pytest --cov-report term --cov=s2wav --cov-config=.coveragerc 
-    #     codecov --token 298dc7ee-bb9f-4221-b31f-3576cc6cb702
+    - name: Run tests
+      run: |
+        pytest --cov-report term --cov=s2wav --cov-config=.coveragerc 
+        codecov --token 298dc7ee-bb9f-4221-b31f-3576cc6cb702

tests/test_samples.py

@@ -5,7 +5,7 @@
 from s2fft.sampling import s2_samples as samples
 
 
-nside_to_test = [32, 64, 128]
+nside_to_test = [16, 32]
 
 
 @pytest.mark.parametrize("L", [15, 16])

tests/test_spherical_base.py

@@ -6,10 +6,10 @@
 import healpy as hp
 
 
-L_to_test = [6, 7, 8]
-L_lower_to_test = [0, 1, 2]
-spin_to_test = [-2, -1, 0, 1, 2]
-nside_to_test = [2, 4, 8]
+L_to_test = [6, 7]
+L_lower_to_test = [0, 2]
+spin_to_test = [-2, 0, 1]
+nside_to_test = [4, 5]
 L_to_nside_ratio = [2, 3]
 sampling_to_test = ["mw", "mwss", "dh"]
 method_to_test = ["direct", "sov", "sov_fft", "sov_fft_vectorized"]

tests/test_spherical_custom_grads.py

@@ -2,29 +2,25 @@
 
 config.update("jax_enable_x64", True)
 import pytest
-import numpy as np
 import jax.numpy as jnp
 from jax.test_util import check_grads
 
-from s2fft.sampling import s2_samples as samples
 from s2fft.transforms import spherical
 from s2fft.recursions.price_mcewen import generate_precomputes_jax
 
 L_to_test = [16]
-L_lower_to_test = [0, 2]
+L_lower_to_test = [2]
 spin_to_test = [-2, 0, 1]
-nside_to_test = [8, 10]
+nside_to_test = [8]
 sampling_to_test = ["mw", "mwss", "dh"]
 reality_to_test = [False, True]
-multiple_gpus = [False]
 
 
 @pytest.mark.parametrize("L", L_to_test)
 @pytest.mark.parametrize("L_lower", L_lower_to_test)
 @pytest.mark.parametrize("spin", spin_to_test)
 @pytest.mark.parametrize("sampling", sampling_to_test)
 @pytest.mark.parametrize("reality", reality_to_test)
-@pytest.mark.parametrize("spmd", multiple_gpus)
 @pytest.mark.filterwarnings("ignore::RuntimeWarning")
 def test_inverse_custom_gradients(
     flm_generator,
@@ -33,7 +29,6 @@ def test_inverse_custom_gradients(
     spin: int,
     sampling: str,
     reality: bool,
-    spmd: bool,
 ):
     if reality and spin != 0:
         pytest.skip("Reality only valid for scalar fields (spin=0).")
@@ -61,7 +56,6 @@ def func(flm):
             reality=reality,
             precomps=precomps,
             sampling=sampling,
-            spmd=spmd,
         )
         return jnp.sum(jnp.abs(f - f_target) ** 2)
 
@@ -73,7 +67,6 @@ def func(flm):
 @pytest.mark.parametrize("spin", spin_to_test)
 @pytest.mark.parametrize("sampling", sampling_to_test)
 @pytest.mark.parametrize("reality", reality_to_test)
-@pytest.mark.parametrize("spmd", multiple_gpus)
 @pytest.mark.filterwarnings("ignore::RuntimeWarning")
 def test_forward_custom_gradients(
     flm_generator,
@@ -82,7 +75,6 @@ def test_forward_custom_gradients(
     spin: int,
     sampling: str,
     reality: bool,
-    spmd: bool,
 ):
     if reality and spin != 0:
         pytest.skip("Reality only valid for scalar fields (spin=0).")
@@ -110,7 +102,6 @@ def func(f):
             reality=reality,
             precomps=precomps,
             sampling=sampling,
-            spmd=spmd,
         )
         return jnp.sum(jnp.abs(flm - flm_target) ** 2)
 
@@ -122,7 +113,6 @@ def func(f):
 @pytest.mark.parametrize("spin", spin_to_test)
 @pytest.mark.parametrize("sampling", sampling_to_test)
 @pytest.mark.parametrize("reality", reality_to_test)
-@pytest.mark.parametrize("spmd", multiple_gpus)
 @pytest.mark.filterwarnings("ignore::RuntimeWarning")
 def test_healpix_inverse_custom_gradients(
     flm_generator,
@@ -131,7 +121,6 @@ def test_healpix_inverse_custom_gradients(
     spin: int,
     sampling: str,
     reality: bool,
-    spmd: bool,
 ):
     sampling = "healpix"
     L = 2 * nside
@@ -164,7 +153,6 @@ def func(flm):
             reality=reality,
             precomps=precomps,
             sampling=sampling,
-            spmd=spmd,
         )
         return jnp.sum(jnp.abs(f - f_target) ** 2)
 
@@ -176,7 +164,6 @@ def func(flm):
 @pytest.mark.parametrize("spin", spin_to_test)
 @pytest.mark.parametrize("sampling", sampling_to_test)
 @pytest.mark.parametrize("reality", reality_to_test)
-@pytest.mark.parametrize("spmd", multiple_gpus)
 @pytest.mark.filterwarnings("ignore::RuntimeWarning")
 def test_healpix_forward_custom_gradients(
     flm_generator,
@@ -185,7 +172,6 @@ def test_healpix_forward_custom_gradients(
     spin: int,
     sampling: str,
     reality: bool,
-    spmd: bool,
 ):
     sampling = "healpix"
     L = 2 * nside
@@ -218,7 +204,6 @@ def func(f):
             reality=reality,
             precomps=precomps,
             sampling=sampling,
-            spmd=spmd,
         )
         return jnp.sum(jnp.abs(flm - flm_target) ** 2)
 

tests/test_spherical_precompute.py

@@ -4,9 +4,9 @@
 from s2fft.precompute_transforms.construct import spin_spherical_kernel
 from s2fft.base_transforms import spherical as base
 
-L_to_test = [6, 7, 8]
-spin_to_test = [-2, -1, 0, 1, 2]
-nside_to_test = [2, 4, 8]
+L_to_test = [6, 7]
+spin_to_test = [-2, 0, 1]
+nside_to_test = [4, 5]
 L_to_nside_ratio = [2, 3]
 sampling_to_test = ["mw", "mwss", "dh"]
 reality_to_test = [True, False]

tests/test_spherical_transform.py

@@ -10,10 +10,10 @@
 from s2fft.transforms import spherical
 from s2fft.recursions.price_mcewen import generate_precomputes
 
-L_to_test = [6, 7, 8]
-L_lower_to_test = [0, 1, 2]
-spin_to_test = [-2, -1, 0, 1, 2]
-nside_to_test = [2, 4, 8]
+L_to_test = [6, 7]
+L_lower_to_test = [0, 2]
+spin_to_test = [-2, 0, 1]
+nside_to_test = [4, 5]
 sampling_to_test = ["mw", "mwss", "dh"]
 method_to_test = ["numpy", "jax"]
 reality_to_test = [False, True]

tests/test_wigner_base.py

@@ -5,9 +5,9 @@
 from s2fft.base_transforms import wigner
 
 
-L_to_test = [8, 16]
-N_to_test = [2, 4, 6]
-L_lower_to_test = [0, 2, 4]
+L_to_test = [6, 7]
+N_to_test = [2, 3]
+L_lower_to_test = [0, 2]
 sampling_schemes_so3 = ["mw", "mwss"]
 sampling_schemes = ["mw", "mwss", "dh"]
 reality_to_test = [False, True]

tests/test_wigner_custom_grads.py

@@ -8,20 +8,18 @@
 from s2fft.transforms import wigner
 from s2fft.recursions.price_mcewen import generate_precomputes_wigner_jax
 
-L_to_test = [16]
-N_to_test = [2]
-L_lower_to_test = [0, 2]
+L_to_test = [6]
+N_to_test = [3]
+L_lower_to_test = [1]
 sampling_to_test = ["mw", "mwss", "dh"]
-reality_to_test = [False]
-multiple_gpus = [False]
+reality_to_test = [False, True]
 
 
 @pytest.mark.parametrize("L", L_to_test)
 @pytest.mark.parametrize("N", N_to_test)
 @pytest.mark.parametrize("L_lower", L_lower_to_test)
 @pytest.mark.parametrize("sampling", sampling_to_test)
 @pytest.mark.parametrize("reality", reality_to_test)
-@pytest.mark.parametrize("spmd", multiple_gpus)
 @pytest.mark.filterwarnings("ignore::RuntimeWarning")
 def test_inverse_wigner_custom_gradients(
     flmn_generator,
@@ -30,7 +28,6 @@ def test_inverse_wigner_custom_gradients(
     L_lower: int,
     sampling: str,
     reality: bool,
-    spmd: bool,
 ):
     precomps = generate_precomputes_wigner_jax(
         L, N, sampling, None, False, reality, L_lower
@@ -39,12 +36,12 @@ def test_inverse_wigner_custom_gradients(
     flmn = flmn_generator(L=L, N=N, L_lower=L_lower, reality=reality)
     flmn_target = flmn_generator(L=L, N=N, L_lower=L_lower, reality=reality)
     f_target = wigner.inverse_jax(
-        flmn_target, L, N, None, sampling, reality, precomps, spmd, L_lower
+        flmn_target, L, N, None, sampling, reality, precomps, False, L_lower
     )
 
     def func(flmn):
         f = wigner.inverse_jax(
-            flmn, L, N, None, sampling, reality, precomps, spmd, L_lower
+            flmn, L, N, None, sampling, reality, precomps, False, L_lower
         )
         return jnp.sum(jnp.abs(f - f_target) ** 2)
 
@@ -56,7 +53,6 @@ def func(flmn):
 @pytest.mark.parametrize("L_lower", L_lower_to_test)
 @pytest.mark.parametrize("sampling", sampling_to_test)
 @pytest.mark.parametrize("reality", reality_to_test)
-@pytest.mark.parametrize("spmd", multiple_gpus)
 @pytest.mark.filterwarnings("ignore::RuntimeWarning")
 def test_forward_wigner_custom_gradients(
     flmn_generator,
@@ -65,7 +61,6 @@ def test_forward_wigner_custom_gradients(
     L_lower: int,
     sampling: str,
     reality: bool,
-    spmd: bool,
 ):
     precomps = generate_precomputes_wigner_jax(
         L, N, sampling, None, True, reality, L_lower
@@ -74,12 +69,12 @@ def test_forward_wigner_custom_gradients(
     flmn_target = flmn_generator(L=L, N=N, L_lower=L_lower, reality=reality)
     flmn = flmn_generator(L=L, N=N, L_lower=L_lower, reality=reality)
     f = wigner.inverse_jax(
-        flmn, L, N, None, sampling, reality, None, spmd, L_lower
+        flmn, L, N, None, sampling, reality, None, False, L_lower
     )
 
     def func(f):
         flmn = wigner.forward_jax(
-            f, L, N, None, sampling, reality, precomps, spmd, L_lower
+            f, L, N, None, sampling, reality, precomps, False, L_lower
         )
         return jnp.sum(jnp.abs(flmn - flmn_target) ** 2)
 

tests/test_wigner_precompute.py

@@ -6,9 +6,9 @@
 from s2fft.base_transforms import wigner as base
 
 L_to_test = [8, 10]
-N_to_test = [2, 4]
+N_to_test = [2, 3]
 nside_to_test = [4, 6]
-L_to_nside_ratio = [2, 3]
+L_to_nside_ratio = [2]
 reality_to_test = [False, True]
 sampling_schemes = ["mw", "mwss", "dh"]
 methods_to_test = ["numpy", "jax"]

tests/test_wigner_recursions.py

@@ -8,8 +8,8 @@
 from s2fft.sampling import s2_samples as samples
 import pyssht as ssht
 
-L_to_test = [8, 16]
-spin_to_test = np.arange(-2, 2)
+L_to_test = [6, 7]
+spin_to_test = [-2, 0, 1]
 sampling_schemes = ["mw", "mwss", "dh", "healpix"]