Coverage for pySDC/helpers/spectral_helper.py: 93%
796 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +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 def setup_CPU(self, useFFTW=False):
121 """switch to CPU modules"""
123 self.xp = np
124 self.sparse_lib = scipy.sparse
125 self.linalg = scipy.sparse.linalg
127 if useFFTW:
128 from mpi4py_fft import fftw
130 self.fft_backend = 'fftw'
131 self.fft_lib = fftw
132 else:
133 self.fft_backend = 'scipy'
134 self.fft_lib = scipy.fft
136 self.fft_comm_backend = 'MPI'
137 self.dtype = mesh
139 def get_Id(self):
140 """
141 Get identity matrix
143 Returns:
144 sparse diagonal identity matrix
145 """
146 return self.sparse_lib.eye(self.N)
148 def get_zero(self):
149 """
150 Get a matrix with all zeros of the correct size.
152 Returns:
153 sparse matrix with zeros everywhere
154 """
155 return 0 * self.get_Id()
157 def get_differentiation_matrix(self):
158 """
159 Get the differentiation matrix of this basis. Implemented by the derived bases.
161 Returns:
162 sparse differentiation matrix
163 """
164 raise NotImplementedError()
166 def get_integration_matrix(self):
167 """
168 Get the integration matrix of this basis. Implemented by the derived bases.
170 Returns:
171 sparse integration matrix
172 """
173 raise NotImplementedError()
175 def get_integration_weights(self):
176 """Weights for integration across entire domain"""
177 raise NotImplementedError()
179 def get_wavenumbers(self):
180 """
181 Get the grid in spectral space
182 """
183 raise NotImplementedError
185 def get_empty_operator_matrix(self, S, O):
186 """
187 Return a matrix of operators to be filled with the connections between the solution components.
189 Args:
190 S (int): Number of components in the solution
191 O (sparse matrix): Zero matrix used for initialization
193 Returns:
194 list of lists containing sparse zeros
195 """
196 return [[O for _ in range(S)] for _ in range(S)]
198 def get_basis_change_matrix(self, *args, **kwargs):
199 """
200 Some spectral discretization change the basis during differentiation. This method can be used to transfer
201 between the various bases.
203 This method accepts arbitrary arguments that may not be used in order to provide an easy interface for multi-
204 dimensional bases. For instance, you may combine an FFT discretization with an ultraspherical discretization.
205 The FFT discretization will always be in the same base, but the ultraspherical discretization uses a different
206 base for every derivative. You can then ask all bases for transfer matrices from one ultraspherical derivative
207 base to the next. The FFT discretization will ignore this and return an identity while the ultraspherical
208 discretization will return the desired matrix. After a Kronecker product, you get the 2D version of the matrix
209 you want. This is what the `SpectralHelper` does when you call the method of the same name on it.
211 Returns:
212 sparse bases change matrix
213 """
214 return self.sparse_lib.eye(self.N)
216 def get_BC(self, kind):
217 """
218 To facilitate boundary conditions (BCs) we use either a basis where all functions satisfy the BCs automatically,
219 as is the case in FFT basis for periodic BCs, or boundary bordering. In boundary bordering, specific lines in
220 the matrix are replaced by the boundary conditions as obtained by this method.
222 Args:
223 kind (str): The type of BC you want to implement please refer to the implementations of this method in the
224 individual 1D bases for what is implemented
226 Returns:
227 self.xp.array: Boundary condition
228 """
229 raise NotImplementedError(f'No boundary conditions of {kind=!r} implemented!')
231 def get_filter_matrix(self, kmin=0, kmax=None):
232 """
233 Get a bandpass filter.
235 Args:
236 kmin (int): Lower limit of the bandpass filter
237 kmax (int): Upper limit of the bandpass filter
239 Returns:
240 sparse matrix
241 """
243 k = abs(self.get_wavenumbers())
245 kmax = max(k) if kmax is None else kmax
247 mask = self.xp.logical_or(k >= kmax, k < kmin)
249 if self.useGPU:
250 Id = self.get_Id().get()
251 else:
252 Id = self.get_Id()
253 F = Id.tolil()
254 F[:, mask] = 0
255 return F.tocsc()
257 def get_1dgrid(self):
258 """
259 Get the grid in physical space
261 Returns:
262 self.xp.array: Grid
263 """
264 raise NotImplementedError
267class ChebychevHelper(SpectralHelper1D):
268 """
269 The Chebychev base consists of special kinds of polynomials, with the main advantage that you can easily transform
270 between physical and spectral space by discrete cosine transform.
271 The differentiation in the Chebychev T base is dense, but can be preconditioned to yield a differentiation operator
272 that moves to Chebychev U basis during differentiation, which is sparse. When using this technique, problems need to
273 be formulated in first order formulation.
275 This implementation is largely based on the Dedalus paper (https://doi.org/10.1103/PhysRevResearch.2.023068).
276 """
278 def __init__(self, *args, x0=-1, x1=1, **kwargs):
279 """
280 Constructor.
281 Please refer to the parent class for additional arguments. Notably, you have to supply a resolution `N` and you
282 may choose to run on GPUs via the `useGPU` argument.
284 Args:
285 x0 (float): Coordinate of left boundary. Note that only -1 is currently implented
286 x1 (float): Coordinate of right boundary. Note that only +1 is currently implented
287 """
288 # need linear transformation y = ax + b with a = (x1-x0)/2 and b = (x1+x0)/2
289 self.lin_trf_fac = (x1 - x0) / 2
290 self.lin_trf_off = (x1 + x0) / 2
291 super().__init__(*args, x0=x0, x1=x1, **kwargs)
293 self.norm = self.get_norm()
295 def get_1dgrid(self):
296 '''
297 Generates a 1D grid with Chebychev points. These are clustered at the boundary. You need this kind of grid to
298 use discrete cosine transformation (DCT) to get the Chebychev representation. If you want a different grid, you
299 need to do an affine transformation before any Chebychev business.
301 Returns:
302 numpy.ndarray: 1D grid
303 '''
304 return self.lin_trf_fac * self.xp.cos(np.pi / self.N * (self.xp.arange(self.N) + 0.5)) + self.lin_trf_off
306 def get_wavenumbers(self):
307 """Get the domain in spectral space"""
308 return self.xp.arange(self.N)
310 @cache
311 def get_conv(self, name, N=None):
312 '''
313 Get conversion matrix between different kinds of polynomials. The supported kinds are
314 - T: Chebychev polynomials of first kind
315 - U: Chebychev polynomials of second kind
316 - D: Dirichlet recombination.
318 You get the desired matrix by choosing a name as ``A2B``. I.e. ``T2U`` for the conversion matrix from T to U.
319 Once generates matrices are cached. So feel free to call the method as often as you like.
321 Args:
322 name (str): Conversion code, e.g. 'T2U'
323 N (int): Size of the matrix (optional)
325 Returns:
326 scipy.sparse: Sparse conversion matrix
327 '''
328 N = N if N else self.N
329 sp = self.sparse_lib
331 def get_forward_conv(name):
332 if name == 'T2U':
333 mat = (sp.eye(N) - sp.eye(N, k=2)).tocsc() / 2.0
334 mat[:, 0] *= 2
335 elif name == 'D2T':
336 mat = sp.eye(N) - sp.eye(N, k=2)
337 elif name[0] == name[-1]:
338 mat = self.sparse_lib.eye(self.N)
339 else:
340 raise NotImplementedError(f'Don\'t have conversion matrix {name!r}')
341 return mat
343 try:
344 mat = get_forward_conv(name)
345 except NotImplementedError as E:
346 try:
347 fwd = get_forward_conv(name[::-1])
348 import scipy.sparse as sp
350 if self.sparse_lib == sp:
351 mat = self.sparse_lib.linalg.inv(fwd.tocsc())
352 else:
353 mat = self.sparse_lib.csc_matrix(sp.linalg.inv(fwd.tocsc().get()))
354 except NotImplementedError:
355 raise NotImplementedError from E
357 return mat
359 def get_basis_change_matrix(self, conv='T2T', **kwargs):
360 """
361 As the differentiation matrix in Chebychev-T base is dense but is sparse when simultaneously changing base to
362 Chebychev-U, you may need a basis change matrix to transfer the other matrices as well. This function returns a
363 conversion matrix from `ChebychevHelper.get_conv`. Not that `**kwargs` are used to absorb arguments for other
364 bases, see documentation of `SpectralHelper1D.get_basis_change_matrix`.
366 Args:
367 conv (str): Conversion code, i.e. T2U
369 Returns:
370 Sparse conversion matrix
371 """
372 return self.get_conv(conv)
374 def get_integration_matrix(self, lbnd=0):
375 """
376 Get matrix for integration
378 Args:
379 lbnd (float): Lower bound for integration, only 0 is currently implemented
381 Returns:
382 Sparse integration matrix
383 """
384 S = self.sparse_lib.diags(1 / (self.xp.arange(self.N - 1) + 1), offsets=-1) @ self.get_conv('T2U')
385 n = self.xp.arange(self.N)
386 if lbnd == 0:
387 S = S.tocsc()
388 S[0, 1::2] = (
389 (n / (2 * (self.xp.arange(self.N) + 1)))[1::2]
390 * (-1) ** (self.xp.arange(self.N // 2))
391 / (np.append([1], self.xp.arange(self.N // 2 - 1) + 1))
392 ) * self.lin_trf_fac
393 else:
394 raise NotImplementedError(f'This function allows to integrate only from x=0, you attempted from x={lbnd}.')
395 return S
397 def get_integration_weights(self):
398 """Weights for integration across entire domain"""
399 n = self.xp.arange(self.N, dtype=float)
401 weights = (-1) ** n + 1
402 weights[2:] /= 1 - (n**2)[2:]
404 weights /= 2 / self.L
405 return weights
407 def get_differentiation_matrix(self, p=1):
408 '''
409 Keep in mind that the T2T differentiation matrix is dense.
411 Args:
412 p (int): Derivative you want to compute
414 Returns:
415 numpy.ndarray: Differentiation matrix
416 '''
417 D = self.xp.zeros((self.N, self.N))
418 for j in range(self.N):
419 for k in range(j):
420 D[k, j] = 2 * j * ((j - k) % 2)
422 D[0, :] /= 2
423 return self.sparse_lib.csc_matrix(self.xp.linalg.matrix_power(D, p)) / self.lin_trf_fac**p
425 @cache
426 def get_norm(self, N=None):
427 '''
428 Get normalization for converting Chebychev coefficients and DCT
430 Args:
431 N (int, optional): Resolution
433 Returns:
434 self.xp.array: Normalization
435 '''
436 N = self.N if N is None else N
437 norm = self.xp.ones(N) / N
438 norm[0] /= 2
439 return norm
441 def transform(self, u, *args, axes=None, shape=None, **kwargs):
442 """
443 DCT along axes. `kwargs` will be passed on to the FFT library.
445 Args:
446 u: Data you want to transform
447 axes (tuple): Axes you want to transform along
448 shape (tuple): Shape of the result along `axes`, passed on as `s` to the FFT library
450 Returns:
451 Data in spectral space
452 """
453 axes = axes if axes else tuple(i for i in range(u.ndim))
454 kwargs['s'] = shape
455 kwargs['norm'] = kwargs.get('norm', 'backward')
457 trf = self.fft_lib.dctn(u, *args, axes=axes, type=2, **kwargs)
458 for axis in axes:
460 if self.N < trf.shape[axis]:
461 # mpi4py-fft implements padding only for FFT, where the frequencies are sorted such that the zeros are
462 # removed in the middle rather than the end. We need to resort this here and put the highest frequencies
463 # in the middle.
464 _trf = self.xp.zeros_like(trf)
465 N = self.N
466 N_pad = _trf.shape[axis] - N
467 end_first_half = N // 2 + 1
469 # copy first "half"
470 su = [slice(None)] * trf.ndim
471 su[axis] = slice(0, end_first_half)
472 _trf[tuple(su)] = trf[tuple(su)]
474 # copy second "half"
475 su = [slice(None)] * u.ndim
476 su[axis] = slice(end_first_half + N_pad, None)
477 s_u = [slice(None)] * u.ndim
478 s_u[axis] = slice(end_first_half, N)
479 _trf[tuple(su)] = trf[tuple(s_u)]
481 # # copy values to be cut
482 # su = [slice(None)] * u.ndim
483 # su[axis] = slice(end_first_half, end_first_half + N_pad)
484 # s_u = [slice(None)] * u.ndim
485 # s_u[axis] = slice(-N_pad, None)
486 # _trf[tuple(su)] = trf[tuple(s_u)]
488 trf = _trf
490 expansion = [np.newaxis for _ in u.shape]
491 expansion[axis] = slice(0, u.shape[axis], 1)
492 norm = self.xp.ones(trf.shape[axis]) * self.norm[-1]
493 norm[: self.N] = self.norm
494 trf *= norm[(*expansion,)]
495 return trf
497 def itransform(self, u, *args, axes=None, shape=None, **kwargs):
498 """
499 Inverse DCT along axis.
501 Args:
502 u: Data you want to transform
503 axes (tuple): Axes you want to transform along
504 shape (tuple): Shape of the result along `axes`, passed on as `s` to the FFT library
506 Returns:
507 Data in physical space
508 """
509 axes = axes if axes else tuple(i for i in range(u.ndim))
510 kwargs['s'] = shape
511 kwargs['norm'] = kwargs.get('norm', 'backward')
512 kwargs['overwrite_x'] = kwargs.get('overwrite_x', False)
514 for axis in axes:
516 if self.N == u.shape[axis]:
517 _u = u.copy()
518 else:
519 # mpi4py-fft implements padding only for FFT, where the frequencies are sorted such that the zeros are
520 # added in the middle rather than the end. We need to resort this here and put the padding in the end.
521 N = self.N
522 _u = self.xp.zeros_like(u)
524 # copy first half
525 su = [slice(None)] * u.ndim
526 su[axis] = slice(0, N // 2 + 1)
527 _u[tuple(su)] = u[tuple(su)]
529 # copy second half
530 su = [slice(None)] * u.ndim
531 su[axis] = slice(-(N // 2), None)
532 s_u = [slice(None)] * u.ndim
533 s_u[axis] = slice(N // 2, N // 2 + (N // 2))
534 _u[tuple(s_u)] = u[tuple(su)]
536 if N % 2 == 0:
537 su = [slice(None)] * u.ndim
538 su[axis] = N // 2
539 _u[tuple(su)] *= 2
541 # generate norm
542 expansion = [np.newaxis for _ in u.shape]
543 expansion[axis] = slice(0, u.shape[axis], 1)
544 norm = self.xp.ones(_u.shape[axis])
545 norm[: self.N] = self.norm
546 norm = self.get_norm(u.shape[axis]) * _u.shape[axis] / self.N
548 _u /= norm[(*expansion,)]
550 return self.fft_lib.idctn(_u, *args, axes=axes, type=2, **kwargs)
552 def get_BC(self, kind, **kwargs):
553 """
554 Get boundary condition row for boundary bordering. `kwargs` will be passed on to implementations of the BC of
555 the kind you choose. Specifically, `x` for `'dirichlet'` boundary condition, which is the coordinate at which to
556 set the BC.
558 Args:
559 kind ('integral' or 'dirichlet'): Kind of boundary condition you want
560 """
561 if kind.lower() == 'integral':
562 return self.get_integ_BC_row(**kwargs)
563 elif kind.lower() == 'dirichlet':
564 return self.get_Dirichlet_BC_row(**kwargs)
565 elif kind.lower() == 'neumann':
566 return self.get_Neumann_BC_row(**kwargs)
567 else:
568 return super().get_BC(kind)
570 def get_integ_BC_row(self):
571 """
572 Get a row for generating integral BCs with T polynomials.
573 It returns the values of the integrals of T polynomials over the entire interval.
575 Returns:
576 self.xp.ndarray: Row to put into a matrix
577 """
578 n = self.xp.arange(self.N) + 1
579 me = self.xp.zeros_like(n).astype(float)
580 me[2:] = ((-1) ** n[1:-1] + 1) / (1 - n[1:-1] ** 2)
581 me[0] = 2.0
582 return me
584 def get_Dirichlet_BC_row(self, x):
585 """
586 Get a row for generating Dirichlet BCs at x with T polynomials.
587 It returns the values of the T polynomials at x.
589 Args:
590 x (float): Position of the boundary condition
592 Returns:
593 self.xp.ndarray: Row to put into a matrix
594 """
595 if x == -1:
596 return (-1) ** self.xp.arange(self.N)
597 elif x == 1:
598 return self.xp.ones(self.N)
599 elif x == 0:
600 n = (1 + (-1) ** self.xp.arange(self.N)) / 2
601 n[2::4] *= -1
602 return n
603 else:
604 raise NotImplementedError(f'Don\'t know how to generate Dirichlet BC\'s at {x=}!')
606 def get_Neumann_BC_row(self, x):
607 """
608 Get a row for generating Neumann BCs at x with T polynomials.
610 Args:
611 x (float): Position of the boundary condition
613 Returns:
614 self.xp.ndarray: Row to put into a matrix
615 """
616 n = self.xp.arange(self.N, dtype='D')
617 nn = n**2
618 if x == -1:
619 me = nn
620 me[1:] *= (-1) ** n[:-1]
621 return me
622 elif x == 1:
623 return nn
624 else:
625 raise NotImplementedError(f'Don\'t know how to generate Neumann BC\'s at {x=}!')
627 def get_Dirichlet_recombination_matrix(self):
628 '''
629 Get matrix for Dirichlet recombination, which changes the basis to have sparse boundary conditions.
630 This makes for a good right preconditioner.
632 Returns:
633 scipy.sparse: Sparse conversion matrix
634 '''
635 N = self.N
636 sp = self.sparse_lib
638 return sp.eye(N) - sp.eye(N, k=2)
641class UltrasphericalHelper(ChebychevHelper):
642 """
643 This implementation follows https://doi.org/10.1137/120865458.
644 The ultraspherical method works in Chebychev polynomials as well, but also uses various Gegenbauer polynomials.
645 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.
646 There are also conversion operators from one derivative basis to the next that are sparse.
648 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.
649 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.
650 """
652 def get_differentiation_matrix(self, p=1):
653 """
654 Notice that while sparse, this matrix is not diagonal, which means the inversion cannot be parallelized easily.
656 Args:
657 p (int): Order of the derivative
659 Returns:
660 sparse differentiation matrix
661 """
662 sp = self.sparse_lib
663 xp = self.xp
664 N = self.N
665 l = p
666 return 2 ** (l - 1) * factorial(l - 1) * sp.diags(xp.arange(N - l) + l, offsets=l) / self.lin_trf_fac**p
668 def get_S(self, lmbda):
669 """
670 Get matrix for bumping the derivative base by one from lmbda to lmbda + 1. This is the same language as in
671 https://doi.org/10.1137/120865458.
673 Args:
674 lmbda (int): Ingoing derivative base
676 Returns:
677 sparse matrix: Conversion from derivative base lmbda to lmbda + 1
678 """
679 N = self.N
681 if lmbda == 0:
682 sp = scipy.sparse
683 mat = ((sp.eye(N) - sp.eye(N, k=2)) / 2.0).tolil()
684 mat[:, 0] *= 2
685 else:
686 sp = self.sparse_lib
687 xp = self.xp
688 mat = sp.diags(lmbda / (lmbda + xp.arange(N))) - sp.diags(
689 lmbda / (lmbda + 2 + xp.arange(N - 2)), offsets=+2
690 )
692 return self.sparse_lib.csc_matrix(mat)
694 def get_basis_change_matrix(self, p_in=0, p_out=0, **kwargs):
695 """
696 Get a conversion matrix from derivative base `p_in` to `p_out`.
698 Args:
699 p_out (int): Resulting derivative base
700 p_in (int): Ingoing derivative base
701 """
702 mat_fwd = self.sparse_lib.eye(self.N)
703 for i in range(min([p_in, p_out]), max([p_in, p_out])):
704 mat_fwd = self.get_S(i) @ mat_fwd
706 if p_out > p_in:
707 return mat_fwd
709 else:
710 # We have to invert the matrix on CPU because the GPU equivalent is not implemented in CuPy at the time of writing.
711 import scipy.sparse as sp
713 if self.useGPU:
714 mat_fwd = mat_fwd.get()
716 mat_bck = sp.linalg.inv(mat_fwd.tocsc())
718 return self.sparse_lib.csc_matrix(mat_bck)
720 def get_integration_matrix(self):
721 """
722 Get an integration matrix. Please use `UltrasphericalHelper.get_integration_constant` afterwards to compute the
723 integration constant such that integration starts from x=-1.
725 Example:
727 .. code-block:: python
729 import numpy as np
730 from pySDC.helpers.spectral_helper import UltrasphericalHelper
732 N = 4
733 helper = UltrasphericalHelper(N)
734 coeffs = np.random.random(N)
735 coeffs[-1] = 0
737 poly = np.polynomial.Chebyshev(coeffs)
739 S = helper.get_integration_matrix()
740 U_hat = S @ coeffs
741 U_hat[0] = helper.get_integration_constant(U_hat, axis=-1)
743 assert np.allclose(poly.integ(lbnd=-1).coef[:-1], U_hat)
745 Returns:
746 sparse integration matrix
747 """
748 return (
749 self.sparse_lib.diags(1 / (self.xp.arange(self.N - 1) + 1), offsets=-1)
750 @ self.get_basis_change_matrix(p_out=1, p_in=0)
751 * self.lin_trf_fac
752 )
754 def get_integration_constant(self, u_hat, axis):
755 """
756 Get integration constant for lower bound of -1. See documentation of `UltrasphericalHelper.get_integration_matrix` for details.
758 Args:
759 u_hat: Solution in spectral space
760 axis: Axis you want to integrate over
762 Returns:
763 Integration constant, has one less dimension than `u_hat`
764 """
765 slices = [
766 None,
767 ] * u_hat.ndim
768 slices[axis] = slice(1, u_hat.shape[axis])
769 return self.xp.sum(u_hat[(*slices,)] * (-1) ** (self.xp.arange(u_hat.shape[axis] - 1)), axis=axis)
772class FFTHelper(SpectralHelper1D):
773 distributable = True
775 def __init__(self, *args, x0=0, x1=2 * np.pi, **kwargs):
776 """
777 Constructor.
778 Please refer to the parent class for additional arguments. Notably, you have to supply a resolution `N` and you
779 may choose to run on GPUs via the `useGPU` argument.
781 Args:
782 x0 (float, optional): Coordinate of left boundary
783 x1 (float, optional): Coordinate of right boundary
784 """
785 super().__init__(*args, x0=x0, x1=x1, **kwargs)
787 def get_1dgrid(self):
788 """
789 We use equally spaced points including the left boundary and not including the right one, which is the left boundary.
790 """
791 dx = self.L / self.N
792 return self.xp.arange(self.N) * dx + self.x0
794 def get_wavenumbers(self):
795 """
796 Be careful that this ordering is very unintuitive.
797 """
798 return self.xp.fft.fftfreq(self.N, 1.0 / self.N) * 2 * np.pi / self.L
800 def get_differentiation_matrix(self, p=1):
801 """
802 This matrix is diagonal, allowing to invert concurrently.
804 Args:
805 p (int): Order of the derivative
807 Returns:
808 sparse differentiation matrix
809 """
810 k = self.get_wavenumbers()
812 if self.useGPU:
813 if p > 1:
814 # Have to raise the matrix to power p on CPU because the GPU equivalent is not implemented in CuPy at the time of writing.
815 from scipy.sparse.linalg import matrix_power
817 D = self.sparse_lib.diags(1j * k).get()
818 return self.sparse_lib.csc_matrix(matrix_power(D, p))
819 else:
820 return self.sparse_lib.diags(1j * k)
821 else:
822 return self.linalg.matrix_power(self.sparse_lib.diags(1j * k), p)
824 def get_integration_matrix(self, p=1):
825 """
826 Get integration matrix to compute `p`-th integral over the entire domain.
828 Args:
829 p (int): Order of integral you want to compute
831 Returns:
832 sparse integration matrix
833 """
834 k = self.xp.array(self.get_wavenumbers(), dtype='complex128')
835 k[0] = 1j * self.L
836 return self.linalg.matrix_power(self.sparse_lib.diags(1 / (1j * k)), p)
838 def get_integration_weights(self):
839 """Weights for integration across entire domain"""
840 weights = self.xp.zeros(self.N)
841 weights[0] = self.L / self.N
842 return weights
844 def get_plan(self, u, forward, *args, **kwargs):
845 """
846 Get a callable that computes the FFT of data like `u`.
848 With FFTW, the planned transform is cached in `self.plans` by direction, shape of `u` and the other arguments.
849 Otherwise, this is `fftn` or `ifftn` of the FFT library with the `norm` from `kwargs`, "backward" and "forward"
850 by default, so that neither direction normalizes.
852 Args:
853 u: Data to transform
854 forward (bool): Forward transform if True, inverse transform otherwise
856 Returns:
857 callable transform
858 """
859 if self.fft_lib.__name__ == 'mpi4py_fft.fftw':
860 if 'axes' in kwargs.keys():
861 kwargs['axes'] = tuple(kwargs['axes'])
862 key = (forward, u.shape, u.dtype, args, *(me for me in kwargs.values()))
863 if key in self.plans.keys():
864 return self.plans[key]
865 else:
866 self.logger.debug(f'Generating FFT plan for {key=}')
867 transform = self.fft_lib.fftn(u, *args, **kwargs) if forward else self.fft_lib.ifftn(u, *args, **kwargs)
868 self.plans[key] = transform
870 return self.plans[key]
871 else:
872 if forward:
873 return partial(self.fft_lib.fftn, norm=kwargs.get('norm', 'backward'))
874 else:
875 return partial(self.fft_lib.ifftn, norm=kwargs.get('norm', 'forward'))
877 def transform(self, u, *args, axes=None, shape=None, **kwargs):
878 """
879 FFT along axes. `kwargs` are passed on to the FFT library.
881 Args:
882 u: Data you want to transform
883 axes (tuple): Axes you want to transform over
884 shape (tuple): Shape of the result along `axes`, passed on as `s` to the FFT library
886 Returns:
887 transformed data
888 """
889 axes = axes if axes else tuple(i for i in range(u.ndim))
890 kwargs['s'] = shape
891 plan = self.get_plan(u, *args, forward=True, axes=axes, **kwargs)
892 return plan(u, *args, axes=axes, **kwargs)
894 def itransform(self, u, *args, axes=None, shape=None, **kwargs):
895 """
896 Inverse FFT.
898 Args:
899 u: Data you want to transform
900 axes (tuple): Axes over which to transform
901 shape (tuple): Shape of the result along `axes`, passed on as `s` to the FFT library
903 Returns:
904 transformed data
905 """
906 axes = axes if axes else tuple(i for i in range(u.ndim))
907 kwargs['s'] = shape
908 plan = self.get_plan(u, *args, forward=False, axes=axes, **kwargs)
909 return plan(u, *args, axes=axes, **kwargs) / np.prod([u.shape[axis] for axis in axes])
911 def get_BC(self, kind):
912 """
913 Get a sort of boundary condition. You can use `kind=integral`, to fix the integral, or you can use `kind=Nyquist`.
914 The latter is not really a boundary condition, but is used to set the Nyquist mode to some value, preferably zero.
915 You should set the Nyquist mode zero when the solution in physical space is real and the resolution is even.
917 Args:
918 kind ('integral' or 'nyquist'): Kind of BC
920 Returns:
921 self.xp.ndarray: Boundary condition row
922 """
923 if kind.lower() == 'integral':
924 return self.get_integ_BC_row()
925 elif kind.lower() == 'nyquist':
926 assert (
927 self.N % 2 == 0
928 ), f'Do not eliminate the Nyquist mode with odd resolution as it is fully resolved. You chose {self.N} in this axis'
929 BC = self.xp.zeros(self.N)
930 BC[self.get_Nyquist_mode_index()] = 1
931 return BC
932 else:
933 return super().get_BC(kind)
935 def get_Nyquist_mode_index(self):
936 """
937 Compute the index of the Nyquist mode, i.e. the mode with the lowest wavenumber, which doesn't have a positive
938 counterpart for even resolution. This means real waves of this wave number cannot be properly resolved and you
939 are best advised to set this mode zero if representing real functions on even-resolution grids is what you're
940 after.
942 Returns:
943 int: Index of the Nyquist mode
944 """
945 k = self.get_wavenumbers()
946 Nyquist_mode = min(k)
947 return self.xp.where(k == Nyquist_mode)[0][0]
949 def get_integ_BC_row(self):
950 """
951 Only the 0-mode has non-zero integral with FFT basis in periodic BCs
952 """
953 me = self.xp.zeros(self.N)
954 me[0] = self.L / self.N
955 return me
958class SpectralHelper:
959 """
960 This class has three functions:
961 - Easily assemble matrices containing multiple equations
962 - Direct product of 1D bases to solve problems in more dimensions
963 - Distribute the FFTs to facilitate concurrency.
965 Attributes:
966 comm (mpi4py.Intracomm): MPI communicator
967 debug (bool): Perform additional checks at extra computational cost
968 useGPU (bool): Whether to use GPUs
969 axes (list): List of 1D bases
970 components (list): List of strings of the names of components in the equations
971 full_BCs (list): List of Dictionaries containing all information about the boundary conditions
972 BC_mat (list): List of lists of sparse matrices to put BCs into and eventually assemble the BC matrix from
973 BCs (sparse matrix): Matrix containing only the BCs
974 fft_cache (dict): Cache FFTs of various shapes here to facilitate padding and so on
975 BC_rhs_mask (self.xp.ndarray): Mask values that contain boundary conditions in the right hand side
976 BC_zero_index (self.xp.ndarray): Indeces of rows in the matrix that are replaced by BCs
977 BC_line_zero_matrix (sparse matrix): Matrix that zeros rows where we can then add the BCs in using `BCs`
978 rhs_BCs_hat (self.xp.ndarray): Boundary conditions in spectral space
979 global_shape (tuple): Global shape of the solution as in `mpi4py-fft`
980 fft_obj: When using distributed FFTs, this will be a parallel transform object from `mpi4py-fft`
981 init (tuple): This is the same `init` that is used throughout the problem classes
982 init_forward (tuple): This is the equivalent of `init` in spectral space
983 """
985 xp = np
986 fft_lib = scipy.fft
987 sparse_lib = scipy.sparse
988 linalg = scipy.sparse.linalg
989 dtype = mesh
990 fft_backend = 'scipy'
991 fft_comm_backend = 'MPI'
993 def setup_GPU(self):
994 """switch to GPU modules"""
995 import cupy as cp
996 import cupyx.scipy.sparse as sparse_lib
997 import cupyx.scipy.sparse.linalg as linalg
998 import cupyx.scipy.fft as fft_lib
999 from pySDC.implementations.datatype_classes.cupy_mesh import cupy_mesh
1001 self.xp = cp
1002 self.sparse_lib = sparse_lib
1003 self.linalg = linalg
1005 self.fft_lib = fft_lib
1006 self.fft_backend = 'cupyx-scipy'
1007 self.fft_comm_backend = 'NCCL'
1009 self.dtype = cupy_mesh
1011 def setup_CPU(self, useFFTW=False):
1012 """switch to CPU modules"""
1014 self.xp = np
1015 self.sparse_lib = scipy.sparse
1016 self.linalg = scipy.sparse.linalg
1018 if useFFTW:
1019 from mpi4py_fft import fftw
1021 self.fft_backend = 'fftw'
1022 self.fft_lib = fftw
1023 else:
1024 self.fft_backend = 'scipy'
1025 self.fft_lib = scipy.fft
1027 self.fft_comm_backend = 'MPI'
1028 self.dtype = mesh
1030 def __init__(self, comm=None, useGPU=False, debug=False):
1031 """
1032 Constructor
1034 Args:
1035 comm (mpi4py.Intracomm): MPI communicator
1036 useGPU (bool): Whether to use GPUs
1037 debug (bool): Perform additional checks at extra computational cost
1038 """
1039 self.comm = comm
1040 self.debug = debug
1041 self.useGPU = useGPU
1043 if useGPU:
1044 self.setup_GPU()
1045 else:
1046 self.setup_CPU()
1048 self.axes = []
1049 self.components = []
1051 self.full_BCs = []
1052 self.BC_mat = None
1053 self.BCs = None
1055 self.fft_cache = {}
1057 self.logger = logging.getLogger(name='Spectral Discretization')
1058 if debug:
1059 self.logger.setLevel(logging.DEBUG)
1061 @property
1062 def u_init(self):
1063 """
1064 Get empty data container in physical space
1065 """
1066 return self.dtype(self.init)
1068 @property
1069 def u_init_forward(self):
1070 """
1071 Get empty data container in spectral space
1072 """
1073 return self.dtype(self.init_forward)
1075 @property
1076 def u_init_physical(self):
1077 """
1078 Get empty data container in physical space
1079 """
1080 return self.dtype(self.init_physical)
1082 @property
1083 def shape(self):
1084 """
1085 Get shape of individual solution component
1086 """
1087 return self.init[0][1:]
1089 @property
1090 def ndim(self):
1091 """Number of dimensions, i.e. of 1D bases"""
1092 return len(self.axes)
1094 @property
1095 def ncomponents(self):
1096 """Number of solution components"""
1097 return len(self.components)
1099 @property
1100 def V(self):
1101 """
1102 Get domain volume
1103 """
1104 return np.prod([me.L for me in self.axes])
1106 def add_axis(self, base, *args, **kwargs):
1107 """
1108 Add an axis to the domain by deciding on suitable 1D base.
1109 Arguments to the bases are forwarded using `*args` and `**kwargs`. Please refer to the documentation of the 1D
1110 bases for possible arguments.
1112 Args:
1113 base (str): 1D spectral method
1114 """
1115 kwargs['useGPU'] = self.useGPU
1117 if base.lower() in ['chebychov', 'chebychev', 'cheby', 'chebychovhelper']:
1118 self.axes.append(ChebychevHelper(*args, **kwargs))
1119 elif base.lower() in ['fft', 'fourier', 'ffthelper']:
1120 self.axes.append(FFTHelper(*args, **kwargs))
1121 elif base.lower() in ['ultraspherical', 'gegenbauer']:
1122 self.axes.append(UltrasphericalHelper(*args, **kwargs))
1123 else:
1124 raise NotImplementedError(f'{base=!r} is not implemented!')
1125 self.axes[-1].xp = self.xp
1126 self.axes[-1].sparse_lib = self.sparse_lib
1128 def add_component(self, name):
1129 """
1130 Add solution component(s).
1132 Args:
1133 name (str or list of strings): Name(s) of component(s)
1134 """
1135 if type(name) in [list, tuple]:
1136 for me in name:
1137 self.add_component(me)
1138 elif type(name) in [str]:
1139 if name in self.components:
1140 raise Exception(f'{name=!r} is already added to this problem!')
1141 self.components.append(name)
1142 else:
1143 raise NotImplementedError
1145 def index(self, name):
1146 """
1147 Get the index of component `name`.
1149 Args:
1150 name (str or list of strings): Name(s) of component(s)
1152 Returns:
1153 int: Index of the component
1154 """
1155 if type(name) in [str, int]:
1156 return self.components.index(name)
1157 elif type(name) in [list, tuple]:
1158 return (self.index(me) for me in name)
1159 else:
1160 raise NotImplementedError(f'Don\'t know how to compute index for {type(name)=}')
1162 def get_empty_operator_matrix(self, diag=False):
1163 """
1164 Return a matrix of operators to be filled with the connections between the solution components.
1166 Args:
1167 diag (bool): Whether operator is block-diagonal
1169 Returns:
1170 list containing sparse zeros
1171 """
1172 S = len(self.components)
1173 O = self.get_Id() * 0
1174 if diag:
1175 return [O for _ in range(S)]
1176 else:
1177 return [[O for _ in range(S)] for _ in range(S)]
1179 def get_BC(self, axis, kind, line=-1, scalar=False, **kwargs):
1180 """
1181 Use this method for boundary bordering. It gets the respective matrix row and embeds it into a matrix.
1182 Pay attention that if you have multiple BCs in a single equation, you need to put them in different lines.
1183 Typically, the last line that does not contain a BC is the best choice.
1184 Forward arguments for the boundary conditions using `kwargs`. Refer to documentation of 1D bases for details.
1186 Args:
1187 axis (int): Axis you want to add the BC to
1188 kind (str): kind of BC, e.g. Dirichlet
1189 line (int): Line you want the BC to go in
1190 scalar (bool): Put the BC in all space positions in the other direction
1192 Returns:
1193 sparse matrix containing the BC
1194 """
1195 sp = scipy.sparse
1197 base = self.axes[axis]
1199 BC = sp.eye(base.N).tolil() * 0
1200 if self.useGPU:
1201 BC[line, :] = base.get_BC(kind=kind, **kwargs).get()
1202 else:
1203 BC[line, :] = base.get_BC(kind=kind, **kwargs)
1205 ndim = len(self.axes)
1206 if ndim == 1:
1207 mat = self.sparse_lib.csc_matrix(BC)
1208 elif ndim == 2:
1209 axis2 = (axis + 1) % ndim
1211 if scalar:
1212 _Id = self.sparse_lib.diags(self.xp.append([1], self.xp.zeros(self.axes[axis2].N - 1)))
1213 else:
1214 _Id = self.axes[axis2].get_Id()
1216 Id = self.get_local_slice_of_1D_matrix(self.axes[axis2].get_Id() @ _Id, axis=axis2)
1218 mats = [
1219 None,
1220 ] * ndim
1221 mats[axis] = self.get_local_slice_of_1D_matrix(BC, axis=axis)
1222 mats[axis2] = Id
1223 mat = self.sparse_lib.csc_matrix(self.sparse_lib.kron(*mats))
1224 elif ndim == 3:
1225 mats = [
1226 None,
1227 ] * ndim
1229 for ax in range(ndim):
1230 if ax == axis:
1231 continue
1233 if scalar:
1234 _Id = self.sparse_lib.diags(self.xp.append([1], self.xp.zeros(self.axes[ax].N - 1)))
1235 else:
1236 _Id = self.axes[ax].get_Id()
1238 mats[ax] = self.get_local_slice_of_1D_matrix(self.axes[ax].get_Id() @ _Id, axis=ax)
1240 mats[axis] = self.get_local_slice_of_1D_matrix(BC, axis=axis)
1242 mat = self.sparse_lib.csc_matrix(self.sparse_lib.kron(mats[0], self.sparse_lib.kron(*mats[1:])))
1243 else:
1244 raise NotImplementedError(
1245 f'Matrix expansion for boundary conditions not implemented for {ndim} dimensions!'
1246 )
1247 mat = self.eliminate_zeros(mat)
1248 return mat
1250 def remove_BC(self, component, equation, axis, kind, line=-1, scalar=False, **kwargs):
1251 """
1252 Remove a BC from the matrix. This is useful e.g. when you add a non-scalar BC and then need to selectively
1253 remove single BCs again, as in incompressible Navier-Stokes, for instance.
1254 Forwards arguments for the boundary conditions using `kwargs`. Refer to documentation of 1D bases for details.
1256 Args:
1257 component (str): Name of the component the BC should act on
1258 equation (str): Name of the equation for the component you want to put the BC in
1259 axis (int): Axis you want to add the BC to
1260 kind (str): kind of BC, e.g. Dirichlet
1261 line (int): Line you want the BC to go in
1262 scalar (bool): Put the BC in all space positions in the other direction
1263 """
1264 _BC = self.get_BC(axis=axis, kind=kind, line=line, scalar=scalar, **kwargs)
1265 _BC = self.eliminate_zeros(_BC)
1266 self.BC_mat[self.index(equation)][self.index(component)] -= _BC
1268 if scalar:
1269 slices = [self.index(equation)] + [
1270 0,
1271 ] * self.ndim
1272 slices[axis + 1] = line
1273 else:
1274 slices = (
1275 [self.index(equation)]
1276 + [slice(0, self.init[0][i + 1]) for i in range(axis)]
1277 + [line]
1278 + [slice(0, self.init[0][i + 1]) for i in range(axis + 1, len(self.axes))]
1279 )
1280 N = self.axes[axis].N
1281 if (N + line) % N in self.xp.arange(N)[self.local_slice()[axis]]:
1282 self.BC_rhs_mask[(*slices,)] = False
1284 def add_BC(self, component, equation, axis, kind, v, line=-1, scalar=False, **kwargs):
1285 """
1286 Add a BC to the matrix. Note that you need to convert the list of lists of BCs that this method generates to a
1287 single sparse matrix by calling `setup_BCs` after adding/removing all BCs.
1288 Forward arguments for the boundary conditions using `kwargs`. Refer to documentation of 1D bases for details.
1290 Args:
1291 component (str): Name of the component the BC should act on
1292 equation (str): Name of the equation for the component you want to put the BC in
1293 axis (int): Axis you want to add the BC to
1294 kind (str): kind of BC, e.g. Dirichlet
1295 v: Value of the BC
1296 line (int): Line you want the BC to go in
1297 scalar (bool): Put the BC in all space positions in the other direction
1298 """
1299 _BC = self.get_BC(axis=axis, kind=kind, line=line, scalar=scalar, **kwargs)
1300 self.BC_mat[self.index(equation)][self.index(component)] += _BC
1301 self.full_BCs += [
1302 {
1303 'component': component,
1304 'equation': equation,
1305 'axis': axis,
1306 'kind': kind,
1307 'v': v,
1308 'line': line,
1309 'scalar': scalar,
1310 **kwargs,
1311 }
1312 ]
1314 if scalar:
1315 slices = [self.index(equation)] + [
1316 0,
1317 ] * self.ndim
1318 slices[axis + 1] = line
1319 if self.comm:
1320 if self.comm.rank == 0:
1321 self.BC_rhs_mask[(*slices,)] = True
1322 else:
1323 self.BC_rhs_mask[(*slices,)] = True
1324 else:
1325 slices = [self.index(equation), *self.global_slice(True)]
1326 N = self.axes[axis].N
1327 if (N + line) % N in self.get_indices(True)[axis]:
1328 slices[axis + 1] = (N + line) % N - self.local_slice()[axis].start
1329 self.BC_rhs_mask[(*slices,)] = True
1331 def setup_BCs(self):
1332 """
1333 Convert the list of lists of BCs to the boundary condition operator.
1334 Also, boundary bordering requires to zero out all other entries in the matrix in rows containing a boundary
1335 condition. This method sets up a suitable sparse matrix to do this.
1336 """
1337 sp = self.sparse_lib
1338 self.BCs = self.convert_operator_matrix_to_operator(self.BC_mat)
1339 self.BC_zero_index = self.xp.arange(np.prod(self.init[0]))[self.BC_rhs_mask.flatten()]
1341 diags = self.xp.ones(self.BCs.shape[0])
1342 diags[self.BC_zero_index] = 0
1343 self.BC_line_zero_matrix = sp.diags(diags).tocsc()
1345 # prepare BCs in spectral space to easily add to the RHS
1346 rhs_BCs = self.put_BCs_in_rhs(self.u_init)
1347 self.rhs_BCs_hat = self.transform(rhs_BCs).view(self.xp.ndarray)
1348 del self.BC_rhs_mask
1350 def check_BCs(self, u):
1351 """
1352 Check that the solution satisfies the boundary conditions
1354 Args:
1355 u: The solution you want to check
1356 """
1357 assert self.ndim < 3
1358 for axis in range(self.ndim):
1359 BCs = [me for me in self.full_BCs if me["axis"] == axis and not me["scalar"]]
1361 if len(BCs) > 0:
1362 u_hat = self.transform(u, axes=(axis - self.ndim,))
1363 for BC in BCs:
1364 kwargs = {
1365 key: value
1366 for key, value in BC.items()
1367 if key not in ['component', 'equation', 'axis', 'v', 'line', 'scalar']
1368 }
1370 if axis == 0:
1371 get = self.axes[axis].get_BC(**kwargs) @ u_hat[self.index(BC['component'])]
1372 elif axis == 1:
1373 get = u_hat[self.index(BC['component'])] @ self.axes[axis].get_BC(**kwargs)
1374 want = BC['v']
1375 assert self.xp.allclose(
1376 get, want
1377 ), f'Unexpected BC in {BC["component"]} in equation {BC["equation"]}, line {BC["line"]}! Got {get}, wanted {want}'
1379 def put_BCs_in_matrix(self, A):
1380 """
1381 Put the boundary conditions in a matrix by replacing rows with BCs.
1382 """
1383 return self.BC_line_zero_matrix @ A + self.BCs
1385 def put_BCs_in_rhs_hat(self, rhs_hat, BCs_hat=None):
1386 """
1387 Put the BCs in the right hand side in spectral space for solving.
1388 This function needs no transforms and caches a mask for faster subsequent use.
1390 Args:
1391 rhs_hat: Right hand side in spectral space
1392 BCs_hat: Boundary values to use instead of `rhs_BCs_hat`, zero outside the boundary rows
1394 Returns:
1395 rhs in spectral space with BCs
1396 """
1397 if not hasattr(self, '_rhs_hat_zero_mask'):
1398 """
1399 Generate a mask where we need to set values in the rhs in spectral space to zero, such that can replace them
1400 by the boundary conditions. The mask is then cached.
1401 """
1402 self._rhs_hat_zero_mask = self.newDistArray(forward_output=True).astype(bool).view(self.xp.ndarray)
1404 for axis in range(self.ndim):
1405 for bc in self.full_BCs:
1406 if axis == bc['axis']:
1407 slices = [self.index(bc['equation']), *self.global_slice(True)]
1408 N = self.axes[axis].N
1409 line = bc['line']
1410 if (N + line) % N in self.get_indices(True)[axis]:
1411 slices[axis + 1] = (N + line) % N - self.local_slice()[axis].start
1412 self._rhs_hat_zero_mask[(*slices,)] = True
1414 rhs_hat[self._rhs_hat_zero_mask] = 0
1415 return rhs_hat + (self.rhs_BCs_hat if BCs_hat is None else BCs_hat)
1417 def put_BCs_in_rhs(self, rhs):
1418 """
1419 Put the BCs in the right hand side for solving.
1420 This function will transform along each axis individually and add all BCs in that axis.
1421 Consider `put_BCs_in_rhs_hat` to add BCs with no extra transforms needed.
1423 Args:
1424 rhs: Right hand side in physical space
1426 Returns:
1427 rhs in physical space with BCs
1428 """
1429 assert rhs.ndim > 1, 'rhs must not be flattened here!'
1431 ndim = self.ndim
1433 for axis in range(ndim):
1434 _rhs_hat = self.transform(rhs, axes=(axis - ndim,))
1436 for bc in self.full_BCs:
1438 if axis == bc['axis']:
1439 _slice = [self.index(bc['equation']), *self.global_slice(True)]
1441 N = self.axes[axis].N
1442 line = bc['line']
1443 if (N + line) % N in self.get_indices(True)[axis]:
1444 _slice[axis + 1] = (N + line) % N - self.local_slice()[axis].start
1445 _rhs_hat[(*_slice,)] = bc['v']
1447 rhs = self.itransform(_rhs_hat, axes=(axis - ndim,))
1449 return rhs
1451 def add_equation_lhs(self, A, equation, relations):
1452 """
1453 Add the left hand part (that you want to solve implicitly) of an equation to a list of lists of sparse matrices
1454 that you will convert to an operator later.
1456 Example:
1457 Setup linear operator `L` for 1D heat equation using Chebychev method in first order form and T-to-U
1458 preconditioning:
1460 .. code-block:: python
1462 helper = SpectralHelper()
1464 helper.add_axis(base='chebychev', N=8)
1465 helper.add_component(['u', 'ux'])
1466 helper.setup_fft()
1468 I = helper.get_Id()
1469 Dx = helper.get_differentiation_matrix(axes=(0,))
1470 T2U = helper.get_basis_change_matrix('T2U')
1472 L_lhs = {
1473 'ux': {'u': -T2U @ Dx, 'ux': T2U @ I},
1474 'u': {'ux': -(T2U @ Dx)},
1475 }
1477 operator = helper.get_empty_operator_matrix()
1478 for line, equation in L_lhs.items():
1479 helper.add_equation_lhs(operator, line, equation)
1481 L = helper.convert_operator_matrix_to_operator(operator)
1483 Args:
1484 A (list of lists of sparse matrices): The operator to be
1485 equation (str): The equation of the component you want this in
1486 relations: (dict): Relations between quantities
1487 """
1488 for k, v in relations.items():
1489 A[self.index(equation)][self.index(k)] = v
1491 def eliminate_zeros(self, A):
1492 """
1493 Eliminate zeros from sparse matrix. This can reduce memory footprint of matrices somewhat.
1494 Note: At the time of writing, there are memory problems in the cupy implementation of `eliminate_zeros`.
1495 Therefore, this function copies the matrix to host, eliminates the zeros there and then copies back to GPU.
1497 Args:
1498 A: sparse matrix to be pruned
1500 Returns:
1501 CSC sparse matrix
1502 """
1503 if self.useGPU:
1504 A = A.get()
1505 A = A.tocsc()
1506 A.eliminate_zeros()
1507 if self.useGPU:
1508 A = self.sparse_lib.csc_matrix(A)
1509 return A
1511 def convert_operator_matrix_to_operator(self, M):
1512 """
1513 Promote the list of lists of sparse matrices to a single sparse matrix that can be used as linear operator.
1514 See documentation of `SpectralHelper.add_equation_lhs` for an example.
1516 Args:
1517 M (list of lists of sparse matrices): The operator to be
1519 Returns:
1520 sparse linear operator
1521 """
1522 if len(self.components) == 1:
1523 op = M[0][0]
1524 else:
1525 op = self.sparse_lib.bmat(M, format='csc')
1527 op = self.eliminate_zeros(op)
1528 return op
1530 def get_wavenumbers(self):
1531 """
1532 Get grid in spectral space
1533 """
1534 grids = [self.axes[i].get_wavenumbers()[self.local_slice(True)[i]] for i in range(len(self.axes))]
1535 return self.xp.meshgrid(*grids, indexing='ij')
1537 def get_grid(self, forward_output=False):
1538 """
1539 Get grid in physical space
1540 """
1541 grids = [self.axes[i].get_1dgrid()[self.local_slice(forward_output)[i]] for i in range(len(self.axes))]
1542 return self.xp.meshgrid(*grids, indexing='ij')
1544 def get_indices(self, forward_output=True):
1545 """
1546 Get the global indices of the modes or grid points held by this rank in each axis.
1548 Args:
1549 forward_output (bool): Indices in spectral space if True, in physical space otherwise
1551 Returns:
1552 list of self.xp.ndarray: Indices in each axis
1553 """
1554 return [self.xp.arange(self.axes[i].N)[self.local_slice(forward_output)[i]] for i in range(len(self.axes))]
1556 @cache
1557 def get_pfft(self, axes=None, padding=None, grid=None):
1558 """
1559 Get a cached distributed transform object from mpi4py-fft, planned by one forward and backward transform.
1560 Axes that are not in `axes` are not transformed. The arguments must be hashable, e.g. tuples.
1562 Args:
1563 axes (tuple): Axes to transform, all by default
1564 padding (tuple): Padding factor in each axis for dealiasing, none by default
1565 grid (tuple): Processor grid, passed on to `PFFT`
1567 Returns:
1568 pySDC.helpers.fft_helper.PFFT: Transform object, or None in 1D or without a communicator
1569 """
1570 if self.ndim == 1 or self.comm is None:
1571 return None
1572 from mpi4py_fft import newDistArray
1574 from pySDC.helpers.fft_helper import PFFT
1576 axes = tuple(i for i in range(self.ndim)) if axes is None else axes
1577 padding = list(padding if padding else [1.0 for _ in range(self.ndim)])
1579 def no_transform(u, *args, **kwargs):
1580 return u
1582 transforms = {(i,): (no_transform, no_transform) for i in range(self.ndim)}
1583 for i in axes:
1584 transforms[((i + self.ndim) % self.ndim,)] = (self.axes[i].transform, self.axes[i].itransform)
1586 # "transform" all axes to ensure consistent shapes.
1587 # Transform non-distributable axes last to ensure they are aligned
1588 _axes = tuple(sorted((axis + self.ndim) % self.ndim for axis in axes))
1589 _axes = [axis for axis in _axes if not self.axes[axis].distributable] + sorted(
1590 [axis for axis in _axes if self.axes[axis].distributable]
1591 + [axis for axis in range(self.ndim) if axis not in _axes]
1592 )
1594 pfft = PFFT(
1595 comm=self.comm,
1596 shape=self.global_shape[1:],
1597 axes=_axes, # TODO: control the order of the transforms better
1598 dtype='D',
1599 collapse=False,
1600 backend=self.fft_backend,
1601 comm_backend=self.fft_comm_backend,
1602 padding=padding,
1603 transforms=transforms,
1604 grid=grid,
1605 )
1607 # do a transform to do the planning
1608 _u = newDistArray(pfft, forward_output=False)
1609 pfft.backward(pfft.forward(_u))
1610 return pfft
1612 def get_fft(self, axes=None, direction='object', padding=None, shape=None):
1613 """
1614 When using MPI, we use `PFFT` objects generated by mpi4py-fft
1616 Args:
1617 axes (tuple): Axes you want to transform over
1618 direction (str): use "forward" or "backward" to get functions for performing the transforms or "object" to get the PFFT object
1619 padding (tuple): Padding for dealiasing
1620 shape (tuple): Shape of the transform
1622 Returns:
1623 transform
1624 """
1625 axes = tuple(-i - 1 for i in range(self.ndim)) if axes is None else axes
1626 shape = self.global_shape[1:] if shape is None else shape
1627 padding = (
1628 [
1629 1,
1630 ]
1631 * self.ndim
1632 if padding is None
1633 else padding
1634 )
1635 key = (axes, direction, tuple(padding), tuple(shape))
1637 if key not in self.fft_cache.keys():
1638 if self.comm is None:
1639 assert np.allclose(padding, 1), 'Zero padding is not implemented for non-MPI transforms'
1641 if direction == 'forward':
1642 self.fft_cache[key] = self.xp.fft.fftn
1643 elif direction == 'backward':
1644 self.fft_cache[key] = self.xp.fft.ifftn
1645 elif direction == 'object':
1646 self.fft_cache[key] = None
1647 else:
1648 if direction == 'object':
1649 from pySDC.helpers.fft_helper import PFFT
1651 _fft = PFFT(
1652 comm=self.comm,
1653 shape=shape,
1654 axes=sorted(axes),
1655 dtype='D',
1656 collapse=False,
1657 backend=self.fft_backend,
1658 comm_backend=self.fft_comm_backend,
1659 padding=padding,
1660 )
1661 else:
1662 _fft = self.get_fft(axes=axes, direction='object', padding=padding, shape=shape)
1664 if direction == 'forward':
1665 self.fft_cache[key] = _fft.forward
1666 elif direction == 'backward':
1667 self.fft_cache[key] = _fft.backward
1668 elif direction == 'object':
1669 self.fft_cache[key] = _fft
1671 return self.fft_cache[key]
1673 def local_slice(self, forward_output=True):
1674 """
1675 Get the slices of the global data held by this rank.
1677 Args:
1678 forward_output (bool): Slices in spectral space if True, in physical space otherwise
1680 Returns:
1681 list of slice: Local slice in each axis
1682 """
1683 if self.fft_obj:
1684 return self.get_pfft().local_slice(forward_output=forward_output)
1685 else:
1686 return [slice(0, me.N) for me in self.axes]
1688 def global_slice(self, forward_output=True):
1689 """
1690 Get slices spanning the global data.
1692 Args:
1693 forward_output (bool): Slices in spectral space if True, in physical space otherwise
1695 Returns:
1696 list of slice: Global slice in each axis
1697 """
1698 if self.fft_obj:
1699 return [slice(0, me) for me in self.fft_obj.global_shape(forward_output=forward_output)]
1700 else:
1701 return self.local_slice(forward_output=forward_output)
1703 def setup_fft(self, real_spectral_coefficients=False):
1704 """
1705 This function must be called after all axes have been setup in order to prepare the local shapes of the data.
1706 This must also be called before setting up any BCs.
1708 Args:
1709 real_spectral_coefficients (bool): Allow only real coefficients in spectral space
1710 """
1711 if len(self.components) == 0:
1712 self.add_component('u')
1714 self.global_shape = (len(self.components),) + tuple(me.N for me in self.axes)
1716 axes = tuple(i for i in range(len(self.axes)))
1717 self.fft_obj = self.get_pfft(axes=axes)
1719 self.init = (
1720 np.empty(shape=self.global_shape)[
1721 (
1722 ...,
1723 *self.local_slice(False),
1724 )
1725 ].shape,
1726 self.comm,
1727 np.dtype('float'),
1728 )
1729 self.init_physical = (
1730 np.empty(shape=self.global_shape)[
1731 (
1732 ...,
1733 *self.local_slice(False),
1734 )
1735 ].shape,
1736 self.comm,
1737 np.dtype('float'),
1738 )
1739 self.init_forward = (
1740 np.empty(shape=self.global_shape)[
1741 (
1742 ...,
1743 *self.local_slice(True),
1744 )
1745 ].shape,
1746 self.comm,
1747 np.dtype('float') if real_spectral_coefficients else np.dtype('complex128'),
1748 )
1750 self.BC_mat = self.get_empty_operator_matrix()
1751 self.BC_rhs_mask = self.newDistArray().astype(bool)
1753 def newDistArray(self, pfft=None, forward_output=True, val=0, rank=1, view=False):
1754 """
1755 Get an empty distributed array. This is almost a copy of the function of the same name from mpi4py-fft, but
1756 takes care of all the solution components in the tensor.
1757 """
1758 if self.comm is None:
1759 return self.xp.zeros(self.init[0], dtype=self.init[2])
1760 from mpi4py_fft.distarray import DistArray
1762 pfft = pfft if pfft else self.get_pfft()
1763 if pfft is None:
1764 if forward_output:
1765 return self.u_init_forward
1766 else:
1767 return self.u_init
1769 global_shape = pfft.global_shape(forward_output)
1770 p0 = pfft.pencil[forward_output]
1771 if forward_output is True:
1772 dtype = pfft.forward.output_array.dtype
1773 else:
1774 dtype = pfft.forward.input_array.dtype
1775 global_shape = (self.ncomponents,) * rank + global_shape
1777 if pfft.xfftn[0].backend in ["cupy", "cupyx-scipy"]:
1778 from mpi4py_fft.distarrayCuPy import DistArrayCuPy as darraycls
1779 else:
1780 darraycls = DistArray
1782 z = darraycls(global_shape, subcomm=p0.subcomm, val=val, dtype=dtype, alignment=p0.axis, rank=rank)
1783 return z.v if view else z
1785 def infer_alignment(self, u, forward_output, padding=None, **kwargs):
1786 """
1787 Find the axes in which `u` is not distributed, by comparing its local shape to the global shape of arrays of the
1788 transform object from `get_pfft`. With `padding`, transforms padded in subsets of the axes are tried until one
1789 fits.
1791 Args:
1792 u: Data with the solution components in the first axis
1793 forward_output (bool): Compare to arrays in spectral space if True, in physical space otherwise
1794 padding (tuple): Padding factor in each axis
1795 **kwargs: Passed on to `get_pfft`
1797 Returns:
1798 list of int: Aligned axes, not counting the component axis, or [0] without a communicator
1799 """
1800 if self.comm is None:
1801 return [0]
1803 def _alignment(pfft):
1804 _arr = self.newDistArray(pfft, forward_output=forward_output)
1805 _aligned_axes = [i for i in range(self.ndim) if _arr.global_shape[i + 1] == u.shape[i + 1]]
1806 return _aligned_axes
1808 if padding is None:
1809 pfft = self.get_pfft(**kwargs)
1810 aligned_axes = _alignment(pfft)
1811 else:
1812 if self.ndim == 2:
1813 padding_options = [(1.0, padding[1]), (padding[0], 1.0), padding, (1.0, 1.0)]
1814 elif self.ndim == 3:
1815 padding_options = [
1816 (1.0, 1.0, padding[2]),
1817 (1.0, padding[1], 1.0),
1818 (padding[0], 1.0, 1.0),
1819 (1.0, padding[1], padding[2]),
1820 (padding[0], 1.0, padding[2]),
1821 (padding[0], padding[1], 1.0),
1822 padding,
1823 (1.0, 1.0, 1.0),
1824 ]
1825 else:
1826 raise NotImplementedError(f'Don\'t know how to infer alignment in {self.ndim}D!')
1827 for _padding in padding_options:
1828 pfft = self.get_pfft(padding=_padding, **kwargs)
1829 aligned_axes = _alignment(pfft)
1830 if len(aligned_axes) > 0:
1831 self.logger.debug(
1832 f'Found alignment of array with size {u.shape}: {aligned_axes} using padding {_padding}'
1833 )
1834 break
1836 assert len(aligned_axes) > 0, f'Found no aligned axes for array of size {u.shape}!'
1837 return aligned_axes
1839 def redistribute(self, u, axis, forward_output, **kwargs):
1840 """
1841 Copy `u` into a new distributed array aligned in `axis`, i.e. not distributed in that axis.
1843 Args:
1844 u: Data with the solution components in the first axis
1845 axis (int): Axis to align in
1846 forward_output (bool): Layout of the new array: spectral space if True, physical space otherwise
1847 **kwargs: Passed on to `get_pfft`, e.g. `padding`
1849 Returns:
1850 mpi4py_fft.DistArray: Redistributed copy of `u`, or `u` itself without a communicator
1851 """
1852 if self.comm is None:
1853 return u
1855 pfft = self.get_pfft(**kwargs)
1856 _arr = self.newDistArray(pfft, forward_output=forward_output)
1858 # the transforms are complex to complex, so the local shapes in physical and spectral space agree and the
1859 # alignment can be inferred from the physical shapes either way
1860 u_alignment = self.infer_alignment(u, forward_output=False, **kwargs)
1861 for alignment in u_alignment:
1862 _arr = _arr.redistribute(alignment)
1863 if _arr.shape == u.shape:
1864 _arr[...] = u
1865 return _arr.redistribute(axis)
1867 raise Exception(
1868 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}'
1869 )
1871 def transform(self, u, *args, axes=None, padding=None, **kwargs):
1872 """
1873 Transform all solution components from physical to spectral space along `axes`.
1874 With a distributed transform, `u` is first redistributed if its local shape does not fit.
1876 Args:
1877 u: Data with the solution components in the first axis
1878 axes (tuple): Axes to transform, all by default
1879 padding (tuple): Padding factor in each axis for dealiasing, only used with a distributed transform
1880 **kwargs: Passed on to `get_pfft`
1882 Returns:
1883 Transformed data in a new array, which is distributed if the transform is
1884 """
1885 pfft = self.get_pfft(*args, axes=axes, padding=padding, **kwargs)
1887 if pfft is None:
1888 if padding is not None and any(me != 1 for me in padding):
1889 raise NotImplementedError('Padding is only implemented for distributed FFTs in more than one dimension')
1890 axes = axes if axes else tuple(i for i in range(self.ndim))
1891 u_hat = u.copy()
1892 for i in axes:
1893 _axis = 1 + i if i >= 0 else i
1894 u_hat = self.axes[i].transform(u_hat, axes=(_axis,))
1895 return u_hat
1897 _in = self.newDistArray(pfft, forward_output=False, rank=1)
1898 _out = self.newDistArray(pfft, forward_output=True, rank=1)
1900 if _in.shape == u.shape:
1901 _in[...] = u
1902 else:
1903 _in[...] = self.redistribute(u, axis=_in.alignment, forward_output=False, padding=padding, **kwargs)
1905 for i in range(self.ncomponents):
1906 pfft.forward(_in[i], _out[i], normalize=False)
1908 if padding is not None:
1909 _out /= np.prod(padding)
1910 return _out
1912 def itransform(self, u, *args, axes=None, padding=None, **kwargs):
1913 """
1914 Transform all solution components from spectral to physical space along `axes`, the inverse of `transform`.
1915 With a distributed transform, `u` is first redistributed if its local shape does not fit.
1917 Args:
1918 u: Data with the solution components in the first axis
1919 axes (tuple): Axes to transform, all by default
1920 padding (tuple): Padding factor in each axis for dealiasing, only used with a distributed transform. The
1921 padded
1922 resolution must be an integer.
1923 **kwargs: Passed on to `get_pfft`
1925 Returns:
1926 Transformed data in a new array, which is distributed if the transform is
1927 """
1928 if padding is not None:
1929 assert all(
1930 (self.axes[i].N * padding[i]) % 1 == 0 for i in range(self.ndim)
1931 ), 'Cannot do this padding with this resolution. Resulting resolution must be integer'
1933 pfft = self.get_pfft(*args, axes=axes, padding=padding, **kwargs)
1934 if pfft is None:
1935 if padding is not None and any(me != 1 for me in padding):
1936 raise NotImplementedError('Padding is only implemented for distributed FFTs in more than one dimension')
1937 axes = axes if axes else tuple(i for i in range(self.ndim))
1938 u_hat = u.copy()
1939 for i in axes:
1940 _axis = 1 + i if i >= 0 else i
1941 u_hat = self.axes[i].itransform(u_hat, axes=(_axis,))
1942 return u_hat
1944 _in = self.newDistArray(pfft, forward_output=True, rank=1)
1945 _out = self.newDistArray(pfft, forward_output=False, rank=1)
1947 if _in.shape == u.shape:
1948 _in[...] = u
1949 else:
1950 _in[...] = self.redistribute(u, axis=_in.alignment, forward_output=True, padding=padding, **kwargs)
1952 for i in range(self.ncomponents):
1953 pfft.backward(_in[i], _out[i], normalize=True)
1955 if padding is not None:
1956 _out *= np.prod(padding)
1957 return _out
1959 def get_local_slice_of_1D_matrix(self, M, axis):
1960 """
1961 Get the local version of a 1D matrix. When using distributed FFTs, each rank will carry only a subset of modes,
1962 which you can sort out via the `SpectralHelper.local_slice()` attribute. When constructing a 1D matrix, you can
1963 use this method to get the part corresponding to the modes carried by this rank.
1965 Args:
1966 M (sparse matrix): Global 1D matrix you want to get the local version of
1967 axis (int): Direction in which you want the local version. You will get the global matrix in other directions.
1969 Returns:
1970 sparse local matrix
1971 """
1972 return M.tocsc()[self.local_slice(True)[axis], self.local_slice(True)[axis]]
1974 def expand_matrix_ND(self, matrix, aligned):
1975 """
1976 Expand a global 1D matrix acting along axis `aligned` to the local part of the ND matrix, by Kronecker products
1977 with identities in the other axes. Uses the local slices in spectral space in every axis.
1979 Args:
1980 matrix (sparse matrix): Global 1D matrix
1981 aligned (int): Axis the matrix acts along
1983 Returns:
1984 sparse local ND matrix in CSC format
1985 """
1986 sp = self.sparse_lib
1987 axes = np.delete(np.arange(self.ndim), aligned)
1988 ndim = len(axes) + 1
1990 if ndim == 1:
1991 mat = matrix
1992 elif ndim == 2:
1993 axis = axes[0]
1994 I1D = sp.eye(self.axes[axis].N)
1996 mats = [None] * ndim
1997 mats[aligned] = self.get_local_slice_of_1D_matrix(matrix, aligned)
1998 mats[axis] = self.get_local_slice_of_1D_matrix(I1D, axis)
2000 mat = sp.kron(*mats)
2001 elif ndim == 3:
2003 mats = [None] * ndim
2004 mats[aligned] = self.get_local_slice_of_1D_matrix(matrix, aligned)
2005 for axis in axes:
2006 I1D = sp.eye(self.axes[axis].N)
2007 mats[axis] = self.get_local_slice_of_1D_matrix(I1D, axis)
2009 mat = sp.kron(mats[0], sp.kron(*mats[1:]))
2011 else:
2012 raise NotImplementedError(f'Matrix expansion not implemented for {ndim} dimensions!')
2014 mat = self.eliminate_zeros(mat)
2015 return mat
2017 def get_filter_matrix(self, axis, **kwargs):
2018 """
2019 Get bandpass filter along `axis`. See the documentation `get_filter_matrix` in the 1D bases for what kwargs are
2020 admissible.
2022 Returns:
2023 sparse bandpass matrix
2024 """
2025 if self.ndim == 1:
2026 return self.axes[0].get_filter_matrix(**kwargs)
2028 mats = [base.get_Id() for base in self.axes]
2029 mats[axis] = self.axes[axis].get_filter_matrix(**kwargs)
2030 return self.sparse_lib.kron(*mats)
2032 def get_differentiation_matrix(self, axes, **kwargs):
2033 """
2034 Get differentiation matrix along specified axis. `kwargs` are forwarded to the 1D base implementation.
2036 Args:
2037 axes (tuple): Axes along which to differentiate.
2039 Returns:
2040 sparse differentiation matrix
2041 """
2042 D = self.expand_matrix_ND(self.axes[axes[0]].get_differentiation_matrix(**kwargs), axes[0])
2043 for axis in axes[1:]:
2044 _D = self.axes[axis].get_differentiation_matrix(**kwargs)
2045 D = D @ self.expand_matrix_ND(_D, axis)
2047 self.logger.debug(f'Set up differentiation matrix along axes {axes} with kwargs {kwargs}')
2048 return D
2050 def get_integration_matrix(self, axes):
2051 """
2052 Get integration matrix to integrate along specified axis.
2054 Args:
2055 axes (tuple): Axes along which to integrate over.
2057 Returns:
2058 sparse integration matrix
2059 """
2060 S = self.expand_matrix_ND(self.axes[axes[0]].get_integration_matrix(), axes[0])
2061 for axis in axes[1:]:
2062 _S = self.axes[axis].get_integration_matrix()
2063 S = S @ self.expand_matrix_ND(_S, axis)
2065 return S
2067 def get_Id(self):
2068 """
2069 Get identity matrix
2071 Returns:
2072 sparse identity matrix
2073 """
2074 I = self.expand_matrix_ND(self.axes[0].get_Id(), 0)
2075 for axis in range(1, self.ndim):
2076 _I = self.axes[axis].get_Id()
2077 I = I @ self.expand_matrix_ND(_I, axis)
2078 return I
2080 def get_Dirichlet_recombination_matrix(self, axis=-1):
2081 """
2082 Get Dirichlet recombination matrix along axis. Not that it only makes sense in directions discretized with variations of Chebychev bases.
2084 Args:
2085 axis (int): Axis you discretized with Chebychev
2087 Returns:
2088 sparse matrix
2089 """
2090 C1D = self.axes[axis].get_Dirichlet_recombination_matrix()
2091 return self.expand_matrix_ND(C1D, axis)
2093 def get_basis_change_matrix(self, axes=None, **kwargs):
2094 """
2095 Some spectral bases do a change between bases while differentiating. This method returns matrices that changes the basis to whatever you want.
2096 Refer to the methods of the same name of the 1D bases to learn what parameters you need to pass here as `kwargs`.
2098 Args:
2099 axes (tuple): Axes along which to change basis.
2101 Returns:
2102 sparse basis change matrix
2103 """
2104 axes = tuple(-i - 1 for i in range(self.ndim)) if axes is None else axes
2106 C = self.expand_matrix_ND(self.axes[axes[0]].get_basis_change_matrix(**kwargs), axes[0])
2107 for axis in axes[1:]:
2108 _C = self.axes[axis].get_basis_change_matrix(**kwargs)
2109 C = C @ self.expand_matrix_ND(_C, axis)
2111 self.logger.debug(f'Set up basis change matrix along axes {axes} with kwargs {kwargs}')
2112 return C