Coverage for pySDC/implementations/problem_classes/generic_ND_FD.py: 98%
258 statements
« prev ^ index » next coverage.py v7.17.3, created at 2026-11-13 02:59 +0000
1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3"""
4Created on Sat Feb 11 22:39:30 2023
5"""
7import numpy as np
8import scipy.fft
9import scipy.sparse as sp
10import scipy.sparse.linalg as spla
12from pySDC.core.errors import ProblemError
13from pySDC.core.problem import Problem, WorkCounter
14from pySDC.helpers import problem_helper
15from pySDC.implementations.datatype_classes.mesh import mesh
17_GPU_KERNELS = {}
20def _gpu_kernels():
21 """
22 The fused kernels of the FFT solve on a GPU, compiled once.
24 Each repeats the operations of the unfused code it replaces, in the same order and with the same roundings:
25 ``float16`` arithmetic in CuPy rounds every single operation to half precision.
26 """
27 if not _GPU_KERNELS:
28 import cupy as cp
30 # (b / scale) as the solve format, imaginary part zero
31 _GPU_KERNELS['scale_in'] = cp.ElementwiseKernel(
32 'float64 b, float64 s', 'T y', 'y = T(b / s)', 'pysdc_fft_scale_in'
33 )
34 # the real part, widened to float64, times scale
35 _GPU_KERNELS['scale_out'] = cp.ElementwiseKernel(
36 'T x, float64 s', 'float64 o', 'o = (double)x.real() * s', 'pysdc_fft_scale_out'
37 )
38 # (b / scale) * unit, rounded to half precision, into interleaved (re, im) pairs
39 _GPU_KERNELS['half_in'] = cp.ElementwiseKernel(
40 'raw float64 b, float64 s, float64 u',
41 'raw float16 d',
42 'd[2 * i] = float16((b[i] / s) * u); d[2 * i + 1] = float16(0.0)',
43 'pysdc_fft_half_in',
44 )
45 # the complex product of the spectrum with the factor, each operation rounded to half precision
46 _GPU_KERNELS['half_product'] = cp.ElementwiseKernel(
47 'raw float16 m',
48 'raw float16 s',
49 """
50 float16 re = s[2 * i], im = s[2 * i + 1], ir = m[2 * i], ii = m[2 * i + 1];
51 float16 rr = re * ir, mm = im * ii, ri = re * ii, mr = im * ir;
52 s[2 * i] = rr - mm;
53 s[2 * i + 1] = ri + mr;
54 """,
55 'pysdc_fft_half_product',
56 )
57 # the real part, widened to single precision, times unit (single), widened to double, times scale
58 _GPU_KERNELS['half_out'] = cp.ElementwiseKernel(
59 'raw float16 d, float32 u, float64 s',
60 'raw float64 o',
61 'o[i] = (double)((float)d[2 * i] * u) * s',
62 'pysdc_fft_half_out',
63 )
64 return _GPU_KERNELS
67class GenericNDimFinDiff(Problem):
68 r"""
69 Base class for finite difference spatial discretisation in :math:`N` dimensions
71 .. math::
72 \frac{d u}{dt} = A u,
74 where :math:`A \in \mathbb{R}^{nN \times nN}` is a matrix arising from finite difference discretisation of spatial
75 derivatives with :math:`n` degrees of freedom per dimension and :math:`N` dimensions. This generic class follows the MOL
76 (method-of-lines) approach and can be used to discretize partial differential equations such as the advection
77 equation and the heat equation.
79 Parameters
80 ----------
81 nvars : int, optional
82 Spatial resolution for the ND problem. For :math:`N = 2`,
83 set ``nvars=(16, 16)``.
84 coeff : float, optional
85 Factor for finite difference matrix :math:`A`.
86 derivative : int, optional
87 Order of the spatial derivative.
88 freq : tuple of int, optional
89 Spatial frequency, can be a tuple.
90 stencil_type : str, optional
91 Stencil type for finite differences.
92 order : int, optional
93 Order of accuracy of the finite difference discretization.
94 lintol : float, optional
95 Tolerance for spatial solver.
96 liniter : int, optional
97 Maximum number of iterations for linear solver.
98 dtype : dtype-like, optional
99 Precision the state is stored at. ``float64`` by default, which is what every caller got
100 before this existed. The operators follow at ``promote_types(dtype, float32)``, since SciPy
101 has no half-precision sparse matrix and hardware that stores half precision computes in
102 single anyway -- so ``float16`` here means genuinely half-precision *storage* with
103 single-precision arithmetic, which is the arrangement it has on real hardware too.
104 solver_type : str, optional
105 Type of solver. Can be ``'direct'``, ``'GMRES'``, ``'CG'`` or, on periodic grids, ``'FFT'``:
106 the operator is then circulant, so the FFT diagonalises it and a solve is two transforms
107 and a division -- exact, and the cheapest solve there is on a GPU. ``'direct'`` keeps the
108 factorizations of the last ``max_cached_factorizations`` factors and reuses them. ``'MG'`` is
109 geometric multigrid (:class:`~pySDC.helpers.multigrid.GeometricMultigrid`), for the
110 second-order Laplacian (``derivative=2``, ``order=2``, ``stencil_type='center'``) on periodic
111 grids or with ``bc='dirichlet-zero'``, and real states: V-cycles from a zero guess until the relative
112 residual is below ``lintol``, the accuracy of the solve's format is reached, or after
113 ``liniter`` cycles, whose number the work counter ``'MG'`` adds up. It needs no matrix, and on a
114 GPU on periodic 3D grids it runs on fused kernels.
115 solve_precision : dtype-like or None, optional
116 Precision the implicit solve runs at, independently of the state's ``dtype``. ``None``
117 (default) solves at the operator's precision, exactly as without it. Otherwise the solve
118 genuinely works in that precision: the sparse solvers on a copy of the operator held at it,
119 ``'FFT'`` with transforms at it, and the result is widened back to ``dtype``. Sparse
120 matrices stop at single precision in SciPy and CuPy alike, so ``float16`` needs ``'FFT'`` or
121 ``'MG'``; multigrid stores in half precision and computes in single, on a CPU too.
122 On a GPU those are cuFFT's half-precision transforms (power-of-two grids), with every value
123 stored and every product rounded in half precision. A CPU has no half-precision FFT, so
124 there ``float16`` is *emulated*, by rounding through it around a single-precision
125 transform -- which is optimistic, since the transform's own arithmetic is then single.
126 For a complex state the solve is complex at the same precision.
127 bc : str or tuple of 2 string, optional
128 Type of boundary conditions. Default is ``'periodic'``.
129 To define two different types of boundary condition for each side,
130 you can use a tuple, for instance ``bc=("dirichlet", "neumann")``
131 uses Dirichlet BC on the left side, and Neumann BC on the right side.
132 bcParams : dict, optional
133 Parameters for boundary conditions, that can contains those keys :
135 - **val** : value for the boundary value (Dirichlet) or derivative
136 (Neumann), default to 0
137 - **reduce** : if true, reduce the order of the A matrix close to the
138 boundary. If false (default), use shifted stencils close to the
139 boundary.
140 - **neumann_bc_order** : finite difference order that should be used
141 for the neumann BC derivative. If None (default), uses the same
142 order as the discretization for A.
144 Default is None, which takes the default values for each parameters.
145 You can also define a tuple to set different parameters for each
146 side.
147 useGPU : bool, optional
148 Run on the GPU with CuPy instead of on the CPU with NumPy.
150 Attributes
151 ----------
152 A : sparse matrix (CSC)
153 FD discretization matrix of the ND operator.
154 Id : sparse matrix (CSC)
155 Identity matrix of the same dimension as A.
156 xvalues : np.1darray
157 Values of spatial grid.
158 """
160 dtype_u = mesh
161 dtype_f = mesh
162 xp = np
163 xsp = sp
164 linalg = spla
165 # direct solves keep this many factorizations of I - factor * A, see `_direct_solve`
166 max_cached_factorizations = 12
168 def setup_GPU(self):
169 """
170 Switch to GPU modules
171 """
172 import cupy as cp
173 import cupyx.scipy.sparse as csp
174 import cupyx.scipy.sparse.linalg as cspla
176 from pySDC.implementations.datatype_classes.cupy_mesh import cupy_mesh
178 self.xp = cp
179 self.xsp = csp
180 self.linalg = cspla
181 self.dtype_u = cupy_mesh
182 self.dtype_f = cupy_mesh
184 def __init__(
185 self,
186 nvars=512,
187 coeff=1.0,
188 derivative=1,
189 freq=2,
190 stencil_type='center',
191 order=2,
192 lintol=1e-12,
193 liniter=10000,
194 solver_type='direct',
195 bc='periodic',
196 bcParams=None,
197 dtype='float64',
198 useGPU=False,
199 solve_precision=None,
200 ):
201 """Initialization routine"""
202 if useGPU:
203 self.setup_GPU()
205 # make sure parameters have the correct types
206 if type(nvars) not in [int, tuple]:
207 raise ProblemError('nvars should be either tuple or int')
208 if type(freq) not in [int, tuple]:
209 raise ProblemError('freq should be either tuple or int')
211 # transforms nvars into a tuple
212 if type(nvars) is int:
213 nvars = (nvars,)
215 # automatically determine ndim from nvars
216 ndim = len(nvars)
217 if ndim > 3:
218 raise ProblemError(f'can work with up to three dimensions, got {ndim}')
220 # eventually extend freq to other dimension
221 if type(freq) is int:
222 freq = (freq,) * ndim
223 if len(freq) != ndim:
224 raise ProblemError(f'len(freq)={len(freq)}, different to ndim={ndim}')
226 # check values for freq and nvars
227 for f in freq:
228 if ndim == 1 and f == -1:
229 # use Gaussian initial solution in 1D
230 bc = 'periodic'
231 break
232 if f % 2 != 0 and bc == 'periodic':
233 raise ProblemError('need even number of frequencies due to periodic BCs')
234 for nvar in nvars:
235 if nvar % 2 != 0 and bc == 'periodic':
236 raise ProblemError('the setup requires nvars = 2^p per dimension')
237 if (nvar + 1) % 2 != 0 and bc == 'dirichlet-zero':
238 raise ProblemError('setup requires nvars = 2^p - 1')
239 if ndim > 1 and nvars[1:] != nvars[:-1]:
240 raise ProblemError('need a square domain, got %s' % nvars)
242 # invoke super init, passing number of dofs and the precision to store them at
243 dtype = np.dtype(dtype)
245 # SciPy holds no float16 sparse matrix, and hardware that stores half precision computes in
246 # single anyway, so the operators sit at the smallest single-or-wider type that holds `dtype`
247 operator_dtype = np.promote_types(dtype, np.float32)
249 super().__init__(init=(nvars[0] if ndim == 1 else nvars, None, dtype))
251 dx, xvalues = problem_helper.get_1d_grid(size=nvars[0], bc=bc, left_boundary=0.0, right_boundary=1.0)
253 self.A, _ = problem_helper.get_finite_difference_matrix(
254 derivative=derivative,
255 order=order,
256 stencil_type=stencil_type,
257 dx=dx,
258 size=nvars[0],
259 dim=ndim,
260 bc=bc,
261 cupy=useGPU,
262 )
263 self.A *= coeff
265 self.A = self.A.astype(operator_dtype)
267 # SciPy's sparse direct solver wants CSC and CuPy's wants CSR. Whichever one is handed the
268 # wrong layout converts the whole matrix on every call -- that is what cupyx's
269 # `SparseEfficiencyWarning: CSR format is required` is reporting, once per solve. CSR is
270 # also the better layout for the matrix-vector product in `eval_f`.
271 self.A = self.A.tocsr() if useGPU else self.A.tocsc()
273 # the grid feeds every `u_exact`, so it has to live where the solution does
274 self.xvalues = self.xp.asarray(xvalues)
275 self.Id = self.xsp.eye(np.prod(nvars), format='csr' if useGPU else 'csc', dtype=operator_dtype)
277 # store attribute and register them as parameters
278 self._makeAttributeAndRegister('nvars', 'stencil_type', 'order', 'bc', localVars=locals(), readOnly=True)
279 self._makeAttributeAndRegister('freq', 'lintol', 'liniter', 'solver_type', localVars=locals())
280 self.dtype = dtype
281 self.operator_dtype = operator_dtype
283 if self.solver_type != 'direct':
284 self.work_counters[self.solver_type] = WorkCounter()
286 self._factorizations = {}
288 self._coeff, self._derivative = coeff, derivative # what the multigrid solve needs to know of A
289 solve_precision = None if solve_precision is None else np.dtype(solve_precision)
290 self._makeAttributeAndRegister('solve_precision', localVars=locals(), readOnly=True)
291 self._setup_solve(solver_type, solve_precision, bc)
293 def _setup_solve(self, solver_type, solve_precision, bc):
294 """Prepare the reduced-precision operator, or the eigenvalues the FFT solve divides by."""
295 # half precision is the one format that is not a complex dtype for a complex state
296 self._half = solve_precision == np.float16
297 precision = self.operator_dtype if solve_precision is None or self._half else solve_precision
298 if self.dtype.kind == 'c':
299 precision = np.promote_types(precision, np.complex64)
300 self._solve_dtype = precision
302 if solver_type == 'FFT':
303 if bc != 'periodic':
304 raise ProblemError(f'the FFT solve diagonalises a circulant operator and needs periodic BCs, got {bc}')
305 if self._half and self.xp is not np and any(n & (n - 1) for n in self.nvars):
306 raise ProblemError(f'cuFFT computes in half precision on power-of-two grids only, got {self.nvars}')
307 self._fft = scipy.fft if self.xp is np else self.xp.fft
308 # the first column of a circulant transforms to its eigenvalues
309 first_column = self.xp.zeros(self.A.shape[0], dtype=self.operator_dtype)
310 first_column[0] = 1.0
311 self._eigenvalues = self._fft.fftn((self.A @ first_column).reshape(self.nvars))
312 self._half_plan = None
313 elif solver_type == 'MG':
314 from pySDC.helpers.multigrid import GeometricMultigrid
316 laplacian = (self._derivative, self.order, self.stencil_type) == (2, 2, 'center')
317 if bc not in ('periodic', 'dirichlet-zero') or not laplacian:
318 raise ProblemError(
319 'multigrid solves the second-order Laplacian on periodic or zero-Dirichlet grids: needs '
320 'bc="periodic" or "dirichlet-zero", derivative=2, order=2, stencil_type="center", got '
321 f'{bc}, {self._derivative}, {self.order}, {self.stencil_type}'
322 )
323 if self.dtype.kind == 'c' or np.iscomplexobj(self._coeff):
324 raise ProblemError('multigrid solves real systems only')
325 self._multigrid = GeometricMultigrid(self.nvars[0], self.ndim, xp=self.xp, bc=bc)
326 self._mg_dtype = np.dtype(np.float16) if self._half else precision
327 elif solve_precision is not None:
328 if self._half:
329 raise ProblemError('sparse matrices stop at single precision; float16 needs solver_type="FFT" or "MG"')
330 self._A_solve = self.A.astype(precision)
331 self._Id_solve = self.Id.astype(precision)
333 @property
334 def ndim(self):
335 """Number of dimensions of the spatial problem"""
336 return len(self.nvars)
338 @property
339 def dx(self):
340 """Size of the mesh (in all dimensions)"""
341 return self.xvalues[1] - self.xvalues[0]
343 @property
344 def grids(self):
345 """ND grids associated to the problem"""
346 x = self.xvalues
347 if self.ndim == 1:
348 return x
349 if self.ndim == 2:
350 return x[None, :], x[:, None]
351 if self.ndim == 3:
352 return x[None, :, None], x[:, None, None], x[None, None, :]
354 @classmethod
355 def get_default_sweeper_class(cls):
356 """
357 Default sweeper for these problems, which are treated fully implicitly.
359 Returns
360 -------
361 type
362 The sweeper class ``generic_implicit``.
363 """
364 from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit
366 return generic_implicit
368 def eval_f(self, u, t):
369 """
370 Routine to evaluate the right-hand side of the problem.
372 Parameters
373 ----------
374 u : dtype_u
375 Current values.
376 t : float
377 Current time.
379 Returns
380 -------
381 f : dtype_f
382 Values of the right-hand side of the problem.
383 """
384 f = self.f_init
385 f[:] = self.A.dot(u.flatten()).reshape(self.nvars)
386 return f
388 def eval_f_increment(self, base, delta, t):
389 r"""
390 Evaluate :math:`f(w + \delta) - f(w) = A\delta`, which carries an explicit factor
391 :math:`\delta`.
393 The operator is linear, so the increment is the operator applied to the correction and the
394 base state does not enter. Supplying it means a sweeper never has to form the increment by
395 subtracting two stored right-hand sides, whose cancellation error carries
396 :math:`\varepsilon\|A\|` -- see :class:`pySDC.core.problem.Problem`.
398 Parameters
399 ----------
400 base : dtype_u
401 The base state, unused for a linear operator.
402 delta : dtype_u
403 The correction.
404 t : float
405 Current time, accepted for interface compatibility.
407 Returns
408 -------
409 f : dtype_f
410 The increment.
411 """
412 f = self.dtype_f(self.init)
413 f[:] = self.A.dot(delta.flatten()).reshape(self.nvars)
414 return f
416 def solve_system(self, rhs, factor, u0, t):
417 r"""
418 Simple linear solver for :math:`(I-factor\cdot A)\vec{u}=\vec{rhs}`.
420 Parameters
421 ----------
422 rhs : dtype_f
423 Right-hand side for the linear system.
424 factor : float
425 Abbrev. for the local stepsize (or any other factor required).
426 u0 : dtype_u
427 Initial guess for the iterative solver.
428 t : float
429 Current time (e.g. for time-dependent BCs).
431 Returns
432 -------
433 sol : dtype_u
434 The solution of the linear solver.
435 """
436 if self.solver_type in ['FFT', 'MG'] or self.solve_precision is not None:
437 return self._solve_at_precision(rhs, factor, u0)
439 solver_type, Id, A, nvars, lintol, liniter, sol = (
440 self.solver_type,
441 self.Id,
442 self.A,
443 self.nvars,
444 self.lintol,
445 self.liniter,
446 self.u_init,
447 )
449 if solver_type == 'direct':
450 sol[:] = self._direct_solve(factor, lambda: Id - factor * A, rhs.flatten()).reshape(nvars)
451 elif solver_type == 'GMRES':
452 sol[:] = self.linalg.gmres(
453 Id - factor * A,
454 rhs.flatten(),
455 x0=u0.flatten(),
456 rtol=lintol,
457 maxiter=liniter,
458 atol=0,
459 callback=self.work_counters[solver_type],
460 callback_type='legacy',
461 )[0].reshape(nvars)
462 elif solver_type == 'CG':
463 sol[:] = self.linalg.cg(
464 Id - factor * A,
465 rhs.flatten(),
466 x0=u0.flatten(),
467 rtol=lintol,
468 maxiter=liniter,
469 atol=0,
470 callback=self.work_counters[solver_type],
471 )[0].reshape(nvars)
472 else:
473 raise ValueError(f'solver type "{solver_type}" not known in generic advection-diffusion implementation!')
475 return sol
477 def _direct_solve(self, factor, matrix, b):
478 r"""
479 Solve ``matrix() @ x = b`` directly, factorizing :math:`I - factor\cdot A` once per factor.
481 A sweep repeats the same few factors (one per node and level at a fixed step size), so
482 refactorizing on every call, as ``spsolve`` does, repeats identical work: the cached
483 factorization gives the same solution bit for bit, 14 to 40 times faster per solve in one
484 and two dimensions. Like ``spsolve``, the matrix is promoted to the right-hand side's dtype,
485 and the cache is keyed by that dtype too: a single-precision problem can be handed a
486 double-precision right-hand side, since a NumPy float64 scalar such as a quadrature weight
487 widens the array it multiplies, and a factorization's solve refuses to cast that down. The
488 oldest factorization is evicted beyond ``max_cached_factorizations``, as in
489 ``GenericSpectralLinear``, which bounds memory when an adaptive step size keeps changing the
490 factor. On a GPU the solve is left to CuPy's ``spsolve``.
492 Parameters
493 ----------
494 factor : float
495 The factor in front of :math:`A`, which identifies the matrix.
496 matrix : callable
497 Builds the matrix, only called when it is not factorized yet.
498 b : array
499 Flattened right-hand side.
501 Returns
502 -------
503 array
504 The flattened solution.
505 """
506 if self.xp is not np:
507 return self.linalg.spsolve(matrix(), b)
508 key = (factor, b.dtype)
509 if key not in self._factorizations:
510 if len(self._factorizations) >= self.max_cached_factorizations:
511 self._factorizations.pop(next(iter(self._factorizations)))
512 M = matrix()
513 self._factorizations[key] = self.linalg.factorized(M.astype(np.promote_types(M.dtype, b.dtype)))
514 return self._factorizations[key](b)
516 def _solve_at_precision(self, rhs, factor, u0):
517 r"""
518 Solve :math:`(I-factor\cdot A)\vec{u}=\vec{rhs}` at ``solve_precision``, or by FFT.
520 Returns
521 -------
522 sol : dtype_u
523 The solution, widened to the state's precision.
524 """
525 sol = self.dtype_u(self.init)
526 if self.solver_type == 'FFT':
527 sol[:] = self._solve_fft(rhs.view(self.xp.ndarray).reshape(self.nvars), factor)
528 self.work_counters['FFT']()
529 return sol
530 if self.solver_type == 'MG':
531 sol[:] = self._solve_mg(rhs.view(self.xp.ndarray).reshape(self.nvars), factor)
532 return sol
534 precision = self._solve_dtype
536 def matrix():
537 # the factor at the solve's own precision: a float64 scalar would widen the operator back
538 return self._Id_solve - precision.type(factor) * self._A_solve
540 b = rhs.view(self.xp.ndarray).flatten().astype(precision)
541 if self.solver_type == 'direct':
542 x = self._direct_solve(factor, matrix, b)
543 elif self.solver_type in ['GMRES', 'CG']:
544 x0 = u0.view(self.xp.ndarray).flatten().astype(precision)
545 options = {'callback_type': 'legacy'} if self.solver_type == 'GMRES' else {}
546 # A Krylov solve cannot meet a tolerance below its own precision. Asked to, it iterates to
547 # `liniter` and, in CG, divides by a dot product that has rounded to zero: NaN. A hundred
548 # times the precision's epsilon is what single precision still reaches.
549 rtol = max(self.lintol, 100 * float(np.finfo(precision).eps))
550 x = getattr(self.linalg, self.solver_type.lower())(
551 matrix(),
552 b,
553 x0=x0,
554 rtol=rtol,
555 maxiter=self.liniter,
556 atol=0,
557 callback=self.work_counters[self.solver_type],
558 **options,
559 )[0]
560 else:
561 raise ValueError(
562 f'solver type "{self.solver_type}" not known in generic advection-diffusion implementation!'
563 )
564 sol[:] = x.reshape(self.nvars)
565 return sol
567 def _solve_mg(self, b, factor):
568 """
569 Multigrid for ``(I - factor * A) x = b`` in the solve's format, returned in that format.
571 In half precision, the right-hand side is scaled to a maximum of one first and the result
572 widened and scaled back, as in ``_solve_fft``: a correction of 1e-10 is below ``float16``'s
573 smallest subnormal.
574 """
575 scale = float(abs(b).max()) if self._half else 1.0
576 if scale == 0.0:
577 return self.xp.zeros_like(b)
578 x, cycles = self._multigrid.solve(
579 (b if scale == 1.0 else b / scale).astype(self._mg_dtype, copy=False),
580 factor * self._coeff,
581 self.lintol,
582 self.liniter,
583 )
584 self.work_counters['MG'].niter += cycles
585 return x if scale == 1.0 else x.astype(self.dtype) * scale
587 def _solve_fft(self, b, factor):
588 """
589 ``ifft(fft(b) / (1 - factor * eigenvalues))`` at the solve precision, at the state's precision.
591 The right-hand side is scaled to a maximum of one first and the result scaled back, which is
592 exact for a linear solve and keeps a half-precision solve clear of underflow: a correction of
593 1e-10 is below ``float16``'s smallest subnormal. A real state on a GPU takes
594 :meth:`_solve_fft_gpu`, which computes the same bit for bit with less memory traffic.
595 """
596 if self.xp is not np and self.dtype.kind != 'c':
597 return self._solve_fft_gpu(b, factor)
598 return self._solve_fft_unfused(b, factor)
600 def _solve_fft_unfused(self, b, factor):
601 """:meth:`_solve_fft` with array operations, for any backend and a real or complex state."""
602 scale = float(abs(b).max())
603 if scale == 0.0:
604 return self.xp.zeros_like(b)
605 inverse = 1.0 / (1.0 - factor * self._eigenvalues)
606 if self._half:
607 x = self._solve_fft_half(b / scale, inverse)
608 else:
609 precision = np.promote_types(self._solve_dtype, np.complex64)
610 x = self._fft.ifftn(self._fft.fftn((b / scale).astype(precision)) * inverse.astype(precision))
611 # widen first, scale second: scaling at half precision would underflow on the way out
612 x = x.astype(np.promote_types(self.dtype, np.complex64))
613 x = x if self.dtype.kind == 'c' else x.real
614 return x * scale
616 def _solve_fft_half(self, y, inverse):
617 """
618 ``ifft(fft(y) * inverse)`` in half precision for ``|y| <= 1``, returned in single precision.
620 Both transforms are scaled by ``1 / sqrt(N)``, so the forward one is bounded by ``sqrt(N)``
621 and the inverse one by ``sqrt(N)`` times the solution: 1024 on a 1024 x 1024 grid, where an
622 unscaled transform would reach 1e6 and ``float16`` stops at 65504.
623 """
624 unit = 1.0 / np.sqrt(y.size)
625 if self.xp is np:
626 # an emulation: rounded through float16 at every stage, computed in single precision
628 def rounded(z):
629 return z.real.astype(np.float16).astype(np.float32) + 1j * z.imag.astype(np.float16).astype(np.float32)
631 spectrum = rounded(self._fft.fftn(rounded((y * unit).astype(np.complex64))))
632 spectrum = rounded(spectrum * rounded(inverse.astype(np.complex64)))
633 return rounded(self._fft.ifftn(spectrum, norm='forward')) * unit
635 import cupy as cp
636 from cupy.cuda import cufft
638 shape, n = tuple(self.nvars), y.size
639 if self._half_plan is None:
640 # complex32 in and out, stored as float16 (re, im) pairs along the last axis
641 self._half_plan = cufft.XtPlanNd(
642 shape, shape, 1, n, 'E', shape, 1, n, 'E', 1, 'E', order='C', last_axis=-1, last_size=None
643 )
644 data = cp.zeros(shape[:-1] + (2 * shape[-1],), dtype=cp.float16)
645 data[..., 0::2] = (y * unit).real
646 if y.dtype.kind == 'c':
647 data[..., 1::2] = (y * unit).imag
648 spectrum = cp.empty_like(data)
649 self._half_plan.fft(data, spectrum, cufft.CUFFT_FORWARD)
650 re, im = spectrum[..., 0::2].copy(), spectrum[..., 1::2].copy()
651 ir, ii = inverse.real.astype(cp.float16), inverse.imag.astype(cp.float16)
652 spectrum[..., 0::2] = re * ir - im * ii
653 spectrum[..., 1::2] = re * ii + im * ir
654 self._half_plan.fft(spectrum, data, cufft.CUFFT_INVERSE)
655 return (data[..., 0::2].astype(cp.float32) + 1j * data[..., 1::2].astype(cp.float32)) * unit
657 def _solve_fft_gpu(self, b, factor):
658 """
659 The FFT solve of :meth:`_solve_fft` for a real state on a GPU, with fused kernels.
661 It computes exactly what :meth:`_solve_fft` and :meth:`_solve_fft_half` compute, operation by operation and
662 in the same order, so the results agree bit for bit. Only the memory traffic is less: the factor
663 ``1 / (1 - factor * eigenvalues)`` in the solve format is kept for every factor instead of being formed anew
664 in every call, and the scaling, the conversions and, in half precision, the complex product are one kernel
665 each instead of several.
667 The cache holds one array of the size of the state per factor, complex in the solve format, for the last
668 ``max_cached_factorizations`` factors. SDC with a fixed step size uses one factor per collocation node.
669 """
670 import cupy as cp
672 kernels = _gpu_kernels()
673 scale = float(abs(b).max())
674 if scale == 0.0:
675 return cp.zeros_like(b)
676 key = (factor, self._half)
677 cache = self.__dict__.setdefault('_fft_factors', {})
678 if key not in cache:
679 if len(cache) >= self.max_cached_factorizations:
680 cache.pop(next(iter(cache)))
681 inverse = 1.0 / (1.0 - factor * self._eigenvalues)
682 if self._half:
683 factors = cp.empty(inverse.shape[:-1] + (2 * inverse.shape[-1],), dtype=cp.float16)
684 factors[..., 0::2] = inverse.real.astype(cp.float16)
685 factors[..., 1::2] = inverse.imag.astype(cp.float16)
686 else:
687 factors = inverse.astype(np.promote_types(self._solve_dtype, np.complex64))
688 cache[key] = factors
689 factors = cache[key]
690 if not self._half:
691 y = kernels['scale_in'](
692 b, scale, cp.empty(b.shape, dtype=np.promote_types(self._solve_dtype, np.complex64))
693 )
694 spectrum = self._fft.fftn(y)
695 spectrum *= factors
696 return kernels['scale_out'](self._fft.ifftn(spectrum), scale)
698 from cupy.cuda import cufft
700 shape, n = tuple(self.nvars), b.size
701 unit = 1.0 / np.sqrt(n)
702 if self._half_plan is None:
703 self._half_plan = cufft.XtPlanNd(
704 shape, shape, 1, n, 'E', shape, 1, n, 'E', 1, 'E', order='C', last_axis=-1, last_size=None
705 )
706 data = cp.empty(shape[:-1] + (2 * shape[-1],), dtype=cp.float16)
707 kernels['half_in'](b, scale, unit, data, size=n)
708 spectrum = cp.empty_like(data)
709 self._half_plan.fft(data, spectrum, cufft.CUFFT_FORWARD)
710 kernels['half_product'](factors, spectrum, size=n)
711 self._half_plan.fft(spectrum, data, cufft.CUFFT_INVERSE)
712 x = cp.empty(shape, dtype=self.dtype)
713 kernels['half_out'](data, np.float32(unit), scale, x, size=n)
714 return x