Coverage for pySDC/helpers/multigrid.py: 99%
188 statements
« prev ^ index » next coverage.py v7.17.3, created at 2026-11-12 23:49 +0000
1r"""
2Geometric multigrid for :math:`(I - s A_h)\, x = b` with :math:`A_h = \nu \Delta_h - c \cdot \nabla_h`, on a CPU
3(NumPy) or a GPU (CuPy).
5This is the system an implicit Euler step, and so every node solve of SDC, of an advection-diffusion term
6:math:`\nu \Delta u - c \cdot \nabla u` leads to, with :math:`s` the factor of the step. :math:`\Delta_h` is the
7second-order Laplacian (3, 5 or 7 points) and :math:`\nabla_h` the first-order upwind gradient, each component of
8:math:`c` differenced on the side it comes from, on :math:`N` points per direction of an interval, square or cube of
9side length :math:`L`. The grid is periodic (``bc='periodic'``, spacing :math:`h = L / N`) or has zero Dirichlet
10values (``bc='dirichlet-zero'``: :math:`N` interior points of a vertex-centered grid, :math:`h = L / (N + 1)`). Both
11are the grids, and :math:`A_h` the matrix, of pySDC's finite differences
12(:func:`~pySDC.helpers.problem_helper.get_finite_difference_matrix` with ``order=2, stencil_type='center'`` for
13:math:`\Delta_h` and ``order=1, stencil_type='upwind'`` for positive velocities). Without a velocity it is pure
14diffusion, :math:`A_h = \nu \Delta_h`. :class:`GeometricMultigrid` takes and returns plain arrays and does not depend
15on the rest of pySDC; :class:`~pySDC.implementations.problem_classes.generic_ND_FD.GenericNDimFinDiff` uses it for
16``solver_type='MG'``.
18**The method.** V(2,2)-cycles with weighted Jacobi smoothing (:math:`\omega = 0.8`, with the diagonal of the whole
19operator), full-weighting restriction, linear interpolation, the same stencil on every grid, and an exact solve with a
20dense inverse on the coarsest grid, which has at most ``coarsest`` points (default ``COARSEST``, 1000). A periodic grid
21coarsens from :math:`N` to :math:`N / 2` points per direction, a Dirichlet grid from :math:`N` to :math:`(N - 1) / 2`
22interior points, with zero boundary values in the transfers. The exact coarse solve is needed: the smooth modes of
23:math:`I - s \Delta_h` are damped only by the identity part, which Jacobi does not do. A solve starts from zero and
24stops when the 2-norm of the residual is below ``rtol`` times that of ``b``, or when a V-cycle no longer halves it
25(``stall``), which is where the floating-point format of ``b`` ends (about :math:`\kappa u`, with :math:`\kappa` the
26condition number and :math:`u` the unit roundoff), or after ``maxiter`` cycles. A cycle that does not reduce the
27residual at all is discarded, so that a format too coarse for the system returns its best iterate, zero in the worst
28case, rather than one that is worse. Both rules judge a cycle against the one before from the second cycle on: the first
29cycle is kept unless it overflows, and a solve that ends worse than zero returns zero.
31**Advection.** Point Jacobi smooths reliably while diffusion dominates on every grid, and the grid Peclet number
32:math:`|c_a| H / \nu` is largest on the coarsest grid actually used (spacing :math:`H`, set by ``coarsest``), which is
33what bounds it; at ``coarsest=15``, a 1D grid has :math:`H = 1/16` whatever :math:`N`. Measured with random
34right-hand sides in 1D to 3D, periodic and Dirichlet, with stiffness :math:`s \nu h^{-2}` up to :math:`10^6`, every
35solve reaches ``rtol`` while that Peclet number is at most about 1 to 2. Beyond, a V-cycle reduces the residual by
36only 0.5 to 0.85 in 2D and 3D (the coarse grids, upwinded on their own spacing, approximate the smooth components
37along the flow badly), the stall rule takes that for the end of the format, and the solve stops after a few cycles
38with a residual of :math:`10^{-6}` to :math:`10^{-2}`.
40Raising ``stall`` (to 0.8 or 0.9) lets such slowly contracting solves go on to ``rtol`` (in 35 V-cycles at a
41coarsest-grid Peclet number of 16 in 2D, at the price of about one more cycle at the end of a low format), and is
42needed near the limit. Every solve leaves the relative residual of what it returns in ``relative_residual``: check it
43with advection.
45The first V-cycle is judged by the second, which matters for advection: with a smooth right-hand side and a stiff system
46(:math:`s \nu h^{-2} \gtrsim 10^4`), the first cycle from zero leaves its error in the outflow boundary layer, where the
47residual is large, and reduces the 2-norm of the residual by only 0.5 to 0.9 or increases it (up to four times
48measured), even at coarsest-grid Peclet numbers far below 1 (0.03 at :math:`N = 1023` and :math:`H = 1/512`), in 1D as
49in 2D. The cycles after it contract by 0.15 to 0.7, and the solve converges as usual.
51**Formats.** A solve computes in the floating-point type of its right-hand side: ``float64`` and ``float32`` arrays
52are solved in their own arithmetic, and ``float16`` is a storage format, as on GPUs: every array is stored in half
53precision, every operation computes in single precision. The result has the type of the right-hand side.
55**Speed.** On a GPU, periodic 3D grids run on fused CUDA kernels: one pass over the arrays for the residual, the Jacobi
56step, the restriction and the interpolation each, with the grid size compiled in. With ``fma=False`` (default), the
57stencil kernels compute bit for bit what the array operations compute, which is what the tests check; ``fma=True`` lets
58the compiler contract multiplications and additions, which is faster where the arithmetic and not the memory limits the
59kernels (fp64 on GPUs with few fp64 units). Measured at :math:`640^3` with V(2,2)-cycles, one Jacobi step moves data at
60about 3 TB/s on an H100 in fp64. The fused kernels are for periodic pure diffusion; 1D and 2D grids, Dirichlet grids,
61advection, and every grid on a CPU use array operations.
62"""
64import re
66import numpy as np
68OMEGA = 0.8 # Jacobi weight
69COARSEST = 1000 # default: grids with at most this many points are solved exactly
70STALL = 0.5 # default: a V-cycle that reduces the residual by less than this is at the accuracy of the format
73def compute_type(dtype):
74 """The type a format computes in: its own, but float32 for float16, which is a storage format here."""
75 return np.promote_types(dtype, np.float32)
78class _Ops:
79 """
80 The multigrid operations on one backend ``xp``: array operations everywhere, fused CuPy kernels for 3D grids on a
81 GPU. Both compute in :func:`compute_type` and store in the type of the input.
82 """
84 # the CUDA types of the stored arrays and of their arithmetic
85 _TYPES = {
86 np.dtype(np.float64): ('double', 'double'),
87 np.dtype(np.float32): ('float', 'float'),
88 np.dtype(np.float16): ('__half', 'float'),
89 }
91 def __init__(self, xp, fma=False, periodic=True):
92 self.xp = xp
93 self.gpu = xp is not np
94 self.periodic = periodic
95 self.start = 0 if periodic else 1 # the fine index of the first coarse point
96 # off, the kernels compute bit for bit what the array operations do; on, they are faster where arithmetic
97 # limits them
98 self.options = () if fma else ('--fmad=false',)
99 self._compiled = {}
100 # planes per thread in the streaming stencils: on an H100 at 640^3, 16 to 64 are equally fast, longer sections
101 # lose up to 25 % (too few thread blocks for even waves), shorter ones re-read more planes
102 self.section = 64
104 def fused(self, x):
105 """Whether ``x`` is handled by the fused kernels: a periodic 3D grid of a supported type on a GPU."""
106 return self.gpu and self.periodic and x.ndim == 3 and x.dtype in self._TYPES
108 @staticmethod
109 def omega_over_d(x, sh, q=None):
110 r"""
111 :math:`\omega / (1 + 2 d\, s \nu h^{-2} + \sum_a |s c_a h^{-1}|)`, the Jacobi factor for ``d``-dimensional
112 ``x``, with ``sh`` :math:`= s \nu h^{-2}` and ``q`` the :math:`s c_a h^{-1}`, computed once on the host in the
113 compute type: no division per point, which costs real time where fp64 arithmetic is slow.
114 """
115 ct = compute_type(x.dtype).type
116 diagonal = ct(1) + ct(2 * x.ndim) * ct(sh)
117 if q:
118 diagonal = diagonal + ct(sum(abs(v) for v in q))
119 return ct(OMEGA) / diagonal
121 def shift(self, y, k, axis):
122 """``y`` moved by ``k`` = 1 or -1 along ``axis`` like ``numpy.roll``: periodic or with zero boundary values."""
123 if self.periodic:
124 return self.xp.roll(y, k, axis)
125 out = self.xp.zeros_like(y)
126 inner, before = (slice(1, None), slice(None, -1)), (slice(None),) * axis
127 destination, source = inner if k == 1 else inner[::-1]
128 out[before + (destination,)] = y[before + (source,)]
129 return out
131 # --- the operations ---------------------------------------------------------------------------------------------
132 def stencil(self, kind, x, b, p1, p2=0.0, out=None, q=None):
133 r"""
134 A stencil operation with :math:`B x = p_1 \Delta x - \sum_a q_a D_a x`, ``p1`` the scale of the Laplacian
135 (unscaled :math:`\Delta`) and ``q`` (or None) the signed scales of the first-order upwind differences
136 :math:`D_a`: 0 gives :math:`B x`, 1 the residual :math:`b - (x - B x)`, 2 the Jacobi step
137 :math:`x + p_2 (b - (x - B x))`, with ``p2`` from :meth:`omega_over_d`.
138 """
139 if q is None and self.fused(x):
140 return self._stencil_fused(kind, x, b, p1, p2, out)
141 ct = compute_type(x.dtype)
142 y = x.astype(ct, copy=False)
143 s = sum(self.shift(y, shift, axis) for axis in range(y.ndim) for shift in (1, -1))
144 lap = s - (2 * y.ndim) * y
145 result = ct.type(p1) * lap
146 for axis, qa in enumerate(q or ()):
147 if qa: # q_a D_a x = |q_a| (x_i - x_upwind), the upwind neighbour on the side the velocity comes from
148 result = result - ct.type(abs(qa)) * (y - self.shift(y, 1 if qa > 0 else -1, axis))
149 if kind != 0:
150 result = b.astype(ct, copy=False) - (y - result)
151 if kind == 2:
152 result = y + ct.type(p2) * result
153 if out is None:
154 return result.astype(x.dtype, copy=False)
155 out[...] = result
156 return out
158 def residual(self, x, b, s, h2inv, q=None):
159 return self.stencil(1, x, b, s * h2inv, q=q)
161 def jacobi(self, x, b, s, h2inv, q=None):
162 """One weighted Jacobi step, out of place."""
163 return self.stencil(2, x, b, s * h2inv, self.omega_over_d(x, s * h2inv, q), q=q)
165 def jacobi_from_zero(self, b, s, h2inv, q=None):
166 """The Jacobi step from x = 0, which is (omega / d) b: one pass over b instead of a stencil over zeros."""
167 wd = self.omega_over_d(b, s * h2inv, q)
168 if self.fused(b):
169 return self._kernel('jacobi_zero', b)(b, self._scalar(b, wd), self.xp.empty_like(b))
170 return (wd * b.astype(compute_type(b.dtype), copy=False)).astype(b.dtype, copy=False)
172 def restrict(self, r):
173 """Full weighting, [1/4, 1/2, 1/4] in every direction, onto every second point."""
174 if self.fused(r):
175 n = r.shape[0] // 2
176 return self._kernel('restrict', r, n)(r, self.xp.empty((n, n, n), dtype=r.dtype), size=n**3)
177 y = r.astype(compute_type(r.dtype), copy=False)
178 for axis in range(r.ndim):
179 y = 0.25 * self.shift(y, 1, axis) + 0.5 * y + 0.25 * self.shift(y, -1, axis)
180 return y[(slice(self.start, None, 2),) * r.ndim].astype(r.dtype)
182 def prolong_add(self, x, e):
183 """x += linear interpolation of the coarse correction e, in place."""
184 if self.fused(x):
185 self._kernel('prolong', x)(e, x, size=e.size)
186 return x
187 fine = self.xp.zeros(x.shape, dtype=compute_type(x.dtype))
188 fine[(slice(self.start, None, 2),) * x.ndim] = e
189 for axis in range(x.ndim):
190 fine = 0.5 * self.shift(fine, 1, axis) + fine + 0.5 * self.shift(fine, -1, axis)
191 x[...] = x.astype(fine.dtype, copy=False) + fine
192 return x
194 # --- fused CuPy kernels for 3D grids ----------------------------------------------------------------------------
195 @staticmethod
196 def _scalar(x, value):
197 """A kernel scalar for arrays like x: of their compute type."""
198 return compute_type(x.dtype).type(value)
200 _STENCIL = r"""
201 #include <cuda_fp16.h>
202 template <typename CT, typename T> __device__ __forceinline__ CT ld(const T* p, int i) { return CT(p[i]); }
203 template <> __device__ __forceinline__ float ld<float, __half>(const __half* p, int i) {
204 return __half2float(p[i]);
205 }
206 template <typename T, typename CT> __device__ __forceinline__ T st(CT v) { return T(v); }
207 template <> __device__ __forceinline__ __half st<__half, float>(float v) { return __float2half_rn(v); }
209 // KIND 0: y = p1 lap x; 1: residual y = bb - (x - p1 lap x); 2: Jacobi step with p2 = omega / d.
210 // Every thread owns one column (b, c) of a section of L_ planes and walks along a, keeping the planes a - 1, a,
211 // a + 1 in registers, so that x is read about once from memory (twice at the ends of a section); the sections
212 // give the GPU enough threads to keep enough loads in flight. The operations are those of the array version, in
213 // the same order.
214 template <typename T, typename CT, int KIND>
215 __global__ void stencil(const T* __restrict__ x, const T* __restrict__ bb, T* __restrict__ y, CT p1, CT p2) {
216 const int n = N_, nn = n * n;
217 const int c = blockIdx.x * blockDim.x + threadIdx.x, b = blockIdx.y * blockDim.y + threadIdx.y;
218 if (b >= n || c >= n) return;
219 const int rest = b * n + c;
220 const int bm = (b == 0 ? n - 1 : b - 1) * n + c, bp = (b == n - 1 ? 0 : b + 1) * n + c;
221 const int cm = b * n + (c == 0 ? n - 1 : c - 1), cp = b * n + (c == n - 1 ? 0 : c + 1);
222 const int first = blockIdx.z * L_, last = min(first + L_, n);
223 CT prev = ld<CT>(x, (first == 0 ? n - 1 : first - 1) * nn + rest), cur = ld<CT>(x, first * nn + rest);
224 for (int a = first; a < last; a++) {
225 const int plane = a * nn;
226 const CT next = ld<CT>(x, (a == n - 1 ? 0 : plane + nn) + rest);
227 CT s = prev + next;
228 s = s + ld<CT>(x, plane + bm);
229 s = s + ld<CT>(x, plane + bp);
230 s = s + ld<CT>(x, plane + cm);
231 s = s + ld<CT>(x, plane + cp);
232 const CT xk = cur;
233 const CT lap = s - CT(6) * xk;
234 const int k = plane + rest;
235 if (KIND == 0) {
236 y[k] = st<T>(lap * p1);
237 } else if (KIND == 1) {
238 y[k] = st<T>(ld<CT>(bb, k) - (xk - p1 * lap));
239 } else {
240 y[k] = st<T>(xk + p2 * (ld<CT>(bb, k) - (xk - p1 * lap)));
241 }
242 prev = cur;
243 cur = next;
244 }
245 }
246 """
248 def _stencil_fused(self, kind, x, b, p1, p2, out):
249 """
250 The stencil on a GPU, in its fastest form for the type, measured on an H100 at 640^3: one thread per point
251 for fp64 (2.1 ms for a Jacobi step, against 2.8 ms streaming), streaming columns for fp32 (1.24 against
252 1.38 ms) and fp16 (0.86 against 1.28 ms). Both forms compute the same, bit for bit.
253 """
254 n = x.shape[0]
255 out = self.xp.empty_like(x) if out is None else out
256 if x.dtype == np.float64:
257 name = ('lap', 'residual', 'jacobi')[kind]
258 scalars = (self._scalar(x, p1),) if kind != 2 else (self._scalar(x, p1), self._scalar(x, p2))
259 return self._kernel(name, x)(x, *(() if kind == 0 else (b,)), *scalars, out)
260 section = min(n, self.section)
261 if ('stencil', n, section) not in self._compiled:
262 expressions = [f'stencil<{t},{ct},{k}>' for t, ct in self._TYPES.values() for k in range(3)]
263 self._compiled['stencil', n, section] = self.xp.RawModule(
264 code=self._STENCIL,
265 options=(*self.options, f'-DN_={n}', f'-DL_={section}'),
266 name_expressions=expressions,
267 )
268 t, ct = self._TYPES[x.dtype]
269 kernel = self._compiled['stencil', n, section].get_function(f'stencil<{t},{ct},{kind}>')
270 threads = (32, 8)
271 blocks = ((n + threads[0] - 1) // threads[0], (n + threads[1] - 1) // threads[1], (n + section - 1) // section)
272 kernel(blocks, threads, (x, x if b is None else b, out, self._scalar(x, p1), self._scalar(x, p2)))
273 return out
275 def _kernel(self, name, x, n=None):
276 r"""
277 The fused kernel ``name`` for arrays like ``x`` on grids of :math:`n^3` points (default: the size of ``x``),
278 compiled once per grid size and type.
280 The grid size is a compile-time constant: the index arithmetic then divides by constants, which the compiler
281 turns into multiplications, instead of by a size read at run time, which is slow on a GPU and does not shrink
282 with the format, so that it would bound a kernel in fp32 rather than its memory traffic. Arithmetic is in the
283 type ``CT``: the array's own type, but float32 for float16.
284 """
285 n = x.shape[0] if n is None else n
286 half = x.dtype == np.float16
287 if (name, n, half) not in self._compiled:
288 neighbours = """
289 const int n = N_, nn = n * n;
290 int k = i, a = k / nn, b = (k / n) % n, c = k % n;
291 int am = (a == 0 ? n - 1 : a - 1) * nn, ap = (a == n - 1 ? 0 : a + 1) * nn;
292 int bm = (b == 0 ? n - 1 : b - 1) * n, bp = (b == n - 1 ? 0 : b + 1) * n;
293 int cm = c == 0 ? n - 1 : c - 1, cp = c == n - 1 ? 0 : c + 1;
294 int rest = b * n + c, plane = a * nn;
295 CT s = CT(x[am + rest]) + CT(x[ap + rest]);
296 s = s + CT(x[plane + bm + c]);
297 s = s + CT(x[plane + bp + c]);
298 s = s + CT(x[plane + b * n + cm]);
299 s = s + CT(x[plane + b * n + cp]);
300 CT xk = CT(x[k]);
301 CT lap = s - CT(6) * xk;
302 """
303 sources = {
304 # one thread per point: the fastest in fp64, which reaches the memory bandwidth this way
305 'lap': ('raw T x, S h2inv', 'T y', neighbours + 'y = T(lap * h2inv)'),
306 'residual': ('raw T x, T bb, S sh', 'T r', neighbours + 'r = T(CT(bb) - (xk - sh * lap))'),
307 'jacobi': (
308 'raw T x, T bb, S sh, S wd',
309 'T y',
310 neighbours + 'y = T(xk + wd * (CT(bb) - (xk - sh * lap)))',
311 ),
312 'jacobi_zero': ('T bb, S wd', 'T y', 'y = T(wd * CT(bb));'),
313 'restrict': (
314 'raw T r',
315 'raw T out',
316 """
317 const int nc = N_, n = 2 * nc;
318 int k = i, a = k / (nc * nc), b = (k / nc) % nc, c = k % nc;
319 CT acc = CT(0);
320 for (int da = -1; da <= 1; da++) {
321 int fa = (2 * a + da + n) % n; CT wa = da == 0 ? CT(0.5) : CT(0.25);
322 for (int db = -1; db <= 1; db++) {
323 int fb = (2 * b + db + n) % n; CT wb = db == 0 ? CT(0.5) : CT(0.25);
324 for (int dc = -1; dc <= 1; dc++) {
325 int fc = (2 * c + dc + n) % n; CT wc = dc == 0 ? CT(0.5) : CT(0.25);
326 acc = acc + (wa * wb * wc) * CT(r[(fa * n + fb) * n + fc]);
327 }
328 }
329 }
330 out[k] = T(acc);
331 """,
332 ),
333 'prolong': (
334 'raw T e',
335 'raw T x',
336 """
337 // one thread per coarse point (A, B, C): it reads the eight coarse values of its cell once and
338 // adds the interpolation to the eight fine points (2A + da, 2B + db, 2C + dc), with all eight
339 // terms in a fixed order, those absent at its parity with weight zero
340 const int n = N_, nc = n / 2;
341 int k = i, A = k / (nc * nc), B = (k / nc) % nc, C = k % nc;
342 int A1 = A == nc - 1 ? 0 : A + 1, B1 = B == nc - 1 ? 0 : B + 1, C1 = C == nc - 1 ? 0 : C + 1;
343 CT e000 = CT(e[(A * nc + B) * nc + C]), e001 = CT(e[(A * nc + B) * nc + C1]);
344 CT e010 = CT(e[(A * nc + B1) * nc + C]), e011 = CT(e[(A * nc + B1) * nc + C1]);
345 CT e100 = CT(e[(A1 * nc + B) * nc + C]), e101 = CT(e[(A1 * nc + B) * nc + C1]);
346 CT e110 = CT(e[(A1 * nc + B1) * nc + C]), e111 = CT(e[(A1 * nc + B1) * nc + C1]);
347 #pragma unroll
348 for (int da = 0; da < 2; da++) {
349 #pragma unroll
350 for (int db = 0; db < 2; db++) {
351 #pragma unroll
352 for (int dc = 0; dc < 2; dc++) {
353 CT wa = da ? CT(0.5) : CT(1), wb = db ? CT(0.5) : CT(1), wc = dc ? CT(0.5) : CT(1);
354 CT ma = da ? CT(1) : CT(0), mb = db ? CT(1) : CT(0), mc = dc ? CT(1) : CT(0);
355 CT w = wa * wb * wc, v = CT(0);
356 v = v + (w * CT(1)) * e000;
357 v = v + (w * mc) * e001;
358 v = v + (w * mb) * e010;
359 v = v + (w * (mb * mc)) * e011;
360 v = v + (w * ma) * e100;
361 v = v + (w * (ma * mc)) * e101;
362 v = v + (w * (ma * mb)) * e110;
363 v = v + (w * (ma * mb * mc)) * e111;
364 int f = ((2 * A + da) * n + 2 * B + db) * n + 2 * C + dc;
365 x[f] = T(CT(x[f]) + v);
366 }
367 }
368 }
369 """,
370 ),
371 }
372 params, out, body = sources[name]
373 # the scalars in the compute type: CuPy infers type placeholders from arrays only
374 params = re.sub(r'\bS\b', 'float32' if half else 'T', params)
375 self._compiled[name, n, half] = self.xp.ElementwiseKernel(
376 params,
377 out,
378 body.replace('CT', 'float' if half else 'T'),
379 f'mg_{name}_{n}_{"h" if half else "f"}',
380 options=self.options,
381 preamble=f'#define N_ {n}',
382 )
383 return self._compiled[name, n, half]
386class GeometricMultigrid:
387 r"""
388 Geometric multigrid for :math:`(I - s A_h)\, x = b` with :math:`A_h = \nu \Delta_h - c \cdot \nabla_h` on a grid
389 of :math:`N^d` points with side length :math:`L`, periodic or with zero Dirichlet values: :math:`\Delta_h` is the
390 second-order Laplacian and :math:`\nabla_h` the first-order upwind gradient. See :mod:`pySDC.helpers.multigrid`
391 for the method, the formats and the speed.
393 Parameters
394 ----------
395 N : int
396 Points per direction. Periodic: :math:`N = m 2^k` with :math:`m^d` at most ``coarsest``; zero Dirichlet: the
397 interior points, :math:`N = 2^k - 1` (or :math:`m 2^k - 1` with :math:`(m - 1)^d` at most ``coarsest``).
398 ndim : int
399 Dimension: 1, 2 or 3.
400 L : float
401 Side length of the domain.
402 xp : module
403 ``numpy`` or ``cupy``: where the arrays live.
404 fma : bool
405 On a GPU, let the compiler contract multiplications and additions into fused multiply-adds: faster where
406 arithmetic limits the kernels, but no longer bit for bit what the array operations compute.
407 bc : str
408 ``'periodic'`` (grid spacing :math:`h = L / N`) or ``'dirichlet-zero'``: :math:`N` interior points of a
409 vertex-centered grid with :math:`h = L / (N + 1)` and zero values on the boundary.
410 nu : float
411 Diffusion coefficient :math:`\nu`, positive.
412 velocity : sequence of float or None
413 Constant velocity :math:`c`, one entry per dimension; None or zeros give pure diffusion. Each component is
414 upwinded on the side it comes from.
415 coarsest : int
416 The largest number of points the coarsest grid may have, which is solved exactly with a dense inverse: the grid
417 is halved until it has at most this many. Default ``COARSEST`` (1000); smaller values give more levels, and
418 a coarsest grid with a larger spacing, which is what bounds the Peclet number (see
419 :mod:`pySDC.helpers.multigrid`).
420 stall : float
421 The stall rule: a V-cycle that reduces the residual by a factor larger than this ends the solve, taken for
422 the accuracy of the format. Raise it (to 0.9, say) where V-cycles contract slowly, as on advection-dominated
423 coarse grids, so that they go on to ``rtol``; the price is more cycles spent at the accuracy of a low format
424 (fp16, fp32) before the solve notices it.
426 Attributes
427 ----------
428 relative_residual : float
429 The 2-norm of the residual of the iterate the last :meth:`solve` returned, relative to that of ``b``.
431 Examples
432 --------
433 >>> import numpy as np
434 >>> mg = GeometricMultigrid(64, ndim=2)
435 >>> b = np.random.default_rng(0).standard_normal((64, 64))
436 >>> x, cycles = mg.solve(b, s=1e-2, rtol=1e-10, maxiter=50)
437 >>> bool(np.linalg.norm(b - (x - mg.apply(x, 1e-2))) < 1e-10 * np.linalg.norm(b))
438 True
439 >>> mg = GeometricMultigrid(63, ndim=2, bc='dirichlet-zero', nu=0.1, velocity=(1.0, -0.5))
440 >>> b = np.random.default_rng(0).standard_normal((63, 63))
441 >>> x, cycles = mg.solve(b, s=1e-1, rtol=1e-10, maxiter=50)
442 >>> bool(np.linalg.norm(b - (x - mg.apply(x, 1e-1))) < 1e-10 * np.linalg.norm(b))
443 True
444 """
446 def __init__(
447 self, N, ndim=3, L=1.0, xp=np, fma=False, bc='periodic', nu=1.0, velocity=None, stall=STALL, coarsest=COARSEST
448 ):
449 if ndim not in (1, 2, 3):
450 raise ValueError(f'ndim has to be 1, 2 or 3, got {ndim}')
451 if bc not in ('periodic', 'dirichlet-zero'):
452 raise ValueError(f'bc has to be "periodic" or "dirichlet-zero", got {bc}')
453 if velocity is not None:
454 velocity = tuple(float(c) for c in velocity)
455 if len(velocity) != ndim:
456 raise ValueError(f'velocity needs one entry per dimension, {ndim}, got {len(velocity)}')
457 velocity = velocity if any(velocity) else None # pure diffusion keeps the fused kernels
458 periodic = bc == 'periodic'
459 # halve the grid: n -> n / 2 on periodic grids, n -> (n - 1) / 2 interior points with Dirichlet values
460 self.sizes = [N]
461 while self.sizes[-1] ** ndim > coarsest and self.sizes[-1] % 2 == (0 if periodic else 1):
462 self.sizes.append(self.sizes[-1] // 2)
463 if self.sizes[-1] ** ndim > coarsest:
464 form = f'm 2^k with m^{ndim}' if periodic else f'2^k - 1, or m 2^k - 1 with (m - 1)^{ndim}'
465 raise ValueError(
466 f'N = {N} does not coarsen to at most {coarsest} points in {ndim}D: use N = {form} <= {coarsest}'
467 )
468 self.N, self.ndim, self.L, self.xp, self.ops = N, ndim, L, xp, _Ops(xp, fma, periodic)
469 self.bc, self.nu, self.velocity, self.stall = bc, nu, velocity, stall
470 self.relative_residual = None
471 self._coarse = {}
473 def _scales(self, n, s):
474 r"""
475 For the grid of ``n`` points per direction: :math:`s \nu`, :math:`h^{-2}` and the signed
476 :math:`s c_a h^{-1}` (None without advection), which is how :class:`_Ops` takes the operator.
477 """
478 hinv = (n if self.ops.periodic else n + 1) / self.L
479 q = None if self.velocity is None else tuple(float(s * c * hinv) for c in self.velocity)
480 return s * self.nu, float(hinv**2), q
482 def solve(self, b, s, rtol, maxiter):
483 r"""
484 Solve :math:`(I - s A_h)\, x = b` from a zero guess, in the format of ``b``.
486 Stops when the 2-norm of the residual is below ``rtol`` times that of ``b``, when a V-cycle from the second
487 on reduces it by less than ``stall`` (the accuracy of the format), or after ``maxiter`` V-cycles; a cycle from
488 the second on that does not reduce it at all is discarded, and a solve that ends worse than zero returns zero.
489 A solve that stops on the stall rule may not have reached ``rtol``: where V-cycles contract slowly, as beyond
490 the Peclet limit (see :mod:`pySDC.helpers.multigrid`), the stall rule mistakes that for the end of the format.
491 The relative residual of the returned iterate is in :attr:`relative_residual` after every solve.
493 Parameters
494 ----------
495 b : array
496 Right-hand side, of shape ``(N,) * ndim`` and type float64, float32 or float16.
497 s : float
498 Factor of the operator, positive.
499 rtol : float
500 Relative tolerance for the residual.
501 maxiter : int
502 Largest number of V-cycles.
504 Returns
505 -------
506 x : array
507 The solution, of the shape and type of ``b``.
508 cycles : int
509 The V-cycles it took.
510 """
511 if b.shape != (self.N,) * self.ndim:
512 raise ValueError(f'expected a right-hand side of shape {(self.N,) * self.ndim}, got {b.shape}')
513 norm_b = self._norm(b)
514 if norm_b == 0:
515 self.relative_residual = 0.0
516 return self.xp.zeros_like(b), 0
517 x, cycles, previous = None, 0, norm_b # x = None is zero, which the first smoothing step knows without reading
518 while cycles < maxiter:
519 candidate = self._vcycle(0, x, b, s)
520 cycles += 1
521 norm_r = self._norm(self.ops.residual(candidate, b, *self._scales(self.N, s)))
522 # the first cycle is judged by the second: with advection, it can leave its error in the outflow boundary
523 # layer, where the residual is large, and still be a good start
524 first = cycles == 1
525 # past the accuracy of the format: keep what was there; written to stop on a NaN of an overflow too
526 if not (norm_r < previous or first and np.isfinite(norm_r)):
527 break
528 x, stalled, previous = candidate, norm_r > self.stall * previous and not first, norm_r
529 # converged, or at the accuracy the format allows: a V-cycle otherwise reduces the residual by 0.1 to 0.3
530 if norm_r <= rtol * norm_b or stalled:
531 break
532 if not previous < norm_b: # worse than zero, which only a first cycle can leave: return zero
533 x, previous = None, norm_b
534 self.relative_residual = previous / norm_b
535 return (self.xp.zeros_like(b) if x is None else x), cycles
537 def laplacian(self, x, scale=1.0, out=None):
538 r"""
539 :math:`\text{scale}\cdot\Delta_h x` on the finest grid, without :math:`\nu` and the advection, in the type
540 of ``x``, into ``out`` if given; on a GPU on periodic 3D grids by the fused kernel.
541 """
542 return self.ops.stencil(0, x, None, scale * self._scales(self.N, 1.0)[1], out=out)
544 def apply(self, x, scale=1.0, out=None):
545 r""":math:`\text{scale}\cdot A_h x` on the finest grid, the operator the solve inverts, in the type of ``x``."""
546 sn, h2inv, q = self._scales(self.N, scale)
547 return self.ops.stencil(0, x, None, sn * h2inv, out=out, q=q)
549 def _norm(self, v):
550 """The 2-norm; on a GPU by cuBLAS, in one pass, where cupy.linalg.norm reaches a third of the bandwidth."""
551 flat = v.reshape(-1).astype(compute_type(v.dtype), copy=False)
552 if self.ops.gpu:
553 from cupy import cublas
555 return float(cublas.nrm2(flat))
556 return float(np.linalg.norm(flat))
558 def _vcycle(self, level, x, b, s):
559 if level == len(self.sizes) - 1:
560 return self._coarse_solve(b, s)
561 scales = self._scales(self.sizes[level], s)
562 x = self.ops.jacobi_from_zero(b, *scales) if x is None else self.ops.jacobi(x, b, *scales)
563 x = self.ops.jacobi(x, b, *scales)
564 coarse = self.ops.restrict(self.ops.residual(x, b, *scales))
565 x = self.ops.prolong_add(x, self._vcycle(level + 1, None, coarse, s))
566 for _ in range(2):
567 x = self.ops.jacobi(x, b, *scales)
568 return x
570 def _coarse_solve(self, b, s):
571 """Exact solve on the coarsest grid with the dense inverse of its operator, kept per s and type."""
572 n, ct = self.sizes[-1], compute_type(b.dtype)
573 key = (s, ct.str)
574 if key not in self._coarse:
575 if len(self._coarse) >= 12: # an adaptive step size keeps changing s
576 self._coarse.pop(next(iter(self._coarse)))
577 shape, (sn, h2inv, q) = (n,) * self.ndim, self._scales(n, s)
578 ops = _Ops(np, periodic=self.ops.periodic)
579 columns = [
580 e - ops.stencil(0, e.reshape(shape), None, sn * h2inv, q=q).reshape(-1) for e in np.eye(n**self.ndim)
581 ]
582 self._coarse[key] = self.xp.asarray(np.linalg.inv(np.array(columns).T)).astype(ct)
583 inverse = self._coarse[key]
584 return (inverse @ b.reshape(-1).astype(ct, copy=False)).astype(b.dtype).reshape(b.shape)