Coverage for pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py: 100%
89 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-25 20:28 +0000
1import numpy as np
2from mpi4py import MPI
3from mpi4py_fft import newDistArray
5from pySDC.helpers.fft_helper import PFFT
7from pySDC.core.errors import ProblemError
8from pySDC.core.problem import Problem, WorkCounter
9from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh
12class IMEX_Laplacian_MPIFFT(Problem):
13 r"""
14 Generic base class for IMEX problems using a spectral method to solve the Laplacian implicitly and a possible rest
15 explicitly. The FFTs are done with``mpi4py-fft`` [1]_.
16 Works in two and three dimensions.
18 Parameters
19 ----------
20 nvars : tuple, optional
21 Spatial resolution
22 spectral : bool, optional
23 If True, the solution is computed in spectral space.
24 L : float, optional
25 Denotes the period of the function to be approximated for the Fourier transform.
26 alpha : float, optional
27 Multiplicative factor before the Laplacian
28 comm : MPI.COMM_World
29 Communicator for parallelisation.
31 Attributes
32 ----------
33 fft : PFFT
34 Object for parallel FFT transforms.
35 X : mesh-grid
36 Grid coordinates in real space.
37 K2 : matrix
38 Laplace operator in spectral space.
40 References
41 ----------
42 .. [1] Lisandro Dalcin, Mikael Mortensen, David E. Keyes. Fast parallel multidimensional FFT using advanced MPI.
43 Journal of Parallel and Distributed Computing (2019).
44 """
46 dtype_u = mesh
47 dtype_f = imex_mesh
49 xp = np
50 fft_backend = 'fftw'
51 fft_comm_backend = 'MPI'
53 def setup_GPU(self):
54 """switch to GPU modules"""
55 import cupy as cp
56 from pySDC.implementations.datatype_classes.cupy_mesh import cupy_mesh, imex_cupy_mesh
58 self.xp = cp
60 self.dtype_u = cupy_mesh
61 self.dtype_f = imex_cupy_mesh
63 self.fft_backend = 'cupy'
64 self.fft_comm_backend = 'NCCL'
66 def __init__(
67 self, nvars=None, spectral=False, L=2 * np.pi, alpha=1.0, comm=MPI.COMM_WORLD, dtype='d', useGPU=False, x0=0.0
68 ):
69 if useGPU:
70 self.setup_GPU()
72 if nvars is None:
73 nvars = (128, 128)
75 if not (isinstance(nvars, tuple) and len(nvars) > 1):
76 raise ProblemError('Need at least two dimensions for distributed FFTs')
78 # Creating FFT structure
79 self.ndim = len(nvars)
80 axes = tuple(range(self.ndim))
81 self.fft = PFFT(
82 comm,
83 list(nvars),
84 axes=axes,
85 dtype=dtype,
86 collapse=True,
87 backend=self.fft_backend,
88 comm_backend=self.fft_comm_backend,
89 grid=(-1,),
90 )
92 # get test data to figure out type and dimensions
93 tmp_u = newDistArray(self.fft, spectral)
95 L = np.array([L] * self.ndim, dtype=float)
97 # invoke super init, passing the communicator and the local dimensions as init
98 super().__init__(init=(tmp_u.shape, comm, tmp_u.dtype))
99 self._makeAttributeAndRegister(
100 'nvars', 'spectral', 'L', 'alpha', 'comm', 'x0', 'useGPU', localVars=locals(), readOnly=True
101 )
103 self.getLocalGrid()
104 self.getLaplacian()
106 # Need this for diagnostics
107 self.dx = self.L[0] / nvars[0]
108 self.dy = self.L[1] / nvars[1]
110 # work counters
111 self.work_counters['rhs'] = WorkCounter()
113 def getLocalGrid(self):
114 X = list(self.xp.ogrid[self.fft.local_slice(False)])
115 N = self.fft.global_shape()
116 for i in range(len(N)):
117 X[i] = self.x0 + (X[i] * self.L[i] / N[i])
118 self.X = [self.xp.broadcast_to(x, self.fft.shape(False)) for x in X]
120 def getLaplacian(self):
121 s = self.fft.local_slice()
122 N = self.fft.global_shape()
123 k = [self.xp.fft.fftfreq(n, 1.0 / n).astype(int) for n in N]
124 K = [ki[si] for ki, si in zip(k, s, strict=True)]
125 Ks = list(self.xp.meshgrid(*K, indexing='ij', sparse=True))
126 Lp = 2 * np.pi / self.L
127 for i in range(self.ndim):
128 Ks[i] = (Ks[i] * Lp[i]).astype(float)
129 K = [self.xp.broadcast_to(k, self.fft.shape(True)) for k in Ks]
130 K = self.xp.array(K).astype(float)
131 self.K2 = self.xp.sum(K * K, 0, dtype=float)
133 def eval_f(self, u, t):
134 """
135 Routine to evaluate the right-hand side of the problem.
137 Parameters
138 ----------
139 u : dtype_u
140 Current values of the numerical solution.
141 t : float
142 Current time at which the numerical solution is computed.
144 Returns
145 -------
146 f : dtype_f
147 The right-hand side of the problem.
148 """
150 f = self.dtype_f(self.init)
152 f.impl[:] = self._eval_Laplacian(u, f.impl)
154 if self.spectral:
155 tmp = self.fft.backward(u)
156 tmp[:] = self._eval_explicit_part(tmp, t, tmp)
157 f.expl[:] = self.fft.forward(tmp)
159 else:
160 f.expl[:] = self._eval_explicit_part(u, t, f.expl)
162 self.work_counters['rhs']()
163 return f
165 def _eval_Laplacian(self, u, f_impl, alpha=None):
166 alpha = alpha if alpha else self.alpha
167 if self.spectral:
168 f_impl[:] = -alpha * self.K2 * u
169 else:
170 u_hat = self.fft.forward(u)
171 lap_u_hat = -alpha * self.K2 * u_hat
172 f_impl[:] = self.fft.backward(lap_u_hat, f_impl)
173 return f_impl
175 def _eval_explicit_part(self, u, t, f_expl):
176 return f_expl
178 def solve_system(self, rhs, factor, u0, t):
179 """
180 Simple FFT solver for the diffusion part.
182 Parameters
183 ----------
184 rhs : dtype_f
185 Right-hand side for the linear system.
186 factor : float
187 Abbrev. for the node-to-node stepsize (or any other factor required).
188 u0 : dtype_u
189 Initial guess for the iterative solver (not used here so far).
190 t : float
191 Current time (e.g. for time-dependent BCs).
193 Returns
194 -------
195 me : dtype_u
196 The solution as mesh.
197 """
198 me = self.dtype_u(self.init)
199 me[:] = self._invert_Laplacian(me, factor, rhs)
201 return me
203 def _invert_Laplacian(self, me, factor, rhs, alpha=None):
204 alpha = alpha if alpha else self.alpha
205 if self.spectral:
206 me[:] = rhs / (1.0 + factor * alpha * self.K2)
208 else:
209 rhs_hat = self.fft.forward(rhs)
210 rhs_hat /= 1.0 + factor * alpha * self.K2
211 me[:] = self.fft.backward(rhs_hat)
212 return me