From ea1ed4844206ea4d01a8142e7f9f70ac51951ab8 Mon Sep 17 00:00:00 2001 From: Thomas Baumann <39156931+brownbaerchen@users.noreply.github.com> Date: Mon, 1 Apr 2024 13:47:23 +0200 Subject: [PATCH] Generic MPI FFT class (#408) * Added generic MPIFFT problem class * Fixes * Generalized to `xp` in preparation for GPUs * Fixes * Ported Allen-Cahn to generic MPI FFT implementation --- .../problem_classes/AllenCahn_MPIFFT.py | 122 ++---------- .../problem_classes/Brusselator.py | 106 +++-------- .../NonlinearSchroedinger_MPIFFT.py | 147 ++------------- .../generic_MPIFFT_Laplacian.py | 177 ++++++++++++++++++ .../test_AC/test_simple_forcing.py | 14 +- 5 files changed, 246 insertions(+), 320 deletions(-) create mode 100644 pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py diff --git a/pySDC/implementations/problem_classes/AllenCahn_MPIFFT.py b/pySDC/implementations/problem_classes/AllenCahn_MPIFFT.py index 17f770d218..f0aa47d89e 100644 --- a/pySDC/implementations/problem_classes/AllenCahn_MPIFFT.py +++ b/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 diff --git a/pySDC/implementations/problem_classes/Brusselator.py b/pySDC/implementations/problem_classes/Brusselator.py index d802e9a7f9..946e676781 100644 --- a/pySDC/implementations/problem_classes/Brusselator.py +++ b/pySDC/implementations/problem_classes/Brusselator.py @@ -1,15 +1,10 @@ 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, WorkCounter -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 Brusselator(ptype): +class Brusselator(IMEX_Laplacian_MPIFFT): r""" Two-dimensional Brusselator from [1]_. This is a reaction-diffusion equation with non-autonomous source term: @@ -27,68 +22,29 @@ class Brusselator(ptype): .. [1] https://link.springer.com/book/10.1007/978-3-642-05221-7 """ - dtype_u = mesh - dtype_f = imex_mesh - - def __init__(self, nvars=None, alpha=0.1, comm=MPI.COMM_WORLD): + def __init__(self, alpha=0.1, **kwargs): """Initialization routine""" - nvars = (128,) * 2 if nvars is None else nvars - L = 1.0 - - if not (isinstance(nvars, tuple) and len(nvars) > 1): - raise ProblemError('Need at least two dimensions') - - # Create FFT structure - self.ndim = len(nvars) - axes = tuple(range(self.ndim)) - self.fft = PFFT( - comm, - list(nvars), - axes=axes, - dtype=np.float64, - collapse=True, - backend='fftw', - ) - - # get test data to figure out type and dimensions - tmp_u = newDistArray(self.fft, False) + super().__init__(spectral=False, L=1.0, dtype='d', alpha=alpha, **kwargs) # prepare the array with two components - shape = (2,) + tmp_u.shape + shape = (2,) + (self.init[0]) self.iU = 0 self.iV = 1 + self.init = (shape, self.comm, np.dtype('float')) + + def _eval_explicit_part(self, u, t, f_expl): + iU, iV = self.iU, self.iV + x, y = self.X[0], self.X[1] - super().__init__(init=(shape, comm, tmp_u.dtype)) - self._makeAttributeAndRegister('nvars', 'alpha', 'L', 'comm', localVars=locals(), readOnly=True) - - L = np.array([self.L] * self.ndim, dtype=float) - - # get local mesh for distributed FFT - 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(self.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] - - self.work_counters['rhs'] = WorkCounter() + # evaluate time independent part + f_expl[iU, ...] = 1.0 + u[iU] ** 2 * u[iV] - 4.4 * u[iU] + f_expl[iV, ...] = 3.4 * u[iU] - u[iU] ** 2 * u[iV] + + # add time-dependent part + if t >= 1.1: + mask = (x - 0.3) ** 2 + (y - 0.6) ** 2 <= 0.1**2 + f_expl[iU][mask] += 5.0 + return f_expl def eval_f(self, u, t): """ @@ -106,25 +62,13 @@ def eval_f(self, u, t): f : dtype_f The right-hand side of the problem. """ - iU, iV = self.iU, self.iV - x, y = self.X[0], self.X[1] - f = self.dtype_f(self.init) # evaluate Laplacian to be solved implicitly for i in [self.iU, self.iV]: - u_hat = self.fft.forward(u[i, ...]) - lap_u_hat = -self.alpha * self.K2 * u_hat - f.impl[i, ...] = self.fft.backward(lap_u_hat, f.impl[i, ...]) + f.impl[i, ...] = self._eval_Laplacian(u[i], f.impl[i]) - # evaluate time independent part - f.expl[iU, ...] = 1.0 + u[iU] ** 2 * u[iV] - 4.4 * u[iU] - f.expl[iV, ...] = 3.4 * u[iU] - u[iU] ** 2 * u[iV] - - # add time-dependent part - if t >= 1.1: - mask = (x - 0.3) ** 2 + (y - 0.6) ** 2 <= 0.1**2 - f.expl[iU][mask] += 5.0 + f.expl[:] = self._eval_explicit_part(u, t, f.expl) self.work_counters['rhs']() @@ -153,9 +97,7 @@ def solve_system(self, rhs, factor, u0, t): me = self.dtype_u(self.init) for i in [self.iU, self.iV]: - rhs_hat = self.fft.forward(rhs[i, ...]) - rhs_hat /= 1.0 + factor * self.K2 * self.alpha - me[i, ...] = self.fft.backward(rhs_hat, me[i, ...]) + me[i, ...] = self._invert_Laplacian(me[i], factor, rhs[i]) return me @@ -184,8 +126,8 @@ def u_exact(self, t, u_init=None, t_init=None): me = self.dtype_u(self.init, val=0.0) if t == 0: - me[iU, ...] = 22.0 * y * (1 - y / self.L) ** (3.0 / 2.0) / self.L - me[iV, ...] = 27.0 * x * (1 - x / self.L) ** (3.0 / 2.0) / self.L + me[iU, ...] = 22.0 * y * (1 - y / self.L[0]) ** (3.0 / 2.0) / self.L[0] + me[iV, ...] = 27.0 * x * (1 - x / self.L[0]) ** (3.0 / 2.0) / self.L[0] else: def eval_rhs(t, u): diff --git a/pySDC/implementations/problem_classes/NonlinearSchroedinger_MPIFFT.py b/pySDC/implementations/problem_classes/NonlinearSchroedinger_MPIFFT.py index 66c96c84ae..70c628ec1a 100644 --- a/pySDC/implementations/problem_classes/NonlinearSchroedinger_MPIFFT.py +++ b/pySDC/implementations/problem_classes/NonlinearSchroedinger_MPIFFT.py @@ -1,18 +1,14 @@ import numpy as np -from scipy.optimize import newton_krylov, root +from scipy.optimize import newton_krylov from scipy.optimize.nonlin import NoConvergence -import scipy.sparse as sp -from mpi4py import MPI -from mpi4py_fft import PFFT from pySDC.core.Errors import ProblemError -from pySDC.core.Problem import ptype, WorkCounter -from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh +from pySDC.core.Problem import WorkCounter +from pySDC.implementations.problem_classes.generic_MPIFFT_Laplacian import IMEX_Laplacian_MPIFFT +from pySDC.implementations.datatype_classes.mesh import mesh -from mpi4py_fft import newDistArray - -class nonlinearschroedinger_imex(ptype): +class nonlinearschroedinger_imex(IMEX_Laplacian_MPIFFT): r""" Example implementing the :math:`N`-dimensional nonlinear Schrödinger equation with periodic boundary conditions @@ -51,130 +47,17 @@ class nonlinearschroedinger_imex(ptype): Journal of Parallel and Distributed Computing (2019). """ - dtype_u = mesh - dtype_f = imex_mesh - - def __init__(self, nvars=None, spectral=False, L=2 * np.pi, c=1.0, comm=MPI.COMM_WORLD): + def __init__(self, c=1.0, **kwargs): """Initialization routine""" - - if nvars is None: - nvars = (128, 128) - - if not L == 2.0 * np.pi: - raise ProblemError(f'Setup not implemented, L has to be 2pi, got {L}') + super().__init__(L=2 * np.pi, alpha=1j, dtype='D', **kwargs) if not (c == 0.0 or c == 1.0): raise ProblemError(f'Setup not implemented, c has to be 0 or 1, got {c}') + self._makeAttributeAndRegister('c', localVars=locals(), readOnly=True) - if not (isinstance(nvars, tuple) and len(nvars) > 1): - raise ProblemError('Need at least two dimensions') - - # Creating FFT structure - self.ndim = len(nvars) - axes = tuple(range(self.ndim)) - self.fft = PFFT(comm, list(nvars), axes=axes, dtype=np.complex128, collapse=True) - - # get test data to figure out type and dimensions - tmp_u = newDistArray(self.fft, spectral) - - L = np.array([L] * self.ndim, dtype=float) - - # invoke super init, passing the communicator and the local dimensions as init - super(nonlinearschroedinger_imex, self).__init__(init=(tmp_u.shape, comm, tmp_u.dtype)) - self._makeAttributeAndRegister('nvars', 'spectral', 'L', 'c', 'comm', localVars=locals(), readOnly=True) - - # 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] * self.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] - K = [ki[si] for ki, si in zip(k, s)] - Ks = np.meshgrid(*K, indexing='ij', sparse=True) - Lp = 2 * np.pi / self.L - for i in range(self.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] - - # work counters - self.work_counters['rhs'] = WorkCounter() - - def eval_f(self, u, t): - """ - Routine to evaluate the right-hand side of the problem. - - Parameters - ---------- - u : dtype_u - Current values of the numerical solution. - t : float - Current time at which the numerical solution is computed. - - Returns - ------- - f : dtype_f - The right-hand side of the problem. - """ - - f = self.dtype_f(self.init) - - if self.spectral: - f.impl = -self.K2 * 1j * u - tmp = self.fft.backward(u) - tmpf = self.ndim * self.c * 2j * np.absolute(tmp) ** 2 * tmp - f.expl[:] = self.fft.forward(tmpf) - - else: - u_hat = self.fft.forward(u) - lap_u_hat = -self.K2 * 1j * u_hat - f.impl[:] = self.fft.backward(lap_u_hat, f.impl) - f.expl = self.ndim * self.c * 2j * np.absolute(u) ** 2 * u - - 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 * 1j) - - else: - me = self.dtype_u(self.init) - rhs_hat = self.fft.forward(rhs) - rhs_hat /= 1.0 + factor * self.K2 * 1j - me[:] = self.fft.backward(rhs_hat) - - return me + def _eval_explicit_part(self, u, t, f_expl): + f_expl[:] = self.ndim * self.c * 2j * self.xp.absolute(u) ** 2 * u + return f_expl def u_exact(self, t, **kwargs): r""" @@ -198,9 +81,9 @@ def u_exact(self, t, **kwargs): def nls_exact_1D(t, x, c): ae = 1.0 / np.sqrt(2.0) * np.exp(1j * t) if c != 0: - u = ae * ((np.cosh(t) + 1j * np.sinh(t)) / (np.cosh(t) - 1.0 / np.sqrt(2.0) * np.cos(x)) - 1.0) + u = ae * ((np.cosh(t) + 1j * np.sinh(t)) / (np.cosh(t) - 1.0 / np.sqrt(2.0) * self.xp.cos(x)) - 1.0) else: - u = np.sin(x) * np.exp(-t * 1j) + u = self.xp.sin(x) * np.exp(-t * 1j) return u @@ -261,13 +144,13 @@ def eval_f(self, u, t): if self.spectral: tmp = self.fft.backward(u) - tmpf = self.ndim * self.c * 2j * np.absolute(tmp) ** 2 * tmp + tmpf = self.ndim * self.c * 2j * self.xp.absolute(tmp) ** 2 * tmp f[:] = -self.K2 * 1j * u + self.fft.forward(tmpf) else: u_hat = self.fft.forward(u) lap_u_hat = -self.K2 * 1j * u_hat - f[:] = self.fft.backward(lap_u_hat) + self.ndim * self.c * 2j * np.absolute(u) ** 2 * u + f[:] = self.fft.backward(lap_u_hat) + self.ndim * self.c * 2j * self.xp.absolute(u) ** 2 * u self.work_counters['rhs']() return f diff --git a/pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py b/pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py new file mode 100644 index 0000000000..9e81378b07 --- /dev/null +++ b/pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py @@ -0,0 +1,177 @@ +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, WorkCounter +from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh + +from mpi4py_fft import newDistArray + + +class IMEX_Laplacian_MPIFFT(ptype): + r""" + Generic base class for IMEX problems using a spectral method to solve the Laplacian implicitly and a possible rest + explicitly. The FFTs are done with``mpi4py-fft`` [1]_. + + Parameters + ---------- + nvars : tuple, optional + Spatial resolution + spectral : bool, optional + If True, the solution is computed in spectral space. + L : float, optional + Denotes the period of the function to be approximated for the Fourier transform. + alpha : float, optional + Multiplicative factor before the Laplacian + comm : MPI.COMM_World + Communicator for parallelisation. + + Attributes + ---------- + fft : PFFT + Object for parallel FFT transforms. + X : mesh-grid + Grid coordinates in real space. + K2 : matrix + Laplace operator in spectral space. + + References + ---------- + .. [1] Lisandro Dalcin, Mikael Mortensen, David E. Keyes. Fast parallel multidimensional FFT using advanced MPI. + Journal of Parallel and Distributed Computing (2019). + """ + + dtype_u = mesh + dtype_f = imex_mesh + + xp = np + + def __init__(self, nvars=None, spectral=False, L=2 * np.pi, alpha=1.0, comm=MPI.COMM_WORLD, dtype='d'): + """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 for distributed FFTs') + + # Creating FFT structure + self.ndim = len(nvars) + axes = tuple(range(self.ndim)) + self.fft = PFFT(comm, list(nvars), axes=axes, dtype=dtype, collapse=True) + + # get test data to figure out type and dimensions + tmp_u = newDistArray(self.fft, spectral) + + L = np.array([L] * self.ndim, dtype=float) + + # 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', 'spectral', 'L', 'alpha', 'comm', localVars=locals(), readOnly=True) + + # get local mesh + X = self.xp.ogrid[self.fft.local_slice(False)] + N = self.fft.global_shape() + for i in range(len(N)): + X[i] = 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 + 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] + K = [ki[si] for ki, si in zip(k, s)] + Ks = self.xp.meshgrid(*K, indexing='ij', sparse=True) + Lp = 2 * np.pi / self.L + for i in range(self.ndim): + 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() + + def eval_f(self, u, t): + """ + Routine to evaluate the right-hand side of the problem. + + Parameters + ---------- + u : dtype_u + Current values of the numerical solution. + t : float + Current time at which the numerical solution is computed. + + Returns + ------- + f : dtype_f + The right-hand side of the problem. + """ + + f = self.dtype_f(self.init) + + f.impl[:] = self._eval_Laplacian(u, f.impl) + + if self.spectral: + tmp = self.fft.backward(u) + 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) + + self.work_counters['rhs']() + return f + + def _eval_Laplacian(self, u, f_impl): + if self.spectral: + f_impl[:] = -self.alpha * self.K2 * u + else: + u_hat = self.fft.forward(u) + lap_u_hat = -self.alpha * self.K2 * u_hat + f_impl[:] = self.fft.backward(lap_u_hat, f_impl) + return f_impl + + def _eval_explicit_part(self, u, t, f_expl): + return f_expl + + 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. + """ + me = self.dtype_u(self.init) + me[:] = self._invert_Laplacian(me, factor, rhs) + + return me + + def _invert_Laplacian(self, me, factor, rhs): + if self.spectral: + me[:] = rhs / (1.0 + factor * self.alpha * self.K2) + + else: + rhs_hat = self.fft.forward(rhs) + rhs_hat /= 1.0 + factor * self.alpha * self.K2 + me[:] = self.fft.backward(rhs_hat) + return me diff --git a/pySDC/tests/test_projects/test_AC/test_simple_forcing.py b/pySDC/tests/test_projects/test_AC/test_simple_forcing.py index 0c2a243ce9..c44589d5ca 100644 --- a/pySDC/tests/test_projects/test_AC/test_simple_forcing.py +++ b/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()