Coverage for pySDC/helpers/spectral_helper.py: 93%
797 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-25 20:28 +0000
1import numpy as np
2import scipy
3from pySDC.implementations.datatype_classes.mesh import mesh
4from scipy.special import factorial
5from functools import partial, wraps
6import logging
9def cache(func):
10 """
11 Decorator for caching return values of functions.
12 This is very similar to `functools.cache`, but without the memory leaks (see
13 https://docs.astral.sh/ruff/rules/cached-instance-method/).
15 Example:
17 .. code-block:: python
19 num_calls = 0
21 @cache
22 def increment(x):
23 num_calls += 1
24 return x + 1
26 increment(0) # returns 1, num_calls = 1
27 increment(1) # returns 2, num_calls = 2
28 increment(0) # returns 1, num_calls = 2
31 Args:
32 func (function): The function you want to cache the return value of
34 Returns:
35 return value of func
36 """
37 attr_cache = f"_{func.__name__}_cache"
39 @wraps(func)
40 def wrapper(self, *args, **kwargs):
41 if not hasattr(self, attr_cache):
42 setattr(self, attr_cache, {})
44 cache = getattr(self, attr_cache)
46 key = (args, frozenset(kwargs.items()))
47 if key in cache:
48 return cache[key]
49 result = func(self, *args, **kwargs)
50 cache[key] = result
51 return result
53 return wrapper
56class SpectralHelper1D:
57 """
58 Abstract base class for 1D spectral discretizations. Defines a common interface with parameters and functions that
59 all bases need to have.
61 When implementing new bases, please take care to use the modules that are supplied as class attributes to enable
62 the code for GPUs.
64 Attributes:
65 N (int): Resolution
66 x0 (float): Coordinate of left boundary
67 x1 (float): Coordinate of right boundary
68 L (float): Length of the domain
69 useGPU (bool): Whether to use GPUs
71 """
73 fft_lib = scipy.fft
74 sparse_lib = scipy.sparse
75 linalg = scipy.sparse.linalg
76 xp = np
77 distributable = False
79 def __init__(self, N, x0=None, x1=None, useGPU=False, useFFTW=False):
80 """
81 Constructor
83 Args:
84 N (int): Resolution
85 x0 (float): Coordinate of left boundary
86 x1 (float): Coordinate of right boundary
87 useGPU (bool): Whether to use GPUs
88 useFFTW (bool): Whether to use FFTW for the transforms
89 """
90 self.N = N
91 self.x0 = x0
92 self.x1 = x1
93 self.L = x1 - x0
94 self.useGPU = useGPU
95 self.plans = {}
96 self.logger = logging.getLogger(name=type(self).__name__)
98 if useGPU:
99 self.setup_GPU()
100 self.logger.debug('Set up for GPU')
101 else:
102 self.setup_CPU(useFFTW=useFFTW)
104 if useGPU and useFFTW:
105 raise ValueError('Please run either on GPUs or with FFTW, not both!')
107 def setup_GPU(self):
108 """switch to GPU modules"""
109 import cupy as cp
110 import cupyx.scipy.sparse as sparse_lib
111 import cupyx.scipy.sparse.linalg as linalg
112 import cupyx.scipy.fft as fft_lib
113 from pySDC.implementations.datatype_classes.cupy_mesh import cupy_mesh
115 self.xp = cp
116 self.sparse_lib = sparse_lib
117 self.linalg = linalg
118 self.fft_lib = fft_lib
120 @classmethod
121 def setup_CPU(cls, useFFTW=False):
122 """switch to CPU modules"""
124 cls.xp = np
125 cls.sparse_lib = scipy.sparse
126 cls.linalg = scipy.sparse.linalg
128 if useFFTW:
129 from mpi4py_fft import fftw
131 cls.fft_backend = 'fftw'
132 cls.fft_lib = fftw
133 else:
134 cls.fft_backend = 'scipy'
135 cls.fft_lib = scipy.fft
137 cls.fft_comm_backend = 'MPI'
138 cls.dtype = mesh
140 def get_Id(self):
141 """
142 Get identity matrix
144 Returns:
145 sparse diagonal identity matrix
146 """
147 return self.sparse_lib.eye(self.N)
149 def get_zero(self):
150 """
151 Get a matrix with all zeros of the correct size.
153 Returns:
154 sparse matrix with zeros everywhere
155 """
156 return 0 * self.get_Id()
158 def get_differentiation_matrix(self):
159 raise NotImplementedError()
161 def get_integration_matrix(self):
162 raise NotImplementedError()
164 def get_integration_weights(self):
165 """Weights for integration across entire domain"""
166 raise NotImplementedError()
168 def get_wavenumbers(self):
169 """
170 Get the grid in spectral space
171 """
172 raise NotImplementedError
174 def get_empty_operator_matrix(self, S, O):
175 """
176 Return a matrix of operators to be filled with the connections between the solution components.
178 Args:
179 S (int): Number of components in the solution
180 O (sparse matrix): Zero matrix used for initialization
182 Returns:
183 list of lists containing sparse zeros
184 """
185 return [[O for _ in range(S)] for _ in range(S)]
187 def get_basis_change_matrix(self, *args, **kwargs):
188 """
189 Some spectral discretization change the basis during differentiation. This method can be used to transfer
190 between the various bases.
192 This method accepts arbitrary arguments that may not be used in order to provide an easy interface for multi-
193 dimensional bases. For instance, you may combine an FFT discretization with an ultraspherical discretization.
194 The FFT discretization will always be in the same base, but the ultraspherical discretization uses a different
195 base for every derivative. You can then ask all bases for transfer matrices from one ultraspherical derivative
196 base to the next. The FFT discretization will ignore this and return an identity while the ultraspherical
197 discretization will return the desired matrix. After a Kronecker product, you get the 2D version of the matrix
198 you want. This is what the `SpectralHelper` does when you call the method of the same name on it.
200 Returns:
201 sparse bases change matrix
202 """
203 return self.sparse_lib.eye(self.N)
205 def get_BC(self, kind):
206 """
207 To facilitate boundary conditions (BCs) we use either a basis where all functions satisfy the BCs automatically,
208 as is the case in FFT basis for periodic BCs, or boundary bordering. In boundary bordering, specific lines in
209 the matrix are replaced by the boundary conditions as obtained by this method.
211 Args:
212 kind (str): The type of BC you want to implement please refer to the implementations of this method in the
213 individual 1D bases for what is implemented
215 Returns:
216 self.xp.array: Boundary condition
217 """
218 raise NotImplementedError(f'No boundary conditions of {kind=!r} implemented!')
220 def get_filter_matrix(self, kmin=0, kmax=None):
221 """
222 Get a bandpass filter.
224 Args:
225 kmin (int): Lower limit of the bandpass filter
226 kmax (int): Upper limit of the bandpass filter
228 Returns:
229 sparse matrix
230 """
232 k = abs(self.get_wavenumbers())
234 kmax = max(k) if kmax is None else kmax
236 mask = self.xp.logical_or(k >= kmax, k < kmin)
238 if self.useGPU:
239 Id = self.get_Id().get()
240 else:
241 Id = self.get_Id()
242 F = Id.tolil()
243 F[:, mask] = 0
244 return F.tocsc()
246 def get_1dgrid(self):
247 """
248 Get the grid in physical space
250 Returns:
251 self.xp.array: Grid
252 """
253 raise NotImplementedError
256class ChebychevHelper(SpectralHelper1D):
257 """
258 The Chebychev base consists of special kinds of polynomials, with the main advantage that you can easily transform
259 between physical and spectral space by discrete cosine transform.
260 The differentiation in the Chebychev T base is dense, but can be preconditioned to yield a differentiation operator
261 that moves to Chebychev U basis during differentiation, which is sparse. When using this technique, problems need to
262 be formulated in first order formulation.
264 This implementation is largely based on the Dedalus paper (https://doi.org/10.1103/PhysRevResearch.2.023068).
265 """
267 def __init__(self, *args, x0=-1, x1=1, **kwargs):
268 """
269 Constructor.
270 Please refer to the parent class for additional arguments. Notably, you have to supply a resolution `N` and you
271 may choose to run on GPUs via the `useGPU` argument.
273 Args:
274 x0 (float): Coordinate of left boundary. Note that only -1 is currently implented
275 x1 (float): Coordinate of right boundary. Note that only +1 is currently implented
276 """
277 # need linear transformation y = ax + b with a = (x1-x0)/2 and b = (x1+x0)/2
278 self.lin_trf_fac = (x1 - x0) / 2
279 self.lin_trf_off = (x1 + x0) / 2
280 super().__init__(*args, x0=x0, x1=x1, **kwargs)
282 self.norm = self.get_norm()
284 def get_1dgrid(self):
285 '''
286 Generates a 1D grid with Chebychev points. These are clustered at the boundary. You need this kind of grid to
287 use discrete cosine transformation (DCT) to get the Chebychev representation. If you want a different grid, you
288 need to do an affine transformation before any Chebychev business.
290 Returns:
291 numpy.ndarray: 1D grid
292 '''
293 return self.lin_trf_fac * self.xp.cos(np.pi / self.N * (self.xp.arange(self.N) + 0.5)) + self.lin_trf_off
295 def get_wavenumbers(self):
296 """Get the domain in spectral space"""
297 return self.xp.arange(self.N)
299 @cache
300 def get_conv(self, name, N=None):
301 '''
302 Get conversion matrix between different kinds of polynomials. The supported kinds are
303 - T: Chebychev polynomials of first kind
304 - U: Chebychev polynomials of second kind
305 - D: Dirichlet recombination.
307 You get the desired matrix by choosing a name as ``A2B``. I.e. ``T2U`` for the conversion matrix from T to U.
308 Once generates matrices are cached. So feel free to call the method as often as you like.
310 Args:
311 name (str): Conversion code, e.g. 'T2U'
312 N (int): Size of the matrix (optional)
314 Returns:
315 scipy.sparse: Sparse conversion matrix
316 '''
317 N = N if N else self.N
318 sp = self.sparse_lib
320 def get_forward_conv(name):
321 if name == 'T2U':
322 mat = (sp.eye(N) - sp.eye(N, k=2)).tocsc() / 2.0
323 mat[:, 0] *= 2
324 elif name == 'D2T':
325 mat = sp.eye(N) - sp.eye(N, k=2)
326 elif name[0] == name[-1]:
327 mat = self.sparse_lib.eye(self.N)
328 else:
329 raise NotImplementedError(f'Don\'t have conversion matrix {name!r}')
330 return mat
332 try:
333 mat = get_forward_conv(name)
334 except NotImplementedError as E:
335 try:
336 fwd = get_forward_conv(name[::-1])
337 import scipy.sparse as sp
339 if self.sparse_lib == sp:
340 mat = self.sparse_lib.linalg.inv(fwd.tocsc())
341 else:
342 mat = self.sparse_lib.csc_matrix(sp.linalg.inv(fwd.tocsc().get()))
343 except NotImplementedError:
344 raise NotImplementedError from E
346 return mat
348 def get_basis_change_matrix(self, conv='T2T', **kwargs):
349 """
350 As the differentiation matrix in Chebychev-T base is dense but is sparse when simultaneously changing base to
351 Chebychev-U, you may need a basis change matrix to transfer the other matrices as well. This function returns a
352 conversion matrix from `ChebychevHelper.get_conv`. Not that `**kwargs` are used to absorb arguments for other
353 bases, see documentation of `SpectralHelper1D.get_basis_change_matrix`.
355 Args:
356 conv (str): Conversion code, i.e. T2U
358 Returns:
359 Sparse conversion matrix
360 """
361 return self.get_conv(conv)
363 def get_integration_matrix(self, lbnd=0):
364 """
365 Get matrix for integration
367 Args:
368 lbnd (float): Lower bound for integration, only 0 is currently implemented
370 Returns:
371 Sparse integration matrix
372 """
373 S = self.sparse_lib.diags(1 / (self.xp.arange(self.N - 1) + 1), offsets=-1) @ self.get_conv('T2U')
374 n = self.xp.arange(self.N)
375 if lbnd == 0:
376 S = S.tocsc()
377 S[0, 1::2] = (
378 (n / (2 * (self.xp.arange(self.N) + 1)))[1::2]
379 * (-1) ** (self.xp.arange(self.N // 2))
380 / (np.append([1], self.xp.arange(self.N // 2 - 1) + 1))
381 ) * self.lin_trf_fac
382 else:
383 raise NotImplementedError(f'This function allows to integrate only from x=0, you attempted from x={lbnd}.')
384 return S
386 def get_integration_weights(self):
387 """Weights for integration across entire domain"""
388 n = self.xp.arange(self.N, dtype=float)
390 weights = (-1) ** n + 1
391 weights[2:] /= 1 - (n**2)[2:]
393 weights /= 2 / self.L
394 return weights
396 def get_differentiation_matrix(self, p=1):
397 '''
398 Keep in mind that the T2T differentiation matrix is dense.
400 Args:
401 p (int): Derivative you want to compute
403 Returns:
404 numpy.ndarray: Differentiation matrix
405 '''
406 D = self.xp.zeros((self.N, self.N))
407 for j in range(self.N):
408 for k in range(j):
409 D[k, j] = 2 * j * ((j - k) % 2)
411 D[0, :] /= 2
412 return self.sparse_lib.csc_matrix(self.xp.linalg.matrix_power(D, p)) / self.lin_trf_fac**p
414 @cache
415 def get_norm(self, N=None):
416 '''
417 Get normalization for converting Chebychev coefficients and DCT
419 Args:
420 N (int, optional): Resolution
422 Returns:
423 self.xp.array: Normalization
424 '''
425 N = self.N if N is None else N
426 norm = self.xp.ones(N) / N
427 norm[0] /= 2
428 return norm
430 def transform(self, u, *args, axes=None, shape=None, **kwargs):
431 """
432 DCT along axes. `kwargs` will be passed on to the FFT library.
434 Args:
435 u: Data you want to transform
436 axes (tuple): Axes you want to transform along
438 Returns:
439 Data in spectral space
440 """
441 axes = axes if axes else tuple(i for i in range(u.ndim))
442 kwargs['s'] = shape
443 kwargs['norm'] = kwargs.get('norm', 'backward')
445 trf = self.fft_lib.dctn(u, *args, axes=axes, type=2, **kwargs)
446 for axis in axes:
448 if self.N < trf.shape[axis]:
449 # mpi4py-fft implements padding only for FFT, where the frequencies are sorted such that the zeros are
450 # removed in the middle rather than the end. We need to resort this here and put the highest frequencies
451 # in the middle.
452 _trf = self.xp.zeros_like(trf)
453 N = self.N
454 N_pad = _trf.shape[axis] - N
455 end_first_half = N // 2 + 1
457 # copy first "half"
458 su = [slice(None)] * trf.ndim
459 su[axis] = slice(0, end_first_half)
460 _trf[tuple(su)] = trf[tuple(su)]
462 # copy second "half"
463 su = [slice(None)] * u.ndim
464 su[axis] = slice(end_first_half + N_pad, None)
465 s_u = [slice(None)] * u.ndim
466 s_u[axis] = slice(end_first_half, N)
467 _trf[tuple(su)] = trf[tuple(s_u)]
469 # # copy values to be cut
470 # su = [slice(None)] * u.ndim
471 # su[axis] = slice(end_first_half, end_first_half + N_pad)
472 # s_u = [slice(None)] * u.ndim
473 # s_u[axis] = slice(-N_pad, None)
474 # _trf[tuple(su)] = trf[tuple(s_u)]
476 trf = _trf
478 expansion = [np.newaxis for _ in u.shape]
479 expansion[axis] = slice(0, u.shape[axis], 1)
480 norm = self.xp.ones(trf.shape[axis]) * self.norm[-1]
481 norm[: self.N] = self.norm
482 trf *= norm[(*expansion,)]
483 return trf
485 def itransform(self, u, *args, axes=None, shape=None, **kwargs):
486 """
487 Inverse DCT along axis.
489 Args:
490 u: Data you want to transform
491 axes (tuple): Axes you want to transform along
493 Returns:
494 Data in physical space
495 """
496 axes = axes if axes else tuple(i for i in range(u.ndim))
497 kwargs['s'] = shape
498 kwargs['norm'] = kwargs.get('norm', 'backward')
499 kwargs['overwrite_x'] = kwargs.get('overwrite_x', False)
501 for axis in axes:
503 if self.N == u.shape[axis]:
504 _u = u.copy()
505 else:
506 # mpi4py-fft implements padding only for FFT, where the frequencies are sorted such that the zeros are
507 # added in the middle rather than the end. We need to resort this here and put the padding in the end.
508 N = self.N
509 _u = self.xp.zeros_like(u)
511 # copy first half
512 su = [slice(None)] * u.ndim
513 su[axis] = slice(0, N // 2 + 1)
514 _u[tuple(su)] = u[tuple(su)]
516 # copy second half
517 su = [slice(None)] * u.ndim
518 su[axis] = slice(-(N // 2), None)
519 s_u = [slice(None)] * u.ndim
520 s_u[axis] = slice(N // 2, N // 2 + (N // 2))
521 _u[tuple(s_u)] = u[tuple(su)]
523 if N % 2 == 0:
524 su = [slice(None)] * u.ndim
525 su[axis] = N // 2
526 _u[tuple(su)] *= 2
528 # generate norm
529 expansion = [np.newaxis for _ in u.shape]
530 expansion[axis] = slice(0, u.shape[axis], 1)
531 norm = self.xp.ones(_u.shape[axis])
532 norm[: self.N] = self.norm
533 norm = self.get_norm(u.shape[axis]) * _u.shape[axis] / self.N
535 _u /= norm[(*expansion,)]
537 return self.fft_lib.idctn(_u, *args, axes=axes, type=2, **kwargs)
539 def get_BC(self, kind, **kwargs):
540 """
541 Get boundary condition row for boundary bordering. `kwargs` will be passed on to implementations of the BC of
542 the kind you choose. Specifically, `x` for `'dirichlet'` boundary condition, which is the coordinate at which to
543 set the BC.
545 Args:
546 kind ('integral' or 'dirichlet'): Kind of boundary condition you want
547 """
548 if kind.lower() == 'integral':
549 return self.get_integ_BC_row(**kwargs)
550 elif kind.lower() == 'dirichlet':
551 return self.get_Dirichlet_BC_row(**kwargs)
552 elif kind.lower() == 'neumann':
553 return self.get_Neumann_BC_row(**kwargs)
554 else:
555 return super().get_BC(kind)
557 def get_integ_BC_row(self):
558 """
559 Get a row for generating integral BCs with T polynomials.
560 It returns the values of the integrals of T polynomials over the entire interval.
562 Returns:
563 self.xp.ndarray: Row to put into a matrix
564 """
565 n = self.xp.arange(self.N) + 1
566 me = self.xp.zeros_like(n).astype(float)
567 me[2:] = ((-1) ** n[1:-1] + 1) / (1 - n[1:-1] ** 2)
568 me[0] = 2.0
569 return me
571 def get_Dirichlet_BC_row(self, x):
572 """
573 Get a row for generating Dirichlet BCs at x with T polynomials.
574 It returns the values of the T polynomials at x.
576 Args:
577 x (float): Position of the boundary condition
579 Returns:
580 self.xp.ndarray: Row to put into a matrix
581 """
582 if x == -1:
583 return (-1) ** self.xp.arange(self.N)
584 elif x == 1:
585 return self.xp.ones(self.N)
586 elif x == 0:
587 n = (1 + (-1) ** self.xp.arange(self.N)) / 2
588 n[2::4] *= -1
589 return n
590 else:
591 raise NotImplementedError(f'Don\'t know how to generate Dirichlet BC\'s at {x=}!')
593 def get_Neumann_BC_row(self, x):
594 """
595 Get a row for generating Neumann BCs at x with T polynomials.
597 Args:
598 x (float): Position of the boundary condition
600 Returns:
601 self.xp.ndarray: Row to put into a matrix
602 """
603 n = self.xp.arange(self.N, dtype='D')
604 nn = n**2
605 if x == -1:
606 me = nn
607 me[1:] *= (-1) ** n[:-1]
608 return me
609 elif x == 1:
610 return nn
611 else:
612 raise NotImplementedError(f'Don\'t know how to generate Neumann BC\'s at {x=}!')
614 def get_Dirichlet_recombination_matrix(self):
615 '''
616 Get matrix for Dirichlet recombination, which changes the basis to have sparse boundary conditions.
617 This makes for a good right preconditioner.
619 Returns:
620 scipy.sparse: Sparse conversion matrix
621 '''
622 N = self.N
623 sp = self.sparse_lib
625 return sp.eye(N) - sp.eye(N, k=2)
628class UltrasphericalHelper(ChebychevHelper):
629 """
630 This implementation follows https://doi.org/10.1137/120865458.
631 The ultraspherical method works in Chebychev polynomials as well, but also uses various Gegenbauer polynomials.
632 The idea is that for every derivative of Chebychev T polynomials, there is a basis of Gegenbauer polynomials where the differentiation matrix is a single off-diagonal.
633 There are also conversion operators from one derivative basis to the next that are sparse.
635 This basis is used like this: For every equation that you have, look for the highest derivative and bump all matrices to the correct basis. If your highest derivative is 2 and you have an identity, it needs to get bumped from 0 to 1 and from 1 to 2. If you have a first derivative as well, it needs to be bumped from 1 to 2.
636 You don't need the same resulting basis in all equations. You just need to take care that you translate the right hand side to the correct basis as well.
637 """
639 def get_differentiation_matrix(self, p=1):
640 """
641 Notice that while sparse, this matrix is not diagonal, which means the inversion cannot be parallelized easily.
643 Args:
644 p (int): Order of the derivative
646 Returns:
647 sparse differentiation matrix
648 """
649 sp = self.sparse_lib
650 xp = self.xp
651 N = self.N
652 l = p
653 return 2 ** (l - 1) * factorial(l - 1) * sp.diags(xp.arange(N - l) + l, offsets=l) / self.lin_trf_fac**p
655 def get_S(self, lmbda):
656 """
657 Get matrix for bumping the derivative base by one from lmbda to lmbda + 1. This is the same language as in
658 https://doi.org/10.1137/120865458.
660 Args:
661 lmbda (int): Ingoing derivative base
663 Returns:
664 sparse matrix: Conversion from derivative base lmbda to lmbda + 1
665 """
666 N = self.N
668 if lmbda == 0:
669 sp = scipy.sparse
670 mat = ((sp.eye(N) - sp.eye(N, k=2)) / 2.0).tolil()
671 mat[:, 0] *= 2
672 else:
673 sp = self.sparse_lib
674 xp = self.xp
675 mat = sp.diags(lmbda / (lmbda + xp.arange(N))) - sp.diags(
676 lmbda / (lmbda + 2 + xp.arange(N - 2)), offsets=+2
677 )
679 return self.sparse_lib.csc_matrix(mat)
681 def get_basis_change_matrix(self, p_in=0, p_out=0, **kwargs):
682 """
683 Get a conversion matrix from derivative base `p_in` to `p_out`.
685 Args:
686 p_out (int): Resulting derivative base
687 p_in (int): Ingoing derivative base
688 """
689 mat_fwd = self.sparse_lib.eye(self.N)
690 for i in range(min([p_in, p_out]), max([p_in, p_out])):
691 mat_fwd = self.get_S(i) @ mat_fwd
693 if p_out > p_in:
694 return mat_fwd
696 else:
697 # We have to invert the matrix on CPU because the GPU equivalent is not implemented in CuPy at the time of writing.
698 import scipy.sparse as sp
700 if self.useGPU:
701 mat_fwd = mat_fwd.get()
703 mat_bck = sp.linalg.inv(mat_fwd.tocsc())
705 return self.sparse_lib.csc_matrix(mat_bck)
707 def get_integration_matrix(self):
708 """
709 Get an integration matrix. Please use `UltrasphericalHelper.get_integration_constant` afterwards to compute the
710 integration constant such that integration starts from x=-1.
712 Example:
714 .. code-block:: python
716 import numpy as np
717 from pySDC.helpers.spectral_helper import UltrasphericalHelper
719 N = 4
720 helper = UltrasphericalHelper(N)
721 coeffs = np.random.random(N)
722 coeffs[-1] = 0
724 poly = np.polynomial.Chebyshev(coeffs)
726 S = helper.get_integration_matrix()
727 U_hat = S @ coeffs
728 U_hat[0] = helper.get_integration_constant(U_hat, axis=-1)
730 assert np.allclose(poly.integ(lbnd=-1).coef[:-1], U_hat)
732 Returns:
733 sparse integration matrix
734 """
735 return (
736 self.sparse_lib.diags(1 / (self.xp.arange(self.N - 1) + 1), offsets=-1)
737 @ self.get_basis_change_matrix(p_out=1, p_in=0)
738 * self.lin_trf_fac
739 )
741 def get_integration_constant(self, u_hat, axis):
742 """
743 Get integration constant for lower bound of -1. See documentation of `UltrasphericalHelper.get_integration_matrix` for details.
745 Args:
746 u_hat: Solution in spectral space
747 axis: Axis you want to integrate over
749 Returns:
750 Integration constant, has one less dimension than `u_hat`
751 """
752 slices = [
753 None,
754 ] * u_hat.ndim
755 slices[axis] = slice(1, u_hat.shape[axis])
756 return self.xp.sum(u_hat[(*slices,)] * (-1) ** (self.xp.arange(u_hat.shape[axis] - 1)), axis=axis)
759class FFTHelper(SpectralHelper1D):
760 distributable = True
762 def __init__(self, *args, x0=0, x1=2 * np.pi, **kwargs):
763 """
764 Constructor.
765 Please refer to the parent class for additional arguments. Notably, you have to supply a resolution `N` and you
766 may choose to run on GPUs via the `useGPU` argument.
768 Args:
769 x0 (float, optional): Coordinate of left boundary
770 x1 (float, optional): Coordinate of right boundary
771 """
772 super().__init__(*args, x0=x0, x1=x1, **kwargs)
774 def get_1dgrid(self):
775 """
776 We use equally spaced points including the left boundary and not including the right one, which is the left boundary.
777 """
778 dx = self.L / self.N
779 return self.xp.arange(self.N) * dx + self.x0
781 def get_wavenumbers(self):
782 """
783 Be careful that this ordering is very unintuitive.
784 """
785 return self.xp.fft.fftfreq(self.N, 1.0 / self.N) * 2 * np.pi / self.L
787 def get_differentiation_matrix(self, p=1):
788 """
789 This matrix is diagonal, allowing to invert concurrently.
791 Args:
792 p (int): Order of the derivative
794 Returns:
795 sparse differentiation matrix
796 """
797 k = self.get_wavenumbers()
799 if self.useGPU:
800 if p > 1:
801 # Have to raise the matrix to power p on CPU because the GPU equivalent is not implemented in CuPy at the time of writing.
802 from scipy.sparse.linalg import matrix_power
804 D = self.sparse_lib.diags(1j * k).get()
805 return self.sparse_lib.csc_matrix(matrix_power(D, p))
806 else:
807 return self.sparse_lib.diags(1j * k)
808 else:
809 return self.linalg.matrix_power(self.sparse_lib.diags(1j * k), p)
811 def get_integration_matrix(self, p=1):
812 """
813 Get integration matrix to compute `p`-th integral over the entire domain.
815 Args:
816 p (int): Order of integral you want to compute
818 Returns:
819 sparse integration matrix
820 """
821 k = self.xp.array(self.get_wavenumbers(), dtype='complex128')
822 k[0] = 1j * self.L
823 return self.linalg.matrix_power(self.sparse_lib.diags(1 / (1j * k)), p)
825 def get_integration_weights(self):
826 """Weights for integration across entire domain"""
827 weights = self.xp.zeros(self.N)
828 weights[0] = self.L / self.N
829 return weights
831 def get_plan(self, u, forward, *args, **kwargs):
832 if self.fft_lib.__name__ == 'mpi4py_fft.fftw':
833 if 'axes' in kwargs.keys():
834 kwargs['axes'] = tuple(kwargs['axes'])
835 key = (forward, u.shape, args, *(me for me in kwargs.values()))
836 if key in self.plans.keys():
837 return self.plans[key]
838 else:
839 self.logger.debug(f'Generating FFT plan for {key=}')
840 transform = self.fft_lib.fftn(u, *args, **kwargs) if forward else self.fft_lib.ifftn(u, *args, **kwargs)
841 self.plans[key] = transform
843 return self.plans[key]
844 else:
845 if forward:
846 return partial(self.fft_lib.fftn, norm=kwargs.get('norm', 'backward'))
847 else:
848 return partial(self.fft_lib.ifftn, norm=kwargs.get('norm', 'forward'))
850 def transform(self, u, *args, axes=None, shape=None, **kwargs):
851 """
852 FFT along axes. `kwargs` are passed on to the FFT library.
854 Args:
855 u: Data you want to transform
856 axes (tuple): Axes you want to transform over
858 Returns:
859 transformed data
860 """
861 axes = axes if axes else tuple(i for i in range(u.ndim))
862 kwargs['s'] = shape
863 plan = self.get_plan(u, *args, forward=True, axes=axes, **kwargs)
864 return plan(u, *args, axes=axes, **kwargs)
866 def itransform(self, u, *args, axes=None, shape=None, **kwargs):
867 """
868 Inverse FFT.
870 Args:
871 u: Data you want to transform
872 axes (tuple): Axes over which to transform
874 Returns:
875 transformed data
876 """
877 axes = axes if axes else tuple(i for i in range(u.ndim))
878 kwargs['s'] = shape
879 plan = self.get_plan(u, *args, forward=False, axes=axes, **kwargs)
880 return plan(u, *args, axes=axes, **kwargs) / np.prod([u.shape[axis] for axis in axes])
882 def get_BC(self, kind):
883 """
884 Get a sort of boundary condition. You can use `kind=integral`, to fix the integral, or you can use `kind=Nyquist`.
885 The latter is not really a boundary condition, but is used to set the Nyquist mode to some value, preferably zero.
886 You should set the Nyquist mode zero when the solution in physical space is real and the resolution is even.
888 Args:
889 kind ('integral' or 'nyquist'): Kind of BC
891 Returns:
892 self.xp.ndarray: Boundary condition row
893 """
894 if kind.lower() == 'integral':
895 return self.get_integ_BC_row()
896 elif kind.lower() == 'nyquist':
897 assert (
898 self.N % 2 == 0
899 ), f'Do not eliminate the Nyquist mode with odd resolution as it is fully resolved. You chose {self.N} in this axis'
900 BC = self.xp.zeros(self.N)
901 BC[self.get_Nyquist_mode_index()] = 1
902 return BC
903 else:
904 return super().get_BC(kind)
906 def get_Nyquist_mode_index(self):
907 """
908 Compute the index of the Nyquist mode, i.e. the mode with the lowest wavenumber, which doesn't have a positive
909 counterpart for even resolution. This means real waves of this wave number cannot be properly resolved and you
910 are best advised to set this mode zero if representing real functions on even-resolution grids is what you're
911 after.
913 Returns:
914 int: Index of the Nyquist mode
915 """
916 k = self.get_wavenumbers()
917 Nyquist_mode = min(k)
918 return self.xp.where(k == Nyquist_mode)[0][0]
920 def get_integ_BC_row(self):
921 """
922 Only the 0-mode has non-zero integral with FFT basis in periodic BCs
923 """
924 me = self.xp.zeros(self.N)
925 me[0] = self.L / self.N
926 return me
929class SpectralHelper:
930 """
931 This class has three functions:
932 - Easily assemble matrices containing multiple equations
933 - Direct product of 1D bases to solve problems in more dimensions
934 - Distribute the FFTs to facilitate concurrency.
936 Attributes:
937 comm (mpi4py.Intracomm): MPI communicator
938 debug (bool): Perform additional checks at extra computational cost
939 useGPU (bool): Whether to use GPUs
940 axes (list): List of 1D bases
941 components (list): List of strings of the names of components in the equations
942 full_BCs (list): List of Dictionaries containing all information about the boundary conditions
943 BC_mat (list): List of lists of sparse matrices to put BCs into and eventually assemble the BC matrix from
944 BCs (sparse matrix): Matrix containing only the BCs
945 fft_cache (dict): Cache FFTs of various shapes here to facilitate padding and so on
946 BC_rhs_mask (self.xp.ndarray): Mask values that contain boundary conditions in the right hand side
947 BC_zero_index (self.xp.ndarray): Indeces of rows in the matrix that are replaced by BCs
948 BC_line_zero_matrix (sparse matrix): Matrix that zeros rows where we can then add the BCs in using `BCs`
949 rhs_BCs_hat (self.xp.ndarray): Boundary conditions in spectral space
950 global_shape (tuple): Global shape of the solution as in `mpi4py-fft`
951 fft_obj: When using distributed FFTs, this will be a parallel transform object from `mpi4py-fft`
952 init (tuple): This is the same `init` that is used throughout the problem classes
953 init_forward (tuple): This is the equivalent of `init` in spectral space
954 """
956 xp = np
957 fft_lib = scipy.fft
958 sparse_lib = scipy.sparse
959 linalg = scipy.sparse.linalg
960 dtype = mesh
961 fft_backend = 'scipy'
962 fft_comm_backend = 'MPI'
964 def setup_GPU(self):
965 """switch to GPU modules"""
966 import cupy as cp
967 import cupyx.scipy.sparse as sparse_lib
968 import cupyx.scipy.sparse.linalg as linalg
969 import cupyx.scipy.fft as fft_lib
970 from pySDC.implementations.datatype_classes.cupy_mesh import cupy_mesh
972 self.xp = cp
973 self.sparse_lib = sparse_lib
974 self.linalg = linalg
976 self.fft_lib = fft_lib
977 self.fft_backend = 'cupyx-scipy'
978 self.fft_comm_backend = 'NCCL'
980 self.dtype = cupy_mesh
982 @classmethod
983 def setup_CPU(cls, useFFTW=False):
984 """switch to CPU modules"""
986 cls.xp = np
987 cls.sparse_lib = scipy.sparse
988 cls.linalg = scipy.sparse.linalg
990 if useFFTW:
991 from mpi4py_fft import fftw
993 cls.fft_backend = 'fftw'
994 cls.fft_lib = fftw
995 else:
996 cls.fft_backend = 'scipy'
997 cls.fft_lib = scipy.fft
999 cls.fft_comm_backend = 'MPI'
1000 cls.dtype = mesh
1002 def __init__(self, comm=None, useGPU=False, debug=False):
1003 """
1004 Constructor
1006 Args:
1007 comm (mpi4py.Intracomm): MPI communicator
1008 useGPU (bool): Whether to use GPUs
1009 debug (bool): Perform additional checks at extra computational cost
1010 """
1011 self.comm = comm
1012 self.debug = debug
1013 self.useGPU = useGPU
1015 if useGPU:
1016 self.setup_GPU()
1017 else:
1018 self.setup_CPU()
1020 self.axes = []
1021 self.components = []
1023 self.full_BCs = []
1024 self.BC_mat = None
1025 self.BCs = None
1027 self.fft_cache = {}
1029 self.logger = logging.getLogger(name='Spectral Discretization')
1030 if debug:
1031 self.logger.setLevel(logging.DEBUG)
1033 @property
1034 def u_init(self):
1035 """
1036 Get empty data container in physical space
1037 """
1038 return self.dtype(self.init)
1040 @property
1041 def u_init_forward(self):
1042 """
1043 Get empty data container in spectral space
1044 """
1045 return self.dtype(self.init_forward)
1047 @property
1048 def u_init_physical(self):
1049 """
1050 Get empty data container in physical space
1051 """
1052 return self.dtype(self.init_physical)
1054 @property
1055 def shape(self):
1056 """
1057 Get shape of individual solution component
1058 """
1059 return self.init[0][1:]
1061 @property
1062 def ndim(self):
1063 return len(self.axes)
1065 @property
1066 def ncomponents(self):
1067 return len(self.components)
1069 @property
1070 def V(self):
1071 """
1072 Get domain volume
1073 """
1074 return np.prod([me.L for me in self.axes])
1076 def add_axis(self, base, *args, **kwargs):
1077 """
1078 Add an axis to the domain by deciding on suitable 1D base.
1079 Arguments to the bases are forwarded using `*args` and `**kwargs`. Please refer to the documentation of the 1D
1080 bases for possible arguments.
1082 Args:
1083 base (str): 1D spectral method
1084 """
1085 kwargs['useGPU'] = self.useGPU
1087 if base.lower() in ['chebychov', 'chebychev', 'cheby', 'chebychovhelper']:
1088 self.axes.append(ChebychevHelper(*args, **kwargs))
1089 elif base.lower() in ['fft', 'fourier', 'ffthelper']:
1090 self.axes.append(FFTHelper(*args, **kwargs))
1091 elif base.lower() in ['ultraspherical', 'gegenbauer']:
1092 self.axes.append(UltrasphericalHelper(*args, **kwargs))
1093 else:
1094 raise NotImplementedError(f'{base=!r} is not implemented!')
1095 self.axes[-1].xp = self.xp
1096 self.axes[-1].sparse_lib = self.sparse_lib
1098 def add_component(self, name):
1099 """
1100 Add solution component(s).
1102 Args:
1103 name (str or list of strings): Name(s) of component(s)
1104 """
1105 if type(name) in [list, tuple]:
1106 for me in name:
1107 self.add_component(me)
1108 elif type(name) in [str]:
1109 if name in self.components:
1110 raise Exception(f'{name=!r} is already added to this problem!')
1111 self.components.append(name)
1112 else:
1113 raise NotImplementedError
1115 def index(self, name):
1116 """
1117 Get the index of component `name`.
1119 Args:
1120 name (str or list of strings): Name(s) of component(s)
1122 Returns:
1123 int: Index of the component
1124 """
1125 if type(name) in [str, int]:
1126 return self.components.index(name)
1127 elif type(name) in [list, tuple]:
1128 return (self.index(me) for me in name)
1129 else:
1130 raise NotImplementedError(f'Don\'t know how to compute index for {type(name)=}')
1132 def get_empty_operator_matrix(self, diag=False):
1133 """
1134 Return a matrix of operators to be filled with the connections between the solution components.
1136 Args:
1137 diag (bool): Whether operator is block-diagonal
1139 Returns:
1140 list containing sparse zeros
1141 """
1142 S = len(self.components)
1143 O = self.get_Id() * 0
1144 if diag:
1145 return [O for _ in range(S)]
1146 else:
1147 return [[O for _ in range(S)] for _ in range(S)]
1149 def get_BC(self, axis, kind, line=-1, scalar=False, **kwargs):
1150 """
1151 Use this method for boundary bordering. It gets the respective matrix row and embeds it into a matrix.
1152 Pay attention that if you have multiple BCs in a single equation, you need to put them in different lines.
1153 Typically, the last line that does not contain a BC is the best choice.
1154 Forward arguments for the boundary conditions using `kwargs`. Refer to documentation of 1D bases for details.
1156 Args:
1157 axis (int): Axis you want to add the BC to
1158 kind (str): kind of BC, e.g. Dirichlet
1159 line (int): Line you want the BC to go in
1160 scalar (bool): Put the BC in all space positions in the other direction
1162 Returns:
1163 sparse matrix containing the BC
1164 """
1165 sp = scipy.sparse
1167 base = self.axes[axis]
1169 BC = sp.eye(base.N).tolil() * 0
1170 if self.useGPU:
1171 BC[line, :] = base.get_BC(kind=kind, **kwargs).get()
1172 else:
1173 BC[line, :] = base.get_BC(kind=kind, **kwargs)
1175 ndim = len(self.axes)
1176 if ndim == 1:
1177 mat = self.sparse_lib.csc_matrix(BC)
1178 elif ndim == 2:
1179 axis2 = (axis + 1) % ndim
1181 if scalar:
1182 _Id = self.sparse_lib.diags(self.xp.append([1], self.xp.zeros(self.axes[axis2].N - 1)))
1183 else:
1184 _Id = self.axes[axis2].get_Id()
1186 Id = self.get_local_slice_of_1D_matrix(self.axes[axis2].get_Id() @ _Id, axis=axis2)
1188 mats = [
1189 None,
1190 ] * ndim
1191 mats[axis] = self.get_local_slice_of_1D_matrix(BC, axis=axis)
1192 mats[axis2] = Id
1193 mat = self.sparse_lib.csc_matrix(self.sparse_lib.kron(*mats))
1194 elif ndim == 3:
1195 mats = [
1196 None,
1197 ] * ndim
1199 for ax in range(ndim):
1200 if ax == axis:
1201 continue
1203 if scalar:
1204 _Id = self.sparse_lib.diags(self.xp.append([1], self.xp.zeros(self.axes[ax].N - 1)))
1205 else:
1206 _Id = self.axes[ax].get_Id()
1208 mats[ax] = self.get_local_slice_of_1D_matrix(self.axes[ax].get_Id() @ _Id, axis=ax)
1210 mats[axis] = self.get_local_slice_of_1D_matrix(BC, axis=axis)
1212 mat = self.sparse_lib.csc_matrix(self.sparse_lib.kron(mats[0], self.sparse_lib.kron(*mats[1:])))
1213 else:
1214 raise NotImplementedError(
1215 f'Matrix expansion for boundary conditions not implemented for {ndim} dimensions!'
1216 )
1217 mat = self.eliminate_zeros(mat)
1218 return mat
1220 def remove_BC(self, component, equation, axis, kind, line=-1, scalar=False, **kwargs):
1221 """
1222 Remove a BC from the matrix. This is useful e.g. when you add a non-scalar BC and then need to selectively
1223 remove single BCs again, as in incompressible Navier-Stokes, for instance.
1224 Forwards arguments for the boundary conditions using `kwargs`. Refer to documentation of 1D bases for details.
1226 Args:
1227 component (str): Name of the component the BC should act on
1228 equation (str): Name of the equation for the component you want to put the BC in
1229 axis (int): Axis you want to add the BC to
1230 kind (str): kind of BC, e.g. Dirichlet
1231 v: Value of the BC
1232 line (int): Line you want the BC to go in
1233 scalar (bool): Put the BC in all space positions in the other direction
1234 """
1235 _BC = self.get_BC(axis=axis, kind=kind, line=line, scalar=scalar, **kwargs)
1236 _BC = self.eliminate_zeros(_BC)
1237 self.BC_mat[self.index(equation)][self.index(component)] -= _BC
1239 if scalar:
1240 slices = [self.index(equation)] + [
1241 0,
1242 ] * self.ndim
1243 slices[axis + 1] = line
1244 else:
1245 slices = (
1246 [self.index(equation)]
1247 + [slice(0, self.init[0][i + 1]) for i in range(axis)]
1248 + [line]
1249 + [slice(0, self.init[0][i + 1]) for i in range(axis + 1, len(self.axes))]
1250 )
1251 N = self.axes[axis].N
1252 if (N + line) % N in self.xp.arange(N)[self.local_slice()[axis]]:
1253 self.BC_rhs_mask[(*slices,)] = False
1255 def add_BC(self, component, equation, axis, kind, v, line=-1, scalar=False, **kwargs):
1256 """
1257 Add a BC to the matrix. Note that you need to convert the list of lists of BCs that this method generates to a
1258 single sparse matrix by calling `setup_BCs` after adding/removing all BCs.
1259 Forward arguments for the boundary conditions using `kwargs`. Refer to documentation of 1D bases for details.
1261 Args:
1262 component (str): Name of the component the BC should act on
1263 equation (str): Name of the equation for the component you want to put the BC in
1264 axis (int): Axis you want to add the BC to
1265 kind (str): kind of BC, e.g. Dirichlet
1266 v: Value of the BC
1267 line (int): Line you want the BC to go in
1268 scalar (bool): Put the BC in all space positions in the other direction
1269 """
1270 _BC = self.get_BC(axis=axis, kind=kind, line=line, scalar=scalar, **kwargs)
1271 self.BC_mat[self.index(equation)][self.index(component)] += _BC
1272 self.full_BCs += [
1273 {
1274 'component': component,
1275 'equation': equation,
1276 'axis': axis,
1277 'kind': kind,
1278 'v': v,
1279 'line': line,
1280 'scalar': scalar,
1281 **kwargs,
1282 }
1283 ]
1285 if scalar:
1286 slices = [self.index(equation)] + [
1287 0,
1288 ] * self.ndim
1289 slices[axis + 1] = line
1290 if self.comm:
1291 if self.comm.rank == 0:
1292 self.BC_rhs_mask[(*slices,)] = True
1293 else:
1294 self.BC_rhs_mask[(*slices,)] = True
1295 else:
1296 slices = [self.index(equation), *self.global_slice(True)]
1297 N = self.axes[axis].N
1298 if (N + line) % N in self.get_indices(True)[axis]:
1299 slices[axis + 1] = (N + line) % N - self.local_slice()[axis].start
1300 self.BC_rhs_mask[(*slices,)] = True
1302 def setup_BCs(self):
1303 """
1304 Convert the list of lists of BCs to the boundary condition operator.
1305 Also, boundary bordering requires to zero out all other entries in the matrix in rows containing a boundary
1306 condition. This method sets up a suitable sparse matrix to do this.
1307 """
1308 sp = self.sparse_lib
1309 self.BCs = self.convert_operator_matrix_to_operator(self.BC_mat)
1310 self.BC_zero_index = self.xp.arange(np.prod(self.init[0]))[self.BC_rhs_mask.flatten()]
1312 diags = self.xp.ones(self.BCs.shape[0])
1313 diags[self.BC_zero_index] = 0
1314 self.BC_line_zero_matrix = sp.diags(diags).tocsc()
1316 # prepare BCs in spectral space to easily add to the RHS
1317 rhs_BCs = self.put_BCs_in_rhs(self.u_init)
1318 self.rhs_BCs_hat = self.transform(rhs_BCs).view(self.xp.ndarray)
1319 del self.BC_rhs_mask
1321 def check_BCs(self, u):
1322 """
1323 Check that the solution satisfies the boundary conditions
1325 Args:
1326 u: The solution you want to check
1327 """
1328 assert self.ndim < 3
1329 for axis in range(self.ndim):
1330 BCs = [me for me in self.full_BCs if me["axis"] == axis and not me["scalar"]]
1332 if len(BCs) > 0:
1333 u_hat = self.transform(u, axes=(axis - self.ndim,))
1334 for BC in BCs:
1335 kwargs = {
1336 key: value
1337 for key, value in BC.items()
1338 if key not in ['component', 'equation', 'axis', 'v', 'line', 'scalar']
1339 }
1341 if axis == 0:
1342 get = self.axes[axis].get_BC(**kwargs) @ u_hat[self.index(BC['component'])]
1343 elif axis == 1:
1344 get = u_hat[self.index(BC['component'])] @ self.axes[axis].get_BC(**kwargs)
1345 want = BC['v']
1346 assert self.xp.allclose(
1347 get, want
1348 ), f'Unexpected BC in {BC["component"]} in equation {BC["equation"]}, line {BC["line"]}! Got {get}, wanted {want}'
1350 def put_BCs_in_matrix(self, A):
1351 """
1352 Put the boundary conditions in a matrix by replacing rows with BCs.
1353 """
1354 return self.BC_line_zero_matrix @ A + self.BCs
1356 def put_BCs_in_rhs_hat(self, rhs_hat):
1357 """
1358 Put the BCs in the right hand side in spectral space for solving.
1359 This function needs no transforms and caches a mask for faster subsequent use.
1361 Args:
1362 rhs_hat: Right hand side in spectral space
1364 Returns:
1365 rhs in spectral space with BCs
1366 """
1367 if not hasattr(self, '_rhs_hat_zero_mask'):
1368 """
1369 Generate a mask where we need to set values in the rhs in spectral space to zero, such that can replace them
1370 by the boundary conditions. The mask is then cached.
1371 """
1372 self._rhs_hat_zero_mask = self.newDistArray(forward_output=True).astype(bool).view(self.xp.ndarray)
1374 for axis in range(self.ndim):
1375 for bc in self.full_BCs:
1376 if axis == bc['axis']:
1377 slices = [self.index(bc['equation']), *self.global_slice(True)]
1378 N = self.axes[axis].N
1379 line = bc['line']
1380 if (N + line) % N in self.get_indices(True)[axis]:
1381 slices[axis + 1] = (N + line) % N - self.local_slice()[axis].start
1382 self._rhs_hat_zero_mask[(*slices,)] = True
1384 rhs_hat[self._rhs_hat_zero_mask] = 0
1385 return rhs_hat + self.rhs_BCs_hat
1387 def put_BCs_in_rhs(self, rhs):
1388 """
1389 Put the BCs in the right hand side for solving.
1390 This function will transform along each axis individually and add all BCs in that axis.
1391 Consider `put_BCs_in_rhs_hat` to add BCs with no extra transforms needed.
1393 Args:
1394 rhs: Right hand side in physical space
1396 Returns:
1397 rhs in physical space with BCs
1398 """
1399 assert rhs.ndim > 1, 'rhs must not be flattened here!'
1401 ndim = self.ndim
1403 for axis in range(ndim):
1404 _rhs_hat = self.transform(rhs, axes=(axis - ndim,))
1406 for bc in self.full_BCs:
1408 if axis == bc['axis']:
1409 _slice = [self.index(bc['equation']), *self.global_slice(True)]
1411 N = self.axes[axis].N
1412 line = bc['line']
1413 if (N + line) % N in self.get_indices(True)[axis]:
1414 _slice[axis + 1] = (N + line) % N - self.local_slice()[axis].start
1415 _rhs_hat[(*_slice,)] = bc['v']
1417 rhs = self.itransform(_rhs_hat, axes=(axis - ndim,))
1419 return rhs
1421 def add_equation_lhs(self, A, equation, relations):
1422 """
1423 Add the left hand part (that you want to solve implicitly) of an equation to a list of lists of sparse matrices
1424 that you will convert to an operator later.
1426 Example:
1427 Setup linear operator `L` for 1D heat equation using Chebychev method in first order form and T-to-U
1428 preconditioning:
1430 .. code-block:: python
1431 helper = SpectralHelper()
1433 helper.add_axis(base='chebychev', N=8)
1434 helper.add_component(['u', 'ux'])
1435 helper.setup_fft()
1437 I = helper.get_Id()
1438 Dx = helper.get_differentiation_matrix(axes=(0,))
1439 T2U = helper.get_basis_change_matrix('T2U')
1441 L_lhs = {
1442 'ux': {'u': -T2U @ Dx, 'ux': T2U @ I},
1443 'u': {'ux': -(T2U @ Dx)},
1444 }
1446 operator = helper.get_empty_operator_matrix()
1447 for line, equation in L_lhs.items():
1448 helper.add_equation_lhs(operator, line, equation)
1450 L = helper.convert_operator_matrix_to_operator(operator)
1452 Args:
1453 A (list of lists of sparse matrices): The operator to be
1454 equation (str): The equation of the component you want this in
1455 relations: (dict): Relations between quantities
1456 """
1457 for k, v in relations.items():
1458 A[self.index(equation)][self.index(k)] = v
1460 def eliminate_zeros(self, A):
1461 """
1462 Eliminate zeros from sparse matrix. This can reduce memory footprint of matrices somewhat.
1463 Note: At the time of writing, there are memory problems in the cupy implementation of `eliminate_zeros`.
1464 Therefore, this function copies the matrix to host, eliminates the zeros there and then copies back to GPU.
1466 Args:
1467 A: sparse matrix to be pruned
1469 Returns:
1470 CSC sparse matrix
1471 """
1472 if self.useGPU:
1473 A = A.get()
1474 A = A.tocsc()
1475 A.eliminate_zeros()
1476 if self.useGPU:
1477 A = self.sparse_lib.csc_matrix(A)
1478 return A
1480 def convert_operator_matrix_to_operator(self, M):
1481 """
1482 Promote the list of lists of sparse matrices to a single sparse matrix that can be used as linear operator.
1483 See documentation of `SpectralHelper.add_equation_lhs` for an example.
1485 Args:
1486 M (list of lists of sparse matrices): The operator to be
1488 Returns:
1489 sparse linear operator
1490 """
1491 if len(self.components) == 1:
1492 op = M[0][0]
1493 else:
1494 op = self.sparse_lib.bmat(M, format='csc')
1496 op = self.eliminate_zeros(op)
1497 return op
1499 def get_wavenumbers(self):
1500 """
1501 Get grid in spectral space
1502 """
1503 grids = [self.axes[i].get_wavenumbers()[self.local_slice(True)[i]] for i in range(len(self.axes))]
1504 return self.xp.meshgrid(*grids, indexing='ij')
1506 def get_grid(self, forward_output=False):
1507 """
1508 Get grid in physical space
1509 """
1510 grids = [self.axes[i].get_1dgrid()[self.local_slice(forward_output)[i]] for i in range(len(self.axes))]
1511 return self.xp.meshgrid(*grids, indexing='ij')
1513 def get_indices(self, forward_output=True):
1514 return [self.xp.arange(self.axes[i].N)[self.local_slice(forward_output)[i]] for i in range(len(self.axes))]
1516 @cache
1517 def get_pfft(self, axes=None, padding=None, grid=None):
1518 if self.ndim == 1 or self.comm is None:
1519 return None
1520 from mpi4py_fft import newDistArray
1522 from pySDC.helpers.fft_helper import PFFT
1524 axes = tuple(i for i in range(self.ndim)) if axes is None else axes
1525 padding = list(padding if padding else [1.0 for _ in range(self.ndim)])
1527 def no_transform(u, *args, **kwargs):
1528 return u
1530 transforms = {(i,): (no_transform, no_transform) for i in range(self.ndim)}
1531 for i in axes:
1532 transforms[((i + self.ndim) % self.ndim,)] = (self.axes[i].transform, self.axes[i].itransform)
1534 # "transform" all axes to ensure consistent shapes.
1535 # Transform non-distributable axes last to ensure they are aligned
1536 _axes = tuple(sorted((axis + self.ndim) % self.ndim for axis in axes))
1537 _axes = [axis for axis in _axes if not self.axes[axis].distributable] + sorted(
1538 [axis for axis in _axes if self.axes[axis].distributable]
1539 + [axis for axis in range(self.ndim) if axis not in _axes]
1540 )
1542 pfft = PFFT(
1543 comm=self.comm,
1544 shape=self.global_shape[1:],
1545 axes=_axes, # TODO: control the order of the transforms better
1546 dtype='D',
1547 collapse=False,
1548 backend=self.fft_backend,
1549 comm_backend=self.fft_comm_backend,
1550 padding=padding,
1551 transforms=transforms,
1552 grid=grid,
1553 )
1555 # do a transform to do the planning
1556 _u = newDistArray(pfft, forward_output=False)
1557 pfft.backward(pfft.forward(_u))
1558 return pfft
1560 def get_fft(self, axes=None, direction='object', padding=None, shape=None):
1561 """
1562 When using MPI, we use `PFFT` objects generated by mpi4py-fft
1564 Args:
1565 axes (tuple): Axes you want to transform over
1566 direction (str): use "forward" or "backward" to get functions for performing the transforms or "object" to get the PFFT object
1567 padding (tuple): Padding for dealiasing
1568 shape (tuple): Shape of the transform
1570 Returns:
1571 transform
1572 """
1573 axes = tuple(-i - 1 for i in range(self.ndim)) if axes is None else axes
1574 shape = self.global_shape[1:] if shape is None else shape
1575 padding = (
1576 [
1577 1,
1578 ]
1579 * self.ndim
1580 if padding is None
1581 else padding
1582 )
1583 key = (axes, direction, tuple(padding), tuple(shape))
1585 if key not in self.fft_cache.keys():
1586 if self.comm is None:
1587 assert np.allclose(padding, 1), 'Zero padding is not implemented for non-MPI transforms'
1589 if direction == 'forward':
1590 self.fft_cache[key] = self.xp.fft.fftn
1591 elif direction == 'backward':
1592 self.fft_cache[key] = self.xp.fft.ifftn
1593 elif direction == 'object':
1594 self.fft_cache[key] = None
1595 else:
1596 if direction == 'object':
1597 from pySDC.helpers.fft_helper import PFFT
1599 _fft = PFFT(
1600 comm=self.comm,
1601 shape=shape,
1602 axes=sorted(axes),
1603 dtype='D',
1604 collapse=False,
1605 backend=self.fft_backend,
1606 comm_backend=self.fft_comm_backend,
1607 padding=padding,
1608 )
1609 else:
1610 _fft = self.get_fft(axes=axes, direction='object', padding=padding, shape=shape)
1612 if direction == 'forward':
1613 self.fft_cache[key] = _fft.forward
1614 elif direction == 'backward':
1615 self.fft_cache[key] = _fft.backward
1616 elif direction == 'object':
1617 self.fft_cache[key] = _fft
1619 return self.fft_cache[key]
1621 def local_slice(self, forward_output=True):
1622 if self.fft_obj:
1623 return self.get_pfft().local_slice(forward_output=forward_output)
1624 else:
1625 return [slice(0, me.N) for me in self.axes]
1627 def global_slice(self, forward_output=True):
1628 if self.fft_obj:
1629 return [slice(0, me) for me in self.fft_obj.global_shape(forward_output=forward_output)]
1630 else:
1631 return self.local_slice(forward_output=forward_output)
1633 def setup_fft(self, real_spectral_coefficients=False):
1634 """
1635 This function must be called after all axes have been setup in order to prepare the local shapes of the data.
1636 This must also be called before setting up any BCs.
1638 Args:
1639 real_spectral_coefficients (bool): Allow only real coefficients in spectral space
1640 """
1641 if len(self.components) == 0:
1642 self.add_component('u')
1644 self.global_shape = (len(self.components),) + tuple(me.N for me in self.axes)
1646 axes = tuple(i for i in range(len(self.axes)))
1647 self.fft_obj = self.get_pfft(axes=axes)
1649 self.init = (
1650 np.empty(shape=self.global_shape)[
1651 (
1652 ...,
1653 *self.local_slice(False),
1654 )
1655 ].shape,
1656 self.comm,
1657 np.dtype('float'),
1658 )
1659 self.init_physical = (
1660 np.empty(shape=self.global_shape)[
1661 (
1662 ...,
1663 *self.local_slice(False),
1664 )
1665 ].shape,
1666 self.comm,
1667 np.dtype('float'),
1668 )
1669 self.init_forward = (
1670 np.empty(shape=self.global_shape)[
1671 (
1672 ...,
1673 *self.local_slice(True),
1674 )
1675 ].shape,
1676 self.comm,
1677 np.dtype('float') if real_spectral_coefficients else np.dtype('complex128'),
1678 )
1680 self.BC_mat = self.get_empty_operator_matrix()
1681 self.BC_rhs_mask = self.newDistArray().astype(bool)
1683 def newDistArray(self, pfft=None, forward_output=True, val=0, rank=1, view=False):
1684 """
1685 Get an empty distributed array. This is almost a copy of the function of the same name from mpi4py-fft, but
1686 takes care of all the solution components in the tensor.
1687 """
1688 if self.comm is None:
1689 return self.xp.zeros(self.init[0], dtype=self.init[2])
1690 from mpi4py_fft.distarray import DistArray
1692 pfft = pfft if pfft else self.get_pfft()
1693 if pfft is None:
1694 if forward_output:
1695 return self.u_init_forward
1696 else:
1697 return self.u_init
1699 global_shape = pfft.global_shape(forward_output)
1700 p0 = pfft.pencil[forward_output]
1701 if forward_output is True:
1702 dtype = pfft.forward.output_array.dtype
1703 else:
1704 dtype = pfft.forward.input_array.dtype
1705 global_shape = (self.ncomponents,) * rank + global_shape
1707 if pfft.xfftn[0].backend in ["cupy", "cupyx-scipy"]:
1708 from mpi4py_fft.distarrayCuPy import DistArrayCuPy as darraycls
1709 else:
1710 darraycls = DistArray
1712 z = darraycls(global_shape, subcomm=p0.subcomm, val=val, dtype=dtype, alignment=p0.axis, rank=rank)
1713 return z.v if view else z
1715 def infer_alignment(self, u, forward_output, padding=None, **kwargs):
1716 if self.comm is None:
1717 return [0]
1719 def _alignment(pfft):
1720 _arr = self.newDistArray(pfft, forward_output=forward_output)
1721 _aligned_axes = [i for i in range(self.ndim) if _arr.global_shape[i + 1] == u.shape[i + 1]]
1722 return _aligned_axes
1724 if padding is None:
1725 pfft = self.get_pfft(**kwargs)
1726 aligned_axes = _alignment(pfft)
1727 else:
1728 if self.ndim == 2:
1729 padding_options = [(1.0, padding[1]), (padding[0], 1.0), padding, (1.0, 1.0)]
1730 elif self.ndim == 3:
1731 padding_options = [
1732 (1.0, 1.0, padding[2]),
1733 (1.0, padding[1], 1.0),
1734 (padding[0], 1.0, 1.0),
1735 (1.0, padding[1], padding[2]),
1736 (padding[0], 1.0, padding[2]),
1737 (padding[0], padding[1], 1.0),
1738 padding,
1739 (1.0, 1.0, 1.0),
1740 ]
1741 else:
1742 raise NotImplementedError(f'Don\'t know how to infer alignment in {self.ndim}D!')
1743 for _padding in padding_options:
1744 pfft = self.get_pfft(padding=_padding, **kwargs)
1745 aligned_axes = _alignment(pfft)
1746 if len(aligned_axes) > 0:
1747 self.logger.debug(
1748 f'Found alignment of array with size {u.shape}: {aligned_axes} using padding {_padding}'
1749 )
1750 break
1752 assert len(aligned_axes) > 0, f'Found no aligned axes for array of size {u.shape}!'
1753 return aligned_axes
1755 def redistribute(self, u, axis, forward_output, **kwargs):
1756 if self.comm is None:
1757 return u
1759 pfft = self.get_pfft(**kwargs)
1760 _arr = self.newDistArray(pfft, forward_output=forward_output)
1762 if 'Dist' in type(u).__name__ and False:
1763 try:
1764 u.redistribute(out=_arr)
1765 return _arr
1766 except AssertionError:
1767 pass
1769 u_alignment = self.infer_alignment(u, forward_output=False, **kwargs)
1770 for alignment in u_alignment:
1771 _arr = _arr.redistribute(alignment)
1772 if _arr.shape == u.shape:
1773 _arr[...] = u
1774 return _arr.redistribute(axis)
1776 raise Exception(
1777 f'Don\'t know how to align array of local shape {u.shape} and global shape {self.global_shape}, aligned in axes {u_alignment}, to axis {axis}'
1778 )
1780 def transform(self, u, *args, axes=None, padding=None, **kwargs):
1781 pfft = self.get_pfft(*args, axes=axes, padding=padding, **kwargs)
1783 if pfft is None:
1784 axes = axes if axes else tuple(i for i in range(self.ndim))
1785 u_hat = u.copy()
1786 for i in axes:
1787 _axis = 1 + i if i >= 0 else i
1788 u_hat = self.axes[i].transform(u_hat, axes=(_axis,))
1789 return u_hat
1791 _in = self.newDistArray(pfft, forward_output=False, rank=1)
1792 _out = self.newDistArray(pfft, forward_output=True, rank=1)
1794 if _in.shape == u.shape:
1795 _in[...] = u
1796 else:
1797 _in[...] = self.redistribute(u, axis=_in.alignment, forward_output=False, padding=padding, **kwargs)
1799 for i in range(self.ncomponents):
1800 pfft.forward(_in[i], _out[i], normalize=False)
1802 if padding is not None:
1803 _out /= np.prod(padding)
1804 return _out
1806 def itransform(self, u, *args, axes=None, padding=None, **kwargs):
1807 if padding is not None:
1808 assert all(
1809 (self.axes[i].N * padding[i]) % 1 == 0 for i in range(self.ndim)
1810 ), 'Cannot do this padding with this resolution. Resulting resolution must be integer'
1812 pfft = self.get_pfft(*args, axes=axes, padding=padding, **kwargs)
1813 if pfft is None:
1814 axes = axes if axes else tuple(i for i in range(self.ndim))
1815 u_hat = u.copy()
1816 for i in axes:
1817 _axis = 1 + i if i >= 0 else i
1818 u_hat = self.axes[i].itransform(u_hat, axes=(_axis,))
1819 return u_hat
1821 _in = self.newDistArray(pfft, forward_output=True, rank=1)
1822 _out = self.newDistArray(pfft, forward_output=False, rank=1)
1824 if _in.shape == u.shape:
1825 _in[...] = u
1826 else:
1827 _in[...] = self.redistribute(u, axis=_in.alignment, forward_output=True, padding=padding, **kwargs)
1829 for i in range(self.ncomponents):
1830 pfft.backward(_in[i], _out[i], normalize=True)
1832 if padding is not None:
1833 _out *= np.prod(padding)
1834 return _out
1836 def get_local_slice_of_1D_matrix(self, M, axis):
1837 """
1838 Get the local version of a 1D matrix. When using distributed FFTs, each rank will carry only a subset of modes,
1839 which you can sort out via the `SpectralHelper.local_slice()` attribute. When constructing a 1D matrix, you can
1840 use this method to get the part corresponding to the modes carried by this rank.
1842 Args:
1843 M (sparse matrix): Global 1D matrix you want to get the local version of
1844 axis (int): Direction in which you want the local version. You will get the global matrix in other directions.
1846 Returns:
1847 sparse local matrix
1848 """
1849 return M.tocsc()[self.local_slice(True)[axis], self.local_slice(True)[axis]]
1851 def expand_matrix_ND(self, matrix, aligned):
1852 sp = self.sparse_lib
1853 axes = np.delete(np.arange(self.ndim), aligned)
1854 ndim = len(axes) + 1
1856 if ndim == 1:
1857 mat = matrix
1858 elif ndim == 2:
1859 axis = axes[0]
1860 I1D = sp.eye(self.axes[axis].N)
1862 mats = [None] * ndim
1863 mats[aligned] = self.get_local_slice_of_1D_matrix(matrix, aligned)
1864 mats[axis] = self.get_local_slice_of_1D_matrix(I1D, axis)
1866 mat = sp.kron(*mats)
1867 elif ndim == 3:
1869 mats = [None] * ndim
1870 mats[aligned] = self.get_local_slice_of_1D_matrix(matrix, aligned)
1871 for axis in axes:
1872 I1D = sp.eye(self.axes[axis].N)
1873 mats[axis] = self.get_local_slice_of_1D_matrix(I1D, axis)
1875 mat = sp.kron(mats[0], sp.kron(*mats[1:]))
1877 else:
1878 raise NotImplementedError(f'Matrix expansion not implemented for {ndim} dimensions!')
1880 mat = self.eliminate_zeros(mat)
1881 return mat
1883 def get_filter_matrix(self, axis, **kwargs):
1884 """
1885 Get bandpass filter along `axis`. See the documentation `get_filter_matrix` in the 1D bases for what kwargs are
1886 admissible.
1888 Returns:
1889 sparse bandpass matrix
1890 """
1891 if self.ndim == 1:
1892 return self.axes[0].get_filter_matrix(**kwargs)
1894 mats = [base.get_Id() for base in self.axes]
1895 mats[axis] = self.axes[axis].get_filter_matrix(**kwargs)
1896 return self.sparse_lib.kron(*mats)
1898 def get_differentiation_matrix(self, axes, **kwargs):
1899 """
1900 Get differentiation matrix along specified axis. `kwargs` are forwarded to the 1D base implementation.
1902 Args:
1903 axes (tuple): Axes along which to differentiate.
1905 Returns:
1906 sparse differentiation matrix
1907 """
1908 D = self.expand_matrix_ND(self.axes[axes[0]].get_differentiation_matrix(**kwargs), axes[0])
1909 for axis in axes[1:]:
1910 _D = self.axes[axis].get_differentiation_matrix(**kwargs)
1911 D = D @ self.expand_matrix_ND(_D, axis)
1913 self.logger.debug(f'Set up differentiation matrix along axes {axes} with kwargs {kwargs}')
1914 return D
1916 def get_integration_matrix(self, axes):
1917 """
1918 Get integration matrix to integrate along specified axis.
1920 Args:
1921 axes (tuple): Axes along which to integrate over.
1923 Returns:
1924 sparse integration matrix
1925 """
1926 S = self.expand_matrix_ND(self.axes[axes[0]].get_integration_matrix(), axes[0])
1927 for axis in axes[1:]:
1928 _S = self.axes[axis].get_integration_matrix()
1929 S = S @ self.expand_matrix_ND(_S, axis)
1931 return S
1933 def get_Id(self):
1934 """
1935 Get identity matrix
1937 Returns:
1938 sparse identity matrix
1939 """
1940 I = self.expand_matrix_ND(self.axes[0].get_Id(), 0)
1941 for axis in range(1, self.ndim):
1942 _I = self.axes[axis].get_Id()
1943 I = I @ self.expand_matrix_ND(_I, axis)
1944 return I
1946 def get_Dirichlet_recombination_matrix(self, axis=-1):
1947 """
1948 Get Dirichlet recombination matrix along axis. Not that it only makes sense in directions discretized with variations of Chebychev bases.
1950 Args:
1951 axis (int): Axis you discretized with Chebychev
1953 Returns:
1954 sparse matrix
1955 """
1956 C1D = self.axes[axis].get_Dirichlet_recombination_matrix()
1957 return self.expand_matrix_ND(C1D, axis)
1959 def get_basis_change_matrix(self, axes=None, **kwargs):
1960 """
1961 Some spectral bases do a change between bases while differentiating. This method returns matrices that changes the basis to whatever you want.
1962 Refer to the methods of the same name of the 1D bases to learn what parameters you need to pass here as `kwargs`.
1964 Args:
1965 axes (tuple): Axes along which to change basis.
1967 Returns:
1968 sparse basis change matrix
1969 """
1970 axes = tuple(-i - 1 for i in range(self.ndim)) if axes is None else axes
1972 C = self.expand_matrix_ND(self.axes[axes[0]].get_basis_change_matrix(**kwargs), axes[0])
1973 for axis in axes[1:]:
1974 _C = self.axes[axis].get_basis_change_matrix(**kwargs)
1975 C = C @ self.expand_matrix_ND(_C, axis)
1977 self.logger.debug(f'Set up basis change matrix along axes {axes} with kwargs {kwargs}')
1978 return C