Implement requested changes

Author: ASKabalan

s2fft/utils/healpix_ffts.py

@@ -711,11 +711,18 @@ def _healpix_fft_cuda_abstract(f, L, nside, reality, fft_type, norm, adjoint):
         raise ValueError(f"fft_type {fft_type} not recognised.")
 
 
+class MissingCUDASupport(Exception):  # noqa : D107
+    def __init__(self):  # noqa : D107
+        super().__init__("""
+                        S2FFT was compiled without CUDA support. Cuda functions are not supported.
+                        Please make sure that nvcc is in your path and $CUDA_HOME is set then reinstall s2fft using pip.
+                        """)
+
+
 def _healpix_fft_cuda_lowering(ctx, f, *, L, nside, reality, fft_type, norm, adjoint):
-    assert _s2fft.COMPILED_WITH_CUDA, """
-    S2FFT was compiled without CUDA support. Cuda functions are not supported.
-    Please make sure that nvcc is in your path and $CUDA_HOME is set then reinstall s2fft using pip.
-    """
+    if not _s2fft.COMPILED_WITH_CUDA:
+        raise MissingCUDASupport()
+
     (aval_out,) = ctx.avals_out
 
     out_dtype = aval_out.dtype

tests/test_healpix_ffts.py

@@ -6,6 +6,7 @@
 from numpy.testing import assert_allclose
 from packaging.version import Version as _Version
 
+import s2fft
 from s2fft.sampling import s2_samples as samples
 from s2fft.utils.healpix_ffts import (
     healpix_fft_cuda,
@@ -103,8 +104,9 @@ def test_healpix_fft_cuda_transforms(flm_generator, nside):
     # Generate a random bandlimited signal
     def generate_flm():
         flm = flm_generator(L=L, reality=False)
-        flm_hp = samples.flm_2d_to_hp(flm, L)
-        f = hp.sphtfunc.alm2map(flm_hp, nside, lmax=L - 1)
+        f = s2fft.inverse(
+            flm, L=L, nside=nside, reality=False, method="jax", sampling="healpix"
+        )
         return f
 
     f_stacked = jnp.stack([generate_flm() for _ in range(10)], axis=0)
@@ -125,15 +127,15 @@ def healpix_cuda(f):
     )
     # test jacfwd
     assert_allclose(
-        jax.jacfwd(healpix_jax)(f),
-        jax.jacfwd(healpix_cuda)(f),
+        jax.jacfwd(healpix_jax)(f.real),
+        jax.jacfwd(healpix_cuda)(f.real),
         atol=1e-7,
         rtol=1e-7,
     )
     # test jacrev
     assert_allclose(
-        jax.jacrev(healpix_jax)(f),
-        jax.jacrev(healpix_cuda)(f),
+        jax.jacrev(healpix_jax)(f.real),
+        jax.jacrev(healpix_cuda)(f.real),
         atol=1e-7,
         rtol=1e-7,
     )
@@ -147,8 +149,9 @@ def test_healpix_ifft_cuda_transforms(flm_generator, nside):
     # Generate a random bandlimited signal
     def generate_flm():
         flm = flm_generator(L=L, reality=False)
-        flm_hp = samples.flm_2d_to_hp(flm, L)
-        f = hp.sphtfunc.alm2map(flm_hp, nside, lmax=L - 1)
+        f = s2fft.inverse(
+            flm, L=L, nside=nside, reality=False, method="jax", sampling="healpix"
+        )
         ftm = healpix_fft_jax(f, L, nside, False)
         return ftm
 
@@ -164,23 +167,23 @@ def healpix_inv_cuda(f):
     # Test VMAP
     assert_allclose(
         jax.vmap(healpix_inv_jax)(ftm_stacked).flatten(),
-        jax.vmap(healpix_inv_jax)(ftm_stacked).flatten(),
+        jax.vmap(healpix_inv_cuda)(ftm_stacked).flatten(),
         atol=1e-7,
         rtol=1e-7,
     )
 
     # test jacfwd
     assert_allclose(
-        jax.jacfwd(healpix_inv_jax)(ftm).flatten(),
-        jax.jacfwd(healpix_inv_cuda)(ftm).flatten(),
+        jax.jacfwd(healpix_inv_jax)(ftm.real).flatten(),
+        jax.jacfwd(healpix_inv_cuda)(ftm.real).flatten(),
         atol=1e-7,
         rtol=1e-7,
     )
 
     # test jacrev
     assert_allclose(
-        jax.jacrev(healpix_inv_jax)(ftm).flatten(),
-        jax.jacrev(healpix_inv_cuda)(ftm).flatten(),
+        jax.jacrev(healpix_inv_jax)(ftm.real).flatten(),
+        jax.jacrev(healpix_inv_cuda)(ftm.real).flatten(),
         atol=1e-7,
         rtol=1e-7,
     )