Added tests for serial GPU transforms

Author: Thomas Baumann

pySDC/helpers/spectral_helper.py

@@ -416,7 +416,6 @@ def transform(self, u, *args, axes=None, shape=None, **kwargs):
         Returns:
             Data in spectral space
         """
-        assert self.fft_lib == scipy.fft
         axes = axes if axes else tuple(i for i in range(u.ndim))
 
         trf = self.fft_lib.dctn(u, *args, axes=axes, type=2, norm='backward', s=shape, **kwargs)
@@ -459,7 +458,7 @@ def transform(self, u, *args, axes=None, shape=None, **kwargs):
             trf *= norm[(*expansion,)]
         return trf
 
-    def itransform(self, u, axes, *args, shape=None, **kwargs):
+    def itransform(self, u, *args, axes=None, shape=None, **kwargs):
         """
         Inverse DCT along axis.
 
@@ -470,7 +469,7 @@ def itransform(self, u, axes, *args, shape=None, **kwargs):
         Returns:
             Data in physical space
         """
-        assert self.fft_lib == scipy.fft
+        axes = axes if axes else tuple(i for i in range(u.ndim))
         for axis in axes:
 
             if self.N == u.shape[axis]:

pySDC/tests/test_helpers/test_spectral_helper_1d_chebychev.py

@@ -163,6 +163,24 @@ def test_transform(N, d):
     assert np.allclose(cheby.itransform(cheby.transform(u, axes=(-1,)), axes=(-1,)), u)
 
 
+@pytest.mark.cupy
+def test_transform_cupy(N=8):
+    import numpy as np
+    import cupy as cp
+    from pySDC.helpers.spectral_helper import ChebychevHelper
+
+    u = cp.random.random(N).astype('D')
+
+    helper_CPU = ChebychevHelper(N=N, useGPU=False)
+    u_hat_CPU = helper_CPU.transform(u.get())
+
+    helper = ChebychevHelper(N=N, useGPU=True)
+    u_hat = helper.transform(u)
+
+    assert cp.allclose(u, helper.itransform(u_hat))
+    assert np.allclose(u_hat.get(), u_hat_CPU)
+
+
 @pytest.mark.base
 @pytest.mark.parametrize('N', [8, 32])
 def test_integration_BC(N):
@@ -402,4 +420,4 @@ def test_tau_method2D_diffusion(nz, nx, bc_val, plotting=False):
 
 
 if __name__ == '__main__':
-    test_transform(4, 1, 'dct')
+    test_transform_cupy()

pySDC/tests/test_helpers/test_spectral_helper_1d_fft.py

@@ -56,6 +56,24 @@ def test_transform(useFFTW, N=8):
     assert np.allclose(u, helper.itransform(u_hat))
 
 
+@pytest.mark.cupy
+def test_transform_cupy(N=8):
+    import numpy as np
+    import cupy as cp
+    from pySDC.helpers.spectral_helper import FFTHelper
+
+    u = cp.random.random(N).astype('D')
+
+    helper_CPU = FFTHelper(N=N, useGPU=False)
+    u_hat_CPU = helper_CPU.transform(u.get())
+
+    helper = FFTHelper(N=N, useGPU=True)
+    u_hat = helper.transform(u)
+
+    assert np.allclose(u_hat.get(), u_hat_CPU)
+    assert cp.allclose(u, helper.itransform(u_hat))
+
+
 @pytest.mark.base
 @pytest.mark.parametrize('N', [8, 64])
 def test_integration_matrix(N, plot=False):
@@ -120,5 +138,6 @@ def test_tau_method(N, v):
     # test_differentiation_matrix(64, 4, 8, True, True)
     # test_integration_matrix(8, True)
     # test_tau_method(6, 1)
-    test_transform(True)
-    test_transform(False)
+    # test_transform(True)
+    # test_transform(False)
+    test_transform_cupy()