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