Added test for 2D differentiation matrix on GPU

Author: Thomas Baumann

pySDC/helpers/spectral_helper.py

@@ -746,11 +746,14 @@ def get_differentiation_matrix(self, p=1):
         k = self.get_wavenumbers()
 
         if self.useGPU:
-            # Have to raise the matrix to power p on CPU because the GPU equivalent is not implemented in CuPy at the time of writing.
-            import scipy.sparse as sp
+            if p > 1:
+                # Have to raise the matrix to power p on CPU because the GPU equivalent is not implemented in CuPy at the time of writing.
+                from scipy.sparse.linalg import matrix_power
 
-            D = self.sparse_lib.diags(1j * k).get()
-            return self.sparse_lib.csc_matrix(sp.linalg.matrix_power(D, p))
+                D = self.sparse_lib.diags(1j * k).get()
+                return self.sparse_lib.csc_matrix(matrix_power(D, p))
+            else:
+                return self.sparse_lib.diags(1j * k)
         else:
             return self.linalg.matrix_power(self.sparse_lib.diags(1j * k), p)
 

pySDC/tests/test_helpers/test_spectral_helper.py

@@ -52,15 +52,21 @@ def test_integration_matrix2D(nx, nz, variant, axes, useMPI=False, **kwargs):
 @pytest.mark.parametrize('axes', [(-2,), (-1,), (-2, -1)])
 @pytest.mark.parametrize('bx', ['cheby', 'fft'])
 @pytest.mark.parametrize('bz', ['cheby', 'fft'])
-def test_differentiation_matrix2D(nx, nz, variant, axes, bx, bz, **kwargs):
-    import numpy as np
+def test_differentiation_matrix2D(nx, nz, variant, axes, bx, bz, useGPU=False, **kwargs):
     from pySDC.helpers.spectral_helper import SpectralHelper
 
-    helper = SpectralHelper(debug=True)
+    helper = SpectralHelper(debug=True, useGPU=useGPU)
     helper.add_axis(base=bx, N=nx)
     helper.add_axis(base=bz, N=nz)
     helper.setup_fft()
 
+    np = helper.xp
+
+    if useGPU:
+        import cupy
+
+        assert np == cupy
+
     X, Z = helper.get_grid()
     conv = helper.get_basis_change_matrix()
     D = helper.get_differentiation_matrix(axes)
@@ -95,6 +101,14 @@ def test_differentiation_matrix2D(nx, nz, variant, axes, bx, bz, **kwargs):
     assert np.allclose(D_u, expect, atol=1e-11)
 
 
+@pytest.mark.cupy
+@pytest.mark.parametrize('axes', [(-2,), (-1,), (-2, -1)])
+@pytest.mark.parametrize('bx', ['cheby', 'fft'])
+@pytest.mark.parametrize('bz', ['cheby', 'fft'])
+def test_differentiation_matrix2D_GPU(bx, bz, axes):
+    test_differentiation_matrix2D(32, 16, variant='T2U', bx=bx, bz=bz, axes=axes, useGPU=True)
+
+
 @pytest.mark.base
 @pytest.mark.parametrize('nx', [32])
 @pytest.mark.parametrize('nz', [16])