Ported Allen-Cahn to generic MPI FFT implementation

pySDC/implementations/problem_classes/AllenCahn_MPIFFT.py

@@ -1,15 +1,11 @@
 import numpy as np
 from mpi4py import MPI
-from mpi4py_fft import PFFT
-
-from pySDC.core.Errors import ProblemError
-from pySDC.core.Problem import ptype
-from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh
 
+from pySDC.implementations.problem_classes.generic_MPIFFT_Laplacian import IMEX_Laplacian_MPIFFT
 from mpi4py_fft import newDistArray
 
 
-class allencahn_imex(ptype):
+class allencahn_imex(IMEX_Laplacian_MPIFFT):
     r"""
     Example implementing the :math:`N`-dimensional Allen-Cahn equation with periodic boundary conditions :math:`u \in [0, 1]^2`
 
@@ -64,68 +60,21 @@ class allencahn_imex(ptype):
     .. [1] https://mpi4py-fft.readthedocs.io/en/latest/
     """
 
-    dtype_u = mesh
-    dtype_f = imex_mesh
-
     def __init__(
         self,
-        nvars=None,
         eps=0.04,
         radius=0.25,
-        spectral=None,
         dw=0.0,
-        L=1.0,
         init_type='circle',
-        comm=MPI.COMM_WORLD,
+        **kwargs,
     ):
-        """Initialization routine"""
-
-        if nvars is None:
-            nvars = (128, 128)
-
-        if not (isinstance(nvars, tuple) and len(nvars) > 1):
-            raise ProblemError('Need at least two dimensions')
-
-        # Creating FFT structure
-        ndim = len(nvars)
-        axes = tuple(range(ndim))
-        self.fft = PFFT(comm, list(nvars), axes=axes, dtype=np.float64, collapse=True)
-
-        # get test data to figure out type and dimensions
-        tmp_u = newDistArray(self.fft, spectral)
-
-        # invoke super init, passing the communicator and the local dimensions as init
-        super().__init__(init=(tmp_u.shape, comm, tmp_u.dtype))
-        self._makeAttributeAndRegister(
-            'nvars', 'eps', 'radius', 'spectral', 'dw', 'L', 'init_type', 'comm', localVars=locals(), readOnly=True
-        )
-
-        L = np.array([self.L] * ndim, dtype=float)
-
-        # get local mesh
-        X = np.ogrid[self.fft.local_slice(False)]
-        N = self.fft.global_shape()
-        for i in range(len(N)):
-            X[i] = X[i] * L[i] / N[i]
-        self.X = [np.broadcast_to(x, self.fft.shape(False)) for x in X]
-
-        # get local wavenumbers and Laplace operator
-        s = self.fft.local_slice()
-        N = self.fft.global_shape()
-        k = [np.fft.fftfreq(n, 1.0 / n).astype(int) for n in N[:-1]]
-        k.append(np.fft.rfftfreq(N[-1], 1.0 / N[-1]).astype(int))
-        K = [ki[si] for ki, si in zip(k, s)]
-        Ks = np.meshgrid(*K, indexing='ij', sparse=True)
-        Lp = 2 * np.pi / L
-        for i in range(ndim):
-            Ks[i] = (Ks[i] * Lp[i]).astype(float)
-        K = [np.broadcast_to(k, self.fft.shape(True)) for k in Ks]
-        K = np.array(K).astype(float)
-        self.K2 = np.sum(K * K, 0, dtype=float)
-
-        # Need this for diagnostics
-        self.dx = self.L / nvars[0]
-        self.dy = self.L / nvars[1]
+        kwargs['L'] = kwargs.get('L', 1.0)
+        super().__init__(alpha=1.0, dtype=np.dtype('float'), **kwargs)
+        self._makeAttributeAndRegister('eps', 'radius', 'dw', 'init_type', localVars=locals(), readOnly=True)
+
+    def _eval_explicit_part(self, u, t, f_expl):
+        f_expl[:] = -2.0 / self.eps**2 * u * (1.0 - u) * (1.0 - 2.0 * u) - 6.0 * self.dw * u * (1.0 - u)
+        return f_expl
 
     def eval_f(self, u, t):
         """
@@ -146,56 +95,24 @@ def eval_f(self, u, t):
 
         f = self.dtype_f(self.init)
 
+        f.impl[:] = self._eval_Laplacian(u, f.impl)
+
         if self.spectral:
             f.impl = -self.K2 * u
 
             if self.eps > 0:
                 tmp = self.fft.backward(u)
-                tmpf = -2.0 / self.eps**2 * tmp * (1.0 - tmp) * (1.0 - 2.0 * tmp) - 6.0 * self.dw * tmp * (1.0 - tmp)
-                f.expl[:] = self.fft.forward(tmpf)
+                tmp[:] = self._eval_explicit_part(tmp, t, tmp)
+                f.expl[:] = self.fft.forward(tmp)
 
         else:
-            u_hat = self.fft.forward(u)
-            lap_u_hat = -self.K2 * u_hat
-            f.impl[:] = self.fft.backward(lap_u_hat, f.impl)
 
             if self.eps > 0:
-                f.expl = -2.0 / self.eps**2 * u * (1.0 - u) * (1.0 - 2.0 * u) - 6.0 * self.dw * u * (1.0 - u)
+                f.expl[:] = self._eval_explicit_part(u, t, f.expl)
 
+        self.work_counters['rhs']()
         return f
 
-    def solve_system(self, rhs, factor, u0, t):
-        """
-        Simple FFT solver for the diffusion part.
-
-        Parameters
-        ----------
-        rhs : dtype_f
-            Right-hand side for the linear system.
-        factor : float
-            Abbrev. for the node-to-node stepsize (or any other factor required).
-        u0 : dtype_u
-            Initial guess for the iterative solver (not used here so far).
-        t : float
-            Current time (e.g. for time-dependent BCs).
-
-        Returns
-        -------
-        me : dtype_u
-            The solution as mesh.
-        """
-
-        if self.spectral:
-            me = rhs / (1.0 + factor * self.K2)
-
-        else:
-            me = self.dtype_u(self.init)
-            rhs_hat = self.fft.forward(rhs)
-            rhs_hat /= 1.0 + factor * self.K2
-            me[:] = self.fft.backward(rhs_hat)
-
-        return me
-
     def u_exact(self, t):
         r"""
         Routine to compute the exact solution at time :math:`t`.
@@ -289,8 +206,9 @@ def eval_f(self, u, t):
 
         f = self.dtype_f(self.init)
 
+        f.impl[:] = self._eval_Laplacian(u, f.impl)
+
         if self.spectral:
-            f.impl = -self.K2 * u
 
             tmp = newDistArray(self.fft, False)
             tmp[:] = self.fft.backward(u, tmp)
@@ -324,9 +242,6 @@ def eval_f(self, u, t):
             f.expl[:] = self.fft.forward(tmpf)
 
         else:
-            u_hat = self.fft.forward(u)
-            lap_u_hat = -self.K2 * u_hat
-            f.impl[:] = self.fft.backward(lap_u_hat, f.impl)
 
             if self.eps > 0:
                 f.expl = -2.0 / self.eps**2 * u * (1.0 - u) * (1.0 - 2.0 * u)
@@ -353,4 +268,5 @@ def eval_f(self, u, t):
 
             f.expl -= 6.0 * dw * u * (1.0 - u)
 
+        self.work_counters['rhs']()
         return f

pySDC/implementations/problem_classes/Brusselator.py

@@ -2,7 +2,6 @@
 from mpi4py import MPI
 
 from pySDC.implementations.problem_classes.generic_MPIFFT_Laplacian import IMEX_Laplacian_MPIFFT
-from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh
 
 
 class Brusselator(IMEX_Laplacian_MPIFFT):
@@ -23,9 +22,6 @@ class Brusselator(IMEX_Laplacian_MPIFFT):
     .. [1] https://link.springer.com/book/10.1007/978-3-642-05221-7
     """
 
-    dtype_u = mesh
-    dtype_f = imex_mesh
-
     def __init__(self, alpha=0.1, **kwargs):
         """Initialization routine"""
         super().__init__(spectral=False, L=1.0, dtype='d', alpha=alpha, **kwargs)

pySDC/implementations/problem_classes/NonlinearSchroedinger_MPIFFT.py

@@ -4,8 +4,8 @@
 
 from pySDC.core.Errors import ProblemError
 from pySDC.core.Problem import WorkCounter
-from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh
 from pySDC.implementations.problem_classes.generic_MPIFFT_Laplacian import IMEX_Laplacian_MPIFFT
+from pySDC.implementations.datatype_classes.mesh import mesh
 
 
 class nonlinearschroedinger_imex(IMEX_Laplacian_MPIFFT):
@@ -47,9 +47,6 @@ class nonlinearschroedinger_imex(IMEX_Laplacian_MPIFFT):
         Journal of Parallel and Distributed Computing (2019).
     """
 
-    dtype_u = mesh
-    dtype_f = imex_mesh
-
     def __init__(self, c=1.0, **kwargs):
         """Initialization routine"""
         super().__init__(L=2 * np.pi, alpha=1j, dtype='D', **kwargs)

pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py

@@ -91,8 +91,8 @@ def __init__(self, nvars=None, spectral=False, L=2 * np.pi, alpha=1.0, comm=MPI.
         self.K2 = self.xp.sum(K * K, 0, dtype=float)  # Laplacian in spectral space
 
         # Need this for diagnostics
-        self.dx = self.L / nvars[0]
-        self.dy = self.L / nvars[1]
+        self.dx = self.L[0] / nvars[0]
+        self.dy = self.L[1] / nvars[1]
 
         # work counters
         self.work_counters['rhs'] = WorkCounter()
@@ -120,8 +120,8 @@ def eval_f(self, u, t):
 
         if self.spectral:
             tmp = self.fft.backward(u)
-            tmpf = self._eval_explicit_part(tmp, t, tmp)
-            f.expl[:] = self.fft.forward(tmpf)
+            tmp[:] = self._eval_explicit_part(tmp, t, tmp)
+            f.expl[:] = self.fft.forward(tmp)
 
         else:
             f.expl[:] = self._eval_explicit_part(u, t, f.expl)

pySDC/tests/test_projects/test_AC/test_simple_forcing.py

@@ -5,10 +5,18 @@
 
 
 @pytest.mark.mpi4py
-def test_main_serial():
-    from pySDC.projects.AllenCahn_Bayreuth.run_simple_forcing_verification import main, visualize_radii
+@pytest.mark.parametrize('spectral', [True, False])
+@pytest.mark.parametrize('name', ['AC-test-noforce', 'AC-test-constforce', 'AC-test-timeforce'])
+def test_main_serial(name, spectral):
+    from pySDC.projects.AllenCahn_Bayreuth.run_simple_forcing_verification import run_simulation
+
+    run_simulation(name=name, spectral=spectral, nprocs_space=None)
+
+
+@pytest.mark.mpi4py
+def test_visualize_radii():
+    from pySDC.projects.AllenCahn_Bayreuth.run_simple_forcing_verification import visualize_radii
 
-    main()
     visualize_radii()