Expanded tests

Author: brownbaerchen

pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py

@@ -1,13 +1,11 @@
 import numpy as np
 from mpi4py import MPI
-from mpi4py_fft import PFFT
+from mpi4py_fft import PFFT, newDistArray
 
 from pySDC.core.errors import ProblemError
 from pySDC.core.problem import Problem, WorkCounter
 from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh
 
-from mpi4py_fft import newDistArray
-
 
 class IMEX_Laplacian_MPIFFT(Problem):
     r"""
@@ -99,14 +97,24 @@ def __init__(
             'nvars', 'spectral', 'L', 'alpha', 'comm', 'x0', 'useGPU', localVars=locals(), readOnly=True
         )
 
-        # get local mesh
+        self.getLocalGrid()
+        self.getLaplacian()
+
+        # Need this for diagnostics
+        self.dx = self.L[0] / nvars[0]
+        self.dy = self.L[1] / nvars[1]
+
+        # work counters
+        self.work_counters['rhs'] = WorkCounter()
+
+    def getLocalGrid(self):
         X = list(self.xp.ogrid[self.fft.local_slice(False)])
         N = self.fft.global_shape()
         for i in range(len(N)):
-            X[i] = x0 + (X[i] * L[i] / N[i])
+            X[i] = self.x0 + (X[i] * self.L[i] / N[i])
         self.X = [self.xp.broadcast_to(x, self.fft.shape(False)) for x in X]
 
-        # get local wavenumbers and Laplace operator
+    def getLaplacian(self):
         s = self.fft.local_slice()
         N = self.fft.global_shape()
         k = [self.xp.fft.fftfreq(n, 1.0 / n).astype(int) for n in N]
@@ -117,14 +125,7 @@ def __init__(
             Ks[i] = (Ks[i] * Lp[i]).astype(float)
         K = [self.xp.broadcast_to(k, self.fft.shape(True)) for k in Ks]
         K = self.xp.array(K).astype(float)
-        self.K2 = self.xp.sum(K * K, 0, dtype=float)  # Laplacian in spectral space
-
-        # Need this for diagnostics
-        self.dx = self.L[0] / nvars[0]
-        self.dy = self.L[1] / nvars[1]
-
-        # work counters
-        self.work_counters['rhs'] = WorkCounter()
+        self.K2 = self.xp.sum(K * K, 0, dtype=float)
 
     def eval_f(self, u, t):
         """

pySDC/tests/test_problems/test_generic_MPIFFT.py

@@ -2,39 +2,56 @@
 
 
 @pytest.mark.mpi4py
-@pytest.mark.parametrize('nx', [16, 32])
-@pytest.mark.parametrize('ny', [16, 32])
-@pytest.mark.parametrize('nz', [0])
+@pytest.mark.parametrize('nx', [8, 16])
+@pytest.mark.parametrize('ny', [8, 16])
+@pytest.mark.parametrize('nz', [0, 8])
 @pytest.mark.parametrize('f', [1, 3])
+@pytest.mark.parametrize('spectral', [True, False])
 @pytest.mark.parametrize('direction', [0, 1, 10])
-def test_derivative(nx, ny, nz, f, direction):
+def test_derivative(nx, ny, nz, f, spectral, direction):
     from pySDC.implementations.problem_classes.generic_MPIFFT_Laplacian import IMEX_Laplacian_MPIFFT
 
     nvars = (nx, ny)
     if nz > 0:
-        nvars.append(nz)
-    prob = IMEX_Laplacian_MPIFFT(nvars=nvars)
+        nvars += (nz,)
+    prob = IMEX_Laplacian_MPIFFT(nvars=nvars, spectral=spectral)
 
     xp = prob.xp
 
     if direction == 0:
-        u = xp.sin(f * prob.X[0])
+        _u = xp.sin(f * prob.X[0])
         du_expect = -(f**2) * xp.sin(f * prob.X[0])
     elif direction == 1:
-        u = xp.sin(f * prob.X[1])
+        _u = xp.sin(f * prob.X[1])
         du_expect = -(f**2) * xp.sin(f * prob.X[1])
     elif direction == 10:
-        u = xp.sin(f * prob.X[1]) + xp.cos(f * prob.X[0])
+        _u = xp.sin(f * prob.X[1]) + xp.cos(f * prob.X[0])
         du_expect = -(f**2) * xp.sin(f * prob.X[1]) - f**2 * xp.cos(f * prob.X[0])
     else:
         raise
 
-    du = prob.eval_f(u, 0).impl
+    if spectral:
+        u = prob.fft.forward(_u)
+    else:
+        u = _u
+
+    _du = prob.eval_f(u, 0).impl
+
+    if spectral:
+        du = prob.fft.backward(_du)
+    else:
+        du = _du
     assert xp.allclose(du, du_expect), 'Got unexpected derivative'
 
-    _u = prob.solve_system(du, factor=1e8, u0=du, t=0) * -1e8
-    assert xp.allclose(_u, u, atol=1e-7), 'Got unexpected inverse derivative'
+    u2 = prob.solve_system(_du, factor=1e8, u0=du, t=0) * -1e8
+
+    if spectral:
+        _u2 = prob.fft.backward(u2)
+    else:
+        _u2 = u2
+
+    assert xp.allclose(_u2, _u, atol=1e-7), 'Got unexpected inverse derivative'
 
 
 if __name__ == '__main__':
-    test_derivative(32, 32, 0, 1, 1)
+    test_derivative(6, 6, 6, 3, False, 1)