diff --git a/pySDC/helpers/spectral_helper.py b/pySDC/helpers/spectral_helper.py new file mode 100644 index 0000000000..2cf912fbf4 --- /dev/null +++ b/pySDC/helpers/spectral_helper.py @@ -0,0 +1,2025 @@ +import numpy as np +import scipy +from pySDC.implementations.datatype_classes.mesh import mesh +from scipy.special import factorial + + +class SpectralHelper1D: + """ + Abstract base class for 1D spectral discretizations. Defines a common interface with parameters and functions that + all bases need to have. + + When implementing new bases, please take care to use the modules that are supplied as class attributes to enable + the code for GPUs. + + Attributes: + N (int): Resolution + x0 (float): Coordinate of left boundary + x1 (float): Coordinate of right boundary + L (float): Length of the domain + useGPU (bool): Whether to use GPUs + + """ + + fft_lib = scipy.fft + sparse_lib = scipy.sparse + linalg = scipy.sparse.linalg + xp = np + + def __init__(self, N, x0=None, x1=None, useGPU=False): + """ + Constructor + + Args: + N (int): Resolution + x0 (float): Coordinate of left boundary + x1 (float): Coordinate of right boundary + useGPU (bool): Whether to use GPUs + """ + self.N = N + self.x0 = x0 + self.x1 = x1 + self.L = x1 - x0 + self.useGPU = useGPU + + if useGPU: + self.setup_GPU() + + @classmethod + def setup_GPU(cls): + """switch to GPU modules""" + import cupy as cp + import cupyx.scipy.sparse as sparse_lib + import cupyx.scipy.sparse.linalg as linalg + import cupyx.scipy.fft as fft_lib + from pySDC.implementations.datatype_classes.cupy_mesh import cupy_mesh + + cls.xp = cp + cls.sparse_lib = sparse_lib + cls.linalg = linalg + cls.fft_lib = fft_lib + + def get_Id(self): + """ + Get identity matrix + + Returns: + sparse diagonal identity matrix + """ + return self.sparse_lib.eye(self.N) + + def get_zero(self): + """ + Get a matrix with all zeros of the correct size. + + Returns: + sparse matrix with zeros everywhere + """ + return 0 * self.get_Id() + + def get_differentiation_matrix(self): + raise NotImplementedError() + + def get_integration_matrix(self): + raise NotImplementedError() + + def get_wavenumbers(self): + """ + Get the grid in spectral space + """ + raise NotImplementedError + + def get_empty_operator_matrix(self, S, O): + """ + Return a matrix of operators to be filled with the connections between the solution components. + + Args: + S (int): Number of components in the solution + O (sparse matrix): Zero matrix used for initialization + + Returns: + list of lists containing sparse zeros + """ + return [[O for _ in range(S)] for _ in range(S)] + + def get_basis_change_matrix(self, *args, **kwargs): + """ + Some spectral discretization change the basis during differentiation. This method can be used to transfer + between the various bases. + + This method accepts arbitrary arguments that may not be used in order to provide an easy interface for multi- + dimensional bases. For instance, you may combine an FFT discretization with an ultraspherical discretization. + The FFT discretization will always be in the same base, but the ultraspherical discretization uses a different + base for every derivative. You can then ask all bases for transfer matrices from one ultraspherical derivative + base to the next. The FFT discretization will ignore this and return an identity while the ultraspherical + discretization will return the desired matrix. After a Kronecker product, you get the 2D version of the matrix + you want. This is what the `SpectralHelper` does when you call the method of the same name on it. + + Returns: + sparse bases change matrix + """ + return self.sparse_lib.eye(self.N) + + def get_BC(self, kind): + """ + To facilitate boundary conditions (BCs) we use either a basis where all functions satisfy the BCs automatically, + as is the case in FFT basis for periodic BCs, or boundary bordering. In boundary bordering, specific lines in + the matrix are replaced by the boundary conditions as obtained by this method. + + Args: + kind (str): The type of BC you want to implement please refer to the implementations of this method in the + individual 1D bases for what is implemented + + Returns: + self.xp.array: Boundary condition + """ + raise NotImplementedError(f'No boundary conditions of {kind=!r} implemented!') + + def get_filter_matrix(self, kmin=0, kmax=None): + """ + Get a bandpass filter. + + Args: + kmin (int): Lower limit of the bandpass filter + kmax (int): Upper limit of the bandpass filter + + Returns: + sparse matrix + """ + + k = abs(self.get_wavenumbers()) + + kmax = max(k) if kmax is None else kmax + + mask = self.xp.logical_or(k >= kmax, k < kmin) + + if self.useGPU: + Id = self.get_Id().get() + else: + Id = self.get_Id() + F = Id.tolil() + F[:, mask] = 0 + return F.tocsc() + + def get_1dgrid(self): + """ + Get the grid in physical space + + Returns: + self.xp.array: Grid + """ + raise NotImplementedError + + +class ChebychevHelper(SpectralHelper1D): + """ + The Chebychev base consists of special kinds of polynomials, with the main advantage that you can easily transform + between physical and spectral space by discrete cosine transform. + The differentiation in the Chebychev T base is dense, but can be preconditioned to yield a differentiation operator + that moves to Chebychev U basis during differentiation, which is sparse. When using this technique, problems need to + be formulated in first order formulation. + + This implementation is largely based on the Dedalus paper (arXiv:1905.10388). + """ + + def __init__(self, *args, transform_type='fft', x0=-1, x1=1, **kwargs): + """ + Constructor. + Please refer to the parent class for additional arguments. Notably, you have to supply a resolution `N` and you + may choose to run on GPUs via the `useGPU` argument. + + Args: + transform_type ('fft' or 'dct'): Either use DCT functions directly implemented in the transform library or + use the FFT from the library to compute the DCT + x0 (float): Coordinate of left boundary. Note that only -1 is currently implented + x1 (float): Coordinate of right boundary. Note that only +1 is currently implented + """ + assert x0 == -1 + assert x1 == 1 + super().__init__(*args, x0=x0, x1=x1, **kwargs) + self.transform_type = transform_type + + if self.transform_type == 'fft': + self.get_fft_utils() + + self.cache = {} + self.norm = self.get_norm() + + def get_1dgrid(self): + ''' + Generates a 1D grid with Chebychev points. These are clustered at the boundary. You need this kind of grid to + use discrete cosine transformation (DCT) to get the Chebychev representation. If you want a different grid, you + need to do an affine transformation before any Chebychev business. + + Returns: + numpy.ndarray: 1D grid + ''' + return self.xp.cos(np.pi / self.N * (self.xp.arange(self.N) + 0.5)) + + def get_wavenumbers(self): + """Get the domain in spectral space""" + return self.xp.arange(self.N) + + def get_conv(self, name, N=None): + ''' + Get conversion matrix between different kinds of polynomials. The supported kinds are + - T: Chebychev polynomials of first kind + - U: Chebychev polynomials of second kind + - D: Dirichlet recombination. + + You get the desired matrix by choosing a name as ``A2B``. I.e. ``T2U`` for the conversion matrix from T to U. + Once generates matrices are cached. So feel free to call the method as often as you like. + + Args: + name (str): Conversion code, e.g. 'T2U' + N (int): Size of the matrix (optional) + + Returns: + scipy.sparse: Sparse conversion matrix + ''' + if name in self.cache.keys() and not N: + return self.cache[name] + + N = N if N else self.N + sp = self.sparse_lib + xp = self.xp + + def get_forward_conv(name): + if name == 'T2U': + mat = (sp.eye(N) - sp.diags(xp.ones(N - 2), offsets=+2)) / 2.0 + mat[:, 0] *= 2 + elif name == 'D2T': + mat = sp.eye(N) - sp.diags(xp.ones(N - 2), offsets=+2) + elif name[0] == name[-1]: + mat = self.sparse_lib.eye(self.N) + else: + raise NotImplementedError(f'Don\'t have conversion matrix {name!r}') + return mat + + try: + mat = get_forward_conv(name) + except NotImplementedError as E: + try: + fwd = get_forward_conv(name[::-1]) + import scipy.sparse as sp + + if self.sparse_lib == sp: + mat = self.sparse_lib.linalg.inv(fwd.tocsc()) + else: + mat = self.sparse_lib.csc_matrix(sp.linalg.inv(fwd.tocsc().get())) + except NotImplementedError: + raise NotImplementedError from E + + self.cache[name] = mat + return mat + + def get_basis_change_matrix(self, conv='T2T', **kwargs): + """ + As the differentiation matrix in Chebychev-T base is dense but is sparse when simultaneously changing base to + Chebychev-U, you may need a basis change matrix to transfer the other matrices as well. This function returns a + conversion matrix from `ChebychevHelper.get_conv`. Not that `**kwargs` are used to absorb arguments for other + bases, see documentation of `SpectralHelper1D.get_basis_change_matrix`. + + Args: + conv (str): Conversion code, i.e. T2U + + Returns: + Sparse conversion matrix + """ + return self.get_conv(conv) + + def get_integration_matrix(self, lbnd=0): + """ + Get matrix for integration + + Args: + lbnd (float): Lower bound for integration, only 0 is currently implemented + + Returns: + Sparse integration matrix + """ + S = self.sparse_lib.diags(1 / (self.xp.arange(self.N - 1) + 1), offsets=-1) @ self.get_conv('T2U') + n = self.xp.arange(self.N) + if lbnd == 0: + S = S.tocsc() + S[0, 1::2] = ( + (n / (2 * (self.xp.arange(self.N) + 1)))[1::2] + * (-1) ** (self.xp.arange(self.N // 2)) + / (np.append([1], self.xp.arange(self.N // 2 - 1) + 1)) + ) + else: + raise NotImplementedError(f'This function allows to integrate only from x=0, you attempted from x={lbnd}.') + return S + + def get_differentiation_matrix(self, p=1): + ''' + Keep in mind that the T2T differentiation matrix is dense. + + Args: + p (int): Derivative you want to compute + + Returns: + numpy.ndarray: Differentiation matrix + ''' + D = self.xp.zeros((self.N, self.N)) + for j in range(self.N): + for k in range(j): + D[k, j] = 2 * j * ((j - k) % 2) + + D[0, :] /= 2 + return self.sparse_lib.csc_matrix(self.xp.linalg.matrix_power(D, p)) + + def get_norm(self, N=None): + ''' + Get normalization for converting Chebychev coefficients and DCT + + Args: + N (int, optional): Resolution + + Returns: + self.xp.array: Normalization + ''' + N = self.N if N is None else N + norm = self.xp.ones(N) / N + norm[0] /= 2 + return norm + + def get_fft_shuffle(self, forward, N): + """ + In order to more easily parallelize using distributed FFT libraries, we express the DCT via an FFT following + doi.org/10.1109/TASSP.1980.1163351. The idea is based on reshuffling the data to be periodic and rotating it + in the complex plane. This function returns a mask to do the shuffling. + + Args: + forward (bool): Whether you want the shuffle for forward transform or backward transform + N (int): size of the grid + + Returns: + self.xp.array: Use as mask + """ + xp = self.xp + if forward: + return xp.append(xp.arange((N + 1) // 2) * 2, -xp.arange(N // 2) * 2 - 1 - N % 2) + else: + mask = xp.zeros(N, dtype=int) + mask[: N - N % 2 : 2] = xp.arange(N // 2) + mask[1::2] = N - xp.arange(N // 2) - 1 + mask[-1] = N // 2 + return mask + + def get_fft_shift(self, forward, N): + """ + As described in the docstring for `get_fft_shuffle`, we need to rotate in the complex plane in order to use FFT for DCT. + + Args: + forward (bool): Whether you want the rotation for forward transform or backward transform + N (int): size of the grid + + Returns: + self.xp.array: Rotation + """ + k = self.get_wavenumbers() + norm = self.get_norm() + xp = self.xp + if forward: + return 2 * xp.exp(-1j * np.pi * k / (2 * N) + 0j * np.pi / 4) * norm + else: + shift = xp.exp(1j * np.pi * k / (2 * N)) + shift[0] = 0.5 + return shift / norm + + def get_fft_utils(self): + """ + Get the required utilities for using FFT to do DCT as described in the docstring for `get_fft_shuffle` and keep + them cached. + """ + self.fft_utils = { + 'fwd': {}, + 'bck': {}, + } + + # forwards transform + self.fft_utils['fwd']['shuffle'] = self.get_fft_shuffle(True, self.N) + self.fft_utils['fwd']['shift'] = self.get_fft_shift(True, self.N) + + # backwards transform + self.fft_utils['bck']['shuffle'] = self.get_fft_shuffle(False, self.N) + self.fft_utils['bck']['shift'] = self.get_fft_shift(False, self.N) + + return self.fft_utils + + def transform(self, u, axis=-1, **kwargs): + """ + 1D DCT along axis. `kwargs` will be passed on to the FFT library. + + Args: + u: Data you want to transform + axis (int): Axis you want to transform along + + Returns: + Data in spectral space + """ + if self.transform_type.lower() == 'dct': + return self.fft_lib.dct(u, axis=axis, **kwargs) * self.norm + elif self.transform_type.lower() == 'fft': + result = u.copy() + + shuffle = [slice(0, s, 1) for s in u.shape] + shuffle[axis] = self.fft_utils['fwd']['shuffle'] + + v = u[(*shuffle,)] + + V = self.fft_lib.fft(v, axis=axis, **kwargs) + + expansion = [np.newaxis for _ in u.shape] + expansion[axis] = slice(0, u.shape[axis], 1) + + V *= self.fft_utils['fwd']['shift'][(*expansion,)] + + result.real[...] = V.real[...] + return result + else: + raise NotImplementedError(f'Please choose a transform type from fft and dct, not {self.transform_type=}') + + def itransform(self, u, axis=-1): + """ + 1D inverse DCT along axis. + + Args: + u: Data you want to transform + axis (int): Axis you want to transform along + + Returns: + Data in physical space + """ + assert self.norm.shape[0] == u.shape[axis] + + if self.transform_type == 'dct': + return self.fft_lib.idct(u / self.norm, axis=axis) + elif self.transform_type == 'fft': + result = u.copy() + + expansion = [np.newaxis for _ in u.shape] + expansion[axis] = slice(0, u.shape[axis], 1) + + v = self.fft_lib.ifft(u * self.fft_utils['bck']['shift'][(*expansion,)], axis=axis) + + shuffle = [slice(0, s, 1) for s in u.shape] + shuffle[axis] = self.fft_utils['bck']['shuffle'] + V = v[(*shuffle,)] + + result.real[...] = V.real[...] + return result + else: + raise NotImplementedError + + def get_BC(self, kind, **kwargs): + """ + Get boundary condition row for boundary bordering. `kwargs` will be passed on to implementations of the BC of + the kind you choose. Specifically, `x` for `'dirichlet'` boundary condition, which is the coordinate at which to + set the BC. + + Args: + kind ('integral' or 'dirichlet'): Kind of boundary condition you want + """ + if kind.lower() == 'integral': + return self.get_integ_BC_row(**kwargs) + elif kind.lower() == 'dirichlet': + return self.get_Dirichlet_BC_row(**kwargs) + else: + return super().get_BC(kind) + + def get_integ_BC_row(self): + """ + Get a row for generating integral BCs with T polynomials. + It returns the values of the integrals of T polynomials over the entire interval. + + Returns: + self.xp.ndarray: Row to put into a matrix + """ + n = self.xp.arange(self.N) + 1 + me = self.xp.zeros_like(n).astype(float) + me[2:] = ((-1) ** n[1:-1] + 1) / (1 - n[1:-1] ** 2) + me[0] = 2.0 + return me + + def get_Dirichlet_BC_row(self, x): + """ + Get a row for generating Dirichlet BCs at x with T polynomials. + It returns the values of the T polynomials at x. + + Args: + x (float): Position of the boundary condition + + Returns: + self.xp.ndarray: Row to put into a matrix + """ + if x == -1: + return (-1) ** self.xp.arange(self.N) + elif x == 1: + return self.xp.ones(self.N) + elif x == 0: + n = (1 + (-1) ** self.xp.arange(self.N)) / 2 + n[2::4] *= -1 + return n + else: + raise NotImplementedError(f'Don\'t know how to generate Dirichlet BC\'s at {x=}!') + + def get_Dirichlet_recombination_matrix(self): + ''' + Get matrix for Dirichlet recombination, which changes the basis to have sparse boundary conditions. + This makes for a good right preconditioner. + + Returns: + scipy.sparse: Sparse conversion matrix + ''' + N = self.N + sp = self.sparse_lib + xp = self.xp + + return sp.eye(N) - sp.diags(xp.ones(N - 2), offsets=+2) + + +class UltrasphericalHelper(ChebychevHelper): + """ + This implementation follows https://doi.org/10.1137/120865458. + The ultraspherical method works in Chebychev polynomials as well, but also uses various Gegenbauer polynomials. + 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. + There are also conversion operators from one derivative basis to the next that are sparse. + + 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. + 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. + """ + + def get_differentiation_matrix(self, p=1): + """ + Notice that while sparse, this matrix is not diagonal, which means the inversion cannot be parallelized easily. + + Args: + p (int): Order of the derivative + + Returns: + sparse differentiation matrix + """ + sp = self.sparse_lib + xp = self.xp + N = self.N + l = p + return 2 ** (l - 1) * factorial(l - 1) * sp.diags(xp.arange(N - l) + l, offsets=l) + + def get_S(self, lmbda): + """ + Get matrix for bumping the derivative base by one from lmbda to lmbda + 1. This is the same language as in + https://doi.org/10.1137/120865458. + + Args: + lmbda (int): Ingoing derivative base + + Returns: + sparse matrix: Conversion from derivative base lmbda to lmbda + 1 + """ + N = self.N + + if lmbda == 0: + sp = scipy.sparse + mat = ((sp.eye(N) - sp.diags(np.ones(N - 2), offsets=+2)) / 2.0).tolil() + mat[:, 0] *= 2 + else: + sp = self.sparse_lib + xp = self.xp + mat = sp.diags(lmbda / (lmbda + xp.arange(N))) - sp.diags( + lmbda / (lmbda + 2 + xp.arange(N - 2)), offsets=+2 + ) + + return self.sparse_lib.csc_matrix(mat) + + def get_basis_change_matrix(self, p_in=0, p_out=0, **kwargs): + """ + Get a conversion matrix from derivative base `p_in` to `p_out`. + + Args: + p_out (int): Resulting derivative base + p_in (int): Ingoing derivative base + """ + mat_fwd = self.sparse_lib.eye(self.N) + for i in range(min([p_in, p_out]), max([p_in, p_out])): + mat_fwd = self.get_S(i) @ mat_fwd + + if p_out > p_in: + return mat_fwd + + else: + # We have to invert the matrix on CPU because the GPU equivalent is not implemented in CuPy at the time of writing. + import scipy.sparse as sp + + if self.useGPU: + mat_fwd = mat_fwd.get() + + mat_bck = sp.linalg.inv(mat_fwd.tocsc()) + + return self.sparse_lib.csc_matrix(mat_bck) + + def get_integration_matrix(self): + """ + Get an integration matrix. Please use `UltrasphericalHelper.get_integration_constant` afterwards to compute the + integration constant such that integration starts from x=-1. + + Example: + + .. code-block:: python + + import numpy as np + from pySDC.helpers.spectral_helper import UltrasphericalHelper + + N = 4 + helper = UltrasphericalHelper(N) + coeffs = np.random.random(N) + coeffs[-1] = 0 + + poly = np.polynomial.Chebyshev(coeffs) + + S = helper.get_integration_matrix() + U_hat = S @ coeffs + U_hat[0] = helper.get_integration_constant(U_hat, axis=-1) + + assert np.allclose(poly.integ(lbnd=-1).coef[:-1], U_hat) + + Returns: + sparse integration matrix + """ + return self.sparse_lib.diags(1 / (self.xp.arange(self.N - 1) + 1), offsets=-1) @ self.get_basis_change_matrix( + p_out=1, p_in=0 + ) + + def get_integration_constant(self, u_hat, axis): + """ + Get integration constant for lower bound of -1. See documentation of `UltrasphericalHelper.get_integration_matrix` for details. + + Args: + u_hat: Solution in spectral space + axis: Axis you want to integrate over + + Returns: + Integration constant, has one less dimension than `u_hat` + """ + slices = [ + None, + ] * u_hat.ndim + slices[axis] = slice(1, u_hat.shape[axis]) + return self.xp.sum(u_hat[(*slices,)] * (-1) ** (self.xp.arange(u_hat.shape[axis] - 1)), axis=axis) + + +class FFTHelper(SpectralHelper1D): + def __init__(self, *args, x0=0, x1=2 * np.pi, **kwargs): + """ + Constructor. + Please refer to the parent class for additional arguments. Notably, you have to supply a resolution `N` and you + may choose to run on GPUs via the `useGPU` argument. + + Args: + transform_type ('fft' or 'dct'): Either use DCT functions directly implemented in the transform library or + use the FFT from the library to compute the DCT + x0 (float, optional): Coordinate of left boundary + x1 (float, optional): Coordinate of right boundary + """ + super().__init__(*args, x0=x0, x1=x1, **kwargs) + + def get_1dgrid(self): + """ + We use equally spaced points including the left boundary and not including the right one, which is the left boundary. + """ + dx = self.L / self.N + return self.xp.arange(self.N) * dx + self.x0 + + def get_wavenumbers(self): + """ + Be careful that this ordering is very unintuitive. + """ + return self.xp.fft.fftfreq(self.N, 1.0 / self.N) * 2 * np.pi / self.L + + def get_differentiation_matrix(self, p=1): + """ + This matrix is diagonal, allowing to invert concurrently. + + Args: + p (int): Order of the derivative + + Returns: + sparse differentiation matrix + """ + k = self.get_wavenumbers() + + if self.useGPU: + # Have to raise the matrix to power p on CPU because the GPU equivalent is not implemented in CuPy at the time of writing. + import scipy.sparse as sp + + D = self.sparse_lib.diags(1j * k).get() + return self.sparse_lib.csc_matrix(sp.linalg.matrix_power(D, p)) + else: + return self.linalg.matrix_power(self.sparse_lib.diags(1j * k), p) + + def get_integration_matrix(self, p=1): + """ + Get integration matrix to compute `p`-th integral over the entire domain. + + Args: + p (int): Order of integral you want to compute + + Returns: + sparse integration matrix + """ + k = self.xp.array(self.get_wavenumbers(), dtype='complex128') + k[0] = 1j * self.L + return self.linalg.matrix_power(self.sparse_lib.diags(1 / (1j * k)), p) + + def transform(self, u, axis=-1, **kwargs): + """ + 1D FFT along axis. `kwargs` are passed on to the FFT library. + + Args: + u: Data you want to transform + axis (int): Axis you want to transform along + + Returns: + transformed data + """ + return self.fft_lib.fft(u, axis=axis, **kwargs) + + def itransform(self, u, axis=-1): + """ + Inverse 1D FFT. + + Args: + u: Data you want to transform + axis (int): Axis you want to transform along + + Returns: + transformed data + """ + return self.fft_lib.ifft(u, axis=axis) + + def get_BC(self, kind): + """ + Get a sort of boundary condition. You can use `kind=integral`, to fix the integral, or you can use `kind=Nyquist`. + The latter is not really a boundary condition, but is used to set the Nyquist mode to some value, preferably zero. + You should set the Nyquist mode zero when the solution in physical space is real and the resolution is even. + + Args: + kind ('integral' or 'nyquist'): Kind of BC + + Returns: + self.xp.ndarray: Boundary condition row + """ + if kind.lower() == 'integral': + return self.get_integ_BC_row() + elif kind.lower() == 'nyquist': + assert ( + self.N % 2 == 0 + ), f'Do not eliminate the Nyquist mode with odd resolution as it is fully resolved. You chose {self.N} in this axis' + BC = self.xp.zeros(self.N) + BC[self.get_Nyquist_mode_index()] = 1 + return BC + else: + return super().get_BC(kind) + + def get_Nyquist_mode_index(self): + """ + Compute the index of the Nyquist mode, i.e. the mode with the lowest wavenumber, which doesn't have a positive + counterpart for even resolution. This means real waves of this wave number cannot be properly resolved and you + are best advised to set this mode zero if representing real functions on even-resolution grids is what you're + after. + + Returns: + int: Index of the Nyquist mode + """ + k = self.get_wavenumbers() + Nyquist_mode = min(k) + return self.xp.where(k == Nyquist_mode)[0][0] + + def get_integ_BC_row(self): + """ + Only the 0-mode has non-zero integral with FFT basis in periodic BCs + """ + me = self.xp.zeros(self.N) + me[0] = self.L / self.N + return me + + +class SpectralHelper: + """ + This class has three functions: + - Easily assemble matrices containing multiple equations + - Direct product of 1D bases to solve problems in more dimensions + - Distribute the FFTs to facilitate concurrency. + + Attributes: + comm (mpi4py.Intracomm): MPI communicator + debug (bool): Perform additional checks at extra computational cost + useGPU (bool): Whether to use GPUs + axes (list): List of 1D bases + components (list): List of strings of the names of components in the equations + full_BCs (list): List of Dictionaries containing all information about the boundary conditions + BC_mat (list): List of lists of sparse matrices to put BCs into and eventually assemble the BC matrix from + BCs (sparse matrix): Matrix containing only the BCs + fft_cache (dict): Cache FFTs of various shapes here to facilitate padding and so on + BC_rhs_mask (self.xp.ndarray): Mask values that contain boundary conditions in the right hand side + BC_zero_index (self.xp.ndarray): Indeces of rows in the matrix that are replaced by BCs + BC_line_zero_matrix (sparse matrix): Matrix that zeros rows where we can then add the BCs in using `BCs` + rhs_BCs_hat (self.xp.ndarray): Boundary conditions in spectral space + global_shape (tuple): Global shape of the solution as in `mpi4py-fft` + local_slice (slice): Local slice of the solution as in `mpi4py-fft` + fft_obj: When using distributed FFTs, this will be a parallel transform object from `mpi4py-fft` + init (tuple): This is the same `init` that is used throughout the problem classes + init_forward (tuple): This is the equivalent of `init` in spectral space + """ + + xp = np + fft_lib = scipy.fft + sparse_lib = scipy.sparse + linalg = scipy.sparse.linalg + dtype = mesh + fft_backend = 'fftw' + fft_comm_backend = 'MPI' + + @classmethod + def setup_GPU(cls): + """switch to GPU modules""" + import cupy as cp + import cupyx.scipy.sparse as sparse_lib + import cupyx.scipy.sparse.linalg as linalg + from pySDC.implementations.datatype_classes.cupy_mesh import cupy_mesh + + cls.xp = cp + cls.sparse_lib = sparse_lib + cls.linalg = linalg + + cls.fft_backend = 'cupy' + cls.fft_comm_backend = 'NCCL' + + cls.dtype = cupy_mesh + + def __init__(self, comm=None, useGPU=False, debug=False): + """ + Constructor + + Args: + comm (mpi4py.Intracomm): MPI communicator + useGPU (bool): Whether to use GPUs + debug (bool): Perform additional checks at extra computational cost + """ + self.comm = comm + self.debug = debug + self.useGPU = useGPU + + if useGPU: + self.setup_GPU() + + self.axes = [] + self.components = [] + + self.full_BCs = [] + self.BC_mat = None + self.BCs = None + + self.fft_cache = {} + + @property + def u_init(self): + """ + Get empty data container in physical space + """ + return self.dtype(self.init) + + @property + def u_init_forward(self): + """ + Get empty data container in spectral space + """ + return self.dtype(self.init_forward) + + @property + def shape(self): + """ + Get shape of individual solution component + """ + return self.init[0][1:] + + @property + def ndim(self): + return len(self.axes) + + @property + def ncomponents(self): + return len(self.components) + + @property + def V(self): + """ + Get domain volume + """ + return np.prod([me.L for me in self.axes]) + + def add_axis(self, base, *args, **kwargs): + """ + Add an axis to the domain by deciding on suitable 1D base. + Arguments to the bases are forwarded using `*args` and `**kwargs`. Please refer to the documentation of the 1D + bases for possible arguments. + + Args: + base (str): 1D spectral method + """ + kwargs['useGPU'] = self.useGPU + + if base.lower() in ['chebychov', 'chebychev', 'cheby', 'chebychovhelper']: + kwargs['transform_type'] = kwargs.get('transform_type', 'fft') + self.axes.append(ChebychevHelper(*args, **kwargs)) + elif base.lower() in ['fft', 'fourier', 'ffthelper']: + self.axes.append(FFTHelper(*args, **kwargs)) + elif base.lower() in ['ultraspherical', 'gegenbauer']: + self.axes.append(UltrasphericalHelper(*args, **kwargs)) + else: + raise NotImplementedError(f'{base=!r} is not implemented!') + self.axes[-1].xp = self.xp + self.axes[-1].sparse_lib = self.sparse_lib + + def add_component(self, name): + """ + Add solution component(s). + + Args: + name (str or list of strings): Name(s) of component(s) + """ + if type(name) in [list, tuple]: + for me in name: + self.add_component(me) + elif type(name) in [str]: + if name in self.components: + raise Exception(f'{name=!r} is already added to this problem!') + self.components.append(name) + else: + raise NotImplementedError + + def index(self, name): + """ + Get the index of component `name`. + + Args: + name (str or list of strings): Name(s) of component(s) + + Returns: + int: Index of the component + """ + if type(name) in [str, int]: + return self.components.index(name) + elif type(name) in [list, tuple]: + return (self.index(me) for me in name) + else: + raise NotImplementedError(f'Don\'t know how to compute index for {type(name)=}') + + def get_empty_operator_matrix(self): + """ + Return a matrix of operators to be filled with the connections between the solution components. + + Returns: + list containing sparse zeros + """ + S = len(self.components) + O = self.get_Id() * 0 + return [[O for _ in range(S)] for _ in range(S)] + + def get_BC(self, axis, kind, line=-1, scalar=False, **kwargs): + """ + Use this method for boundary bordering. It gets the respective matrix row and embeds it into a matrix. + Pay attention that if you have multiple BCs in a single equation, you need to put them in different lines. + Typically, the last line that does not contain a BC is the best choice. + Forward arguments for the boundary conditions using `kwargs`. Refer to documentation of 1D bases for details. + + Args: + axis (int): Axis you want to add the BC to + kind (str): kind of BC, e.g. Dirichlet + line (int): Line you want the BC to go in + scalar (bool): Put the BC in all space positions in the other direction + + Returns: + sparse matrix containing the BC + """ + sp = scipy.sparse + + base = self.axes[axis] + + BC = sp.eye(base.N).tolil() * 0 + if self.useGPU: + BC[line, :] = base.get_BC(kind=kind, **kwargs).get() + else: + BC[line, :] = base.get_BC(kind=kind, **kwargs) + + ndim = len(self.axes) + if ndim == 1: + return self.sparse_lib.csc_matrix(BC) + elif ndim == 2: + axis2 = (axis + 1) % ndim + + if scalar: + _Id = self.sparse_lib.diags(self.xp.append([1], self.xp.zeros(self.axes[axis2].N - 1))) + else: + _Id = self.axes[axis2].get_Id() + + Id = self.get_local_slice_of_1D_matrix(self.axes[axis2].get_Id() @ _Id, axis=axis2) + + if self.useGPU: + Id = Id.get() + + mats = [ + None, + ] * ndim + mats[axis] = self.get_local_slice_of_1D_matrix(BC, axis=axis) + mats[axis2] = Id + return self.sparse_lib.csc_matrix(sp.kron(*mats)) + + def remove_BC(self, component, equation, axis, kind, line=-1, scalar=False, **kwargs): + """ + Remove a BC from the matrix. This is useful e.g. when you add a non-scalar BC and then need to selectively + remove single BCs again, as in incompressible Navier-Stokes, for instance. + Forward arguments for the boundary conditions using `kwargs`. Refer to documentation of 1D bases for details. + + Args: + component (str): Name of the component the BC should act on + equation (str): Name of the equation for the component you want to put the BC in + axis (int): Axis you want to add the BC to + kind (str): kind of BC, e.g. Dirichlet + v: Value of the BC + line (int): Line you want the BC to go in + scalar (bool): Put the BC in all space positions in the other direction + """ + _BC = self.get_BC(axis=axis, kind=kind, line=line, scalar=scalar, **kwargs) + self.BC_mat[self.index(equation)][self.index(component)] -= _BC + + if scalar: + slices = [self.index(equation)] + [ + 0, + ] * self.ndim + slices[axis + 1] = line + else: + slices = ( + [self.index(equation)] + + [slice(0, self.init[0][i + 1]) for i in range(axis)] + + [line] + + [slice(0, self.init[0][i + 1]) for i in range(axis + 1, len(self.axes))] + ) + N = self.axes[axis].N + if (N + line) % N in self.xp.arange(N)[self.local_slice[axis]]: + self.BC_rhs_mask[(*slices,)] = False + + def add_BC(self, component, equation, axis, kind, v, line=-1, scalar=False, **kwargs): + """ + Add a BC to the matrix. Note that you need to convert the list of lists of BCs that this method generates to a + single sparse matrix by calling `setup_BCs` after adding/removing all BCs. + Forward arguments for the boundary conditions using `kwargs`. Refer to documentation of 1D bases for details. + + Args: + component (str): Name of the component the BC should act on + equation (str): Name of the equation for the component you want to put the BC in + axis (int): Axis you want to add the BC to + kind (str): kind of BC, e.g. Dirichlet + v: Value of the BC + line (int): Line you want the BC to go in + scalar (bool): Put the BC in all space positions in the other direction + """ + _BC = self.get_BC(axis=axis, kind=kind, line=line, scalar=scalar, **kwargs) + self.BC_mat[self.index(equation)][self.index(component)] += _BC + self.full_BCs += [ + { + 'component': component, + 'equation': equation, + 'axis': axis, + 'kind': kind, + 'v': v, + 'line': line, + 'scalar': scalar, + **kwargs, + } + ] + + if scalar: + slices = [self.index(equation)] + [ + 0, + ] * self.ndim + slices[axis + 1] = line + if self.comm: + if self.comm.rank == 0: + self.BC_rhs_mask[(*slices,)] = True + else: + self.BC_rhs_mask[(*slices,)] = True + else: + slices = ( + [self.index(equation)] + + [slice(0, self.init[0][i + 1]) for i in range(axis)] + + [line] + + [slice(0, self.init[0][i + 1]) for i in range(axis + 1, len(self.axes))] + ) + N = self.axes[axis].N + if (N + line) % N in self.xp.arange(N)[self.local_slice[axis]]: + slices[axis + 1] -= self.local_slice[axis].start + self.BC_rhs_mask[(*slices,)] = True + + def setup_BCs(self): + """ + Convert the list of lists of BCs to the boundary condition operator. + Also, boundary bordering requires to zero out all other entries in the matrix in rows containing a boundary + condition. This method sets up a suitable sparse matrix to do this. + """ + sp = self.sparse_lib + self.BCs = self.convert_operator_matrix_to_operator(self.BC_mat) + self.BC_zero_index = self.xp.arange(np.prod(self.init[0]))[self.BC_rhs_mask.flatten()] + + diags = self.xp.ones(self.BCs.shape[0]) + diags[self.BC_zero_index] = 0 + self.BC_line_zero_matrix = sp.diags(diags) + + # prepare BCs in spectral space to easily add to the RHS + rhs_BCs = self.put_BCs_in_rhs(self.u_init) + self.rhs_BCs_hat = self.transform(rhs_BCs) + + def check_BCs(self, u): + """ + Check that the solution satisfies the boundary conditions + + Args: + u: The solution you want to check + """ + assert self.ndim < 3 + for axis in range(self.ndim): + BCs = [me for me in self.full_BCs if me["axis"] == axis and not me["scalar"]] + + if len(BCs) > 0: + u_hat = self.transform(u, axes=(axis - self.ndim,)) + for BC in BCs: + kwargs = { + key: value + for key, value in BC.items() + if key not in ['component', 'equation', 'axis', 'v', 'line', 'scalar'] + } + + if axis == 0: + get = self.axes[axis].get_BC(**kwargs) @ u_hat[self.index(BC['component'])] + elif axis == 1: + get = u_hat[self.index(BC['component'])] @ self.axes[axis].get_BC(**kwargs) + want = BC['v'] + assert self.xp.allclose( + get, want + ), f'Unexpected BC in {BC["component"]} in equation {BC["equation"]}, line {BC["line"]}! Got {get}, wanted {want}' + + def put_BCs_in_matrix(self, A): + """ + Put the boundary conditions in a matrix by replacing rows with BCs. + """ + return self.BC_line_zero_matrix @ A + self.BCs + + def put_BCs_in_rhs_hat(self, rhs_hat): + """ + Put the BCs in the right hand side in spectral space for solving. + This function needs no transforms. + + Args: + rhs_hat: Right hand side in spectral space + + Returns: + rhs in spectral space with BCs + """ + ndim = self.ndim + + for axis in range(ndim): + for bc in self.full_BCs: + slices = ( + [slice(0, self.init[0][i + 1]) for i in range(axis)] + + [bc['line']] + + [slice(0, self.init[0][i + 1]) for i in range(axis + 1, len(self.axes))] + ) + if axis == bc['axis']: + _slice = [self.index(bc['equation'])] + slices + N = self.axes[axis].N + if (N + bc['line']) % N in self.xp.arange(N)[self.local_slice[axis]]: + _slice[axis + 1] -= self.local_slice[axis].start + rhs_hat[(*_slice,)] = 0 + + return rhs_hat + self.rhs_BCs_hat + + def put_BCs_in_rhs(self, rhs): + """ + Put the BCs in the right hand side for solving. + This function will transform along each axis individually and add all BCs in that axis. + Consider `put_BCs_in_rhs_hat` to add BCs with no extra transforms needed. + + Args: + rhs: Right hand side in physical space + + Returns: + rhs in physical space with BCs + """ + assert rhs.ndim > 1, 'rhs must not be flattened here!' + + ndim = self.ndim + + for axis in range(ndim): + _rhs_hat = self.transform(rhs, axes=(axis - ndim,)) + + for bc in self.full_BCs: + slices = ( + [slice(0, self.init[0][i + 1]) for i in range(axis)] + + [bc['line']] + + [slice(0, self.init[0][i + 1]) for i in range(axis + 1, len(self.axes))] + ) + if axis == bc['axis']: + _slice = [self.index(bc['equation'])] + slices + + N = self.axes[axis].N + if (N + bc['line']) % N in self.xp.arange(N)[self.local_slice[axis]]: + _slice[axis + 1] -= self.local_slice[axis].start + + _rhs_hat[(*_slice,)] = bc['v'] + + rhs = self.itransform(_rhs_hat, axes=(axis - ndim,)) + + return rhs + + def add_equation_lhs(self, A, equation, relations): + """ + Add the left hand part (that you want to solve implicitly) of an equation to a list of lists of sparse matrices + that you will convert to an operator later. + + Example: + Setup linear operator `L` for 1D heat equation using Chebychev method in first order form and T-to-U + preconditioning: + + .. code-block:: python + helper = SpectralHelper() + + helper.add_axis(base='chebychev', N=8) + helper.add_component(['u', 'ux']) + helper.setup_fft() + + I = helper.get_Id() + Dx = helper.get_differentiation_matrix(axes=(0,)) + T2U = helper.get_basis_change_matrix('T2U') + + L_lhs = { + 'ux': {'u': -T2U @ Dx, 'ux': T2U @ I}, + 'u': {'ux': -(T2U @ Dx)}, + } + + operator = helper.get_empty_operator_matrix() + for line, equation in L_lhs.items(): + helper.add_equation_lhs(operator, line, equation) + + L = helper.convert_operator_matrix_to_operator(operator) + + Args: + A (list of lists of sparse matrices): The operator to be + equation (str): The equation of the component you want this in + relations: (dict): Relations between quantities + """ + for k, v in relations.items(): + A[self.index(equation)][self.index(k)] = v + + def convert_operator_matrix_to_operator(self, M): + """ + Promote the list of lists of sparse matrices to a single sparse matrix that can be used as linear operator. + See documentation of `SpectralHelper.add_equation_lhs` for an example. + + Args: + M (list of lists of sparse matrices): The operator to be + + Returns: + sparse linear operator + """ + if len(self.components) == 1: + return M[0][0] + else: + return self.sparse_lib.bmat(M, format='csc') + + def get_wavenumbers(self): + """ + Get grid in spectral space + """ + grids = [self.axes[i].get_wavenumbers()[self.local_slice[i]] for i in range(len(self.axes))][::-1] + return self.xp.meshgrid(*grids) + + def get_grid(self): + """ + Get grid in physical space + """ + grids = [self.axes[i].get_1dgrid()[self.local_slice[i]] for i in range(len(self.axes))][::-1] + return self.xp.meshgrid(*grids) + + def get_fft(self, axes=None, direction='object', padding=None, shape=None): + """ + When using MPI, we use `PFFT` objects generated by mpi4py-fft + + Args: + axes (tuple): Axes you want to transform over + direction (str): use "forward" or "backward" to get functions for performing the transforms or "object" to get the PFFT object + padding (tuple): Padding for dealiasing + shape (tuple): Shape of the transform + + Returns: + transform + """ + axes = tuple(-i - 1 for i in range(self.ndim)) if axes is None else axes + shape = self.global_shape[1:] if shape is None else shape + padding = ( + [ + 1, + ] + * self.ndim + if padding is None + else padding + ) + key = (axes, direction, tuple(padding), tuple(shape)) + + if key not in self.fft_cache.keys(): + if self.comm is None: + assert np.allclose(padding, 1), 'Zero padding is not implemented for non-MPI transforms' + + if direction == 'forward': + self.fft_cache[key] = self.xp.fft.fftn + elif direction == 'backward': + self.fft_cache[key] = self.xp.fft.ifftn + elif direction == 'object': + self.fft_cache[key] = None + else: + from mpi4py_fft import PFFT + + _fft = PFFT( + comm=self.comm, + shape=shape, + axes=sorted(axes), + dtype='D', + collapse=False, + backend=self.fft_backend, + comm_backend=self.fft_comm_backend, + padding=padding, + ) + if direction == 'forward': + self.fft_cache[key] = _fft.forward + elif direction == 'backward': + self.fft_cache[key] = _fft.backward + elif direction == 'object': + self.fft_cache[key] = _fft + + return self.fft_cache[key] + + def setup_fft(self, real_spectral_coefficients=False): + """ + This function must be called after all axes have been setup in order to prepare the local shapes of the data. + This must also be called before setting up any BCs. + + Args: + real_spectral_coefficients (bool): Allow only real coefficients in spectral space + """ + if len(self.components) == 0: + self.add_component('u') + + self.global_shape = (len(self.components),) + tuple(me.N for me in self.axes) + self.local_slice = [slice(0, me.N) for me in self.axes] + + axes = tuple(i for i in range(len(self.axes))) + self.fft_obj = self.get_fft(axes=axes, direction='object') + if self.fft_obj is not None: + self.local_slice = self.fft_obj.local_slice(False) + + self.init = ( + np.empty(shape=self.global_shape)[ + ( + ..., + *self.local_slice, + ) + ].shape, + self.comm, + np.dtype('float'), + ) + self.init_forward = ( + np.empty(shape=self.global_shape)[ + ( + ..., + *self.local_slice, + ) + ].shape, + self.comm, + np.dtype('float') if real_spectral_coefficients else np.dtype('complex128'), + ) + + self.BC_mat = self.get_empty_operator_matrix() + self.BC_rhs_mask = self.xp.zeros( + shape=self.init[0], + dtype=bool, + ) + + def _transform_fft(self, u, axes, **kwargs): + """ + FFT along `axes` + + Args: + u: The solution + axes (tuple): Axes you want to transform over + + Returns: + transformed solution + """ + # TODO: clean up and try putting more of this in the 1D bases + fft = self.get_fft(axes, 'forward', **kwargs) + return fft(u, axes=axes) + + def _transform_dct(self, u, axes, padding=None, **kwargs): + ''' + DCT along `axes`. + This will only return real values! + When padding the solution, we cannot just use the mpi4py-fft implementation, because of the unusual ordering of + wavenumbers in FFTs. + + Args: + u: The solution + axes (tuple): Axes you want to transform over + + Returns: + transformed solution + ''' + # TODO: clean up and try putting more of this in the 1D bases + if self.debug: + assert self.xp.allclose(u.imag, 0), 'This function can only handle real input.' + + if len(axes) > 1: + v = self._transform_dct(self._transform_dct(u, axes[1:], **kwargs), (axes[0],), **kwargs) + else: + v = u.copy().astype(complex) + axis = axes[0] + base = self.axes[axis] + + shuffle = [slice(0, s, 1) for s in u.shape] + shuffle[axis] = base.get_fft_shuffle(True, N=v.shape[axis]) + v = v[(*shuffle,)] + + if padding is not None: + shape = list(v.shape) + if self.comm: + shape[0] = self.comm.allreduce(v.shape[0]) + fft = self.get_fft(axes, 'forward', shape=shape) + else: + fft = self.get_fft(axes, 'forward', **kwargs) + + v = fft(v, axes=axes) + + expansion = [np.newaxis for _ in u.shape] + expansion[axis] = slice(0, v.shape[axis], 1) + + if padding is not None: + shift = base.get_fft_shift(True, v.shape[axis]) + + if padding[axis] != 1: + N = int(np.ceil(v.shape[axis] / padding[axis])) + _expansion = [slice(0, n) for n in v.shape] + _expansion[axis] = slice(0, N, 1) + v = v[(*_expansion,)] + else: + shift = base.fft_utils['fwd']['shift'] + + v *= shift[(*expansion,)] + + return v.real + + def transform_single_component(self, u, axes=None, padding=None): + """ + Transform a single component of the solution + + Args: + u data to transform: + axes (tuple): Axes over which to transform + padding (list): Padding factor for transform. E.g. a padding factor of 2 will discard the upper half of modes after transforming + + Returns: + Transformed data + """ + # TODO: clean up and try putting more of this in the 1D bases + trfs = { + ChebychevHelper: self._transform_dct, + UltrasphericalHelper: self._transform_dct, + FFTHelper: self._transform_fft, + } + + axes = tuple(-i - 1 for i in range(self.ndim)[::-1]) if axes is None else axes + padding = ( + [ + 1, + ] + * self.ndim + if padding is None + else padding + ) # You know, sometimes I feel very strongly about Black still. This atrocious formatting is readable by Sauron only. + + result = u.copy().astype(complex) + alignment = self.ndim - 1 + + axes_collapsed = [tuple(sorted(me for me in axes if type(self.axes[me]) == base)) for base in trfs.keys()] + bases = [list(trfs.keys())[i] for i in range(len(axes_collapsed)) if len(axes_collapsed[i]) > 0] + axes_collapsed = [me for me in axes_collapsed if len(me) > 0] + shape = [max(u.shape[i], self.global_shape[1 + i]) for i in range(self.ndim)] + + fft = self.get_fft(axes=axes, padding=padding, direction='object') + if fft is not None: + shape = list(fft.global_shape(False)) + + for trf in range(len(axes_collapsed)): + _axes = axes_collapsed[trf] + base = bases[trf] + + if len(_axes) == 0: + continue + + for _ax in _axes: + shape[_ax] = self.global_shape[1 + self.ndim + _ax] + + fft = self.get_fft(_axes, 'object', padding=padding, shape=shape) + + _in = self.get_aligned( + result, axis_in=alignment, axis_out=self.ndim + _axes[-1], forward=False, fft=fft, shape=shape + ) + + alignment = self.ndim + _axes[-1] + + _out = trfs[base](_in, axes=_axes, padding=padding, shape=shape) + + if self.comm is not None: + _out *= np.prod([self.axes[i].N for i in _axes]) + + axes_next_base = axes_collapsed[(trf + 1) % len(axes_collapsed)] + alignment = alignment if len(axes_next_base) == 0 else self.ndim + axes_next_base[-1] + result = self.get_aligned( + _out, axis_in=self.ndim + _axes[0], axis_out=alignment, fft=fft, forward=True, shape=shape + ) + + fft = self.get_fft(axes=axes, padding=padding) + return self.get_aligned(result, axis_in=alignment, axis_out=self.ndim - 1, fft=fft, forward=True, shape=shape) + + def transform(self, u, axes=None, padding=None): + """ + Transform all components from physical space to spectral space + + Args: + u data to transform: + axes (tuple): Axes over which to transform + padding (list): Padding factor for transform. E.g. a padding factor of 2 will discard the upper half of modes after transforming + + Returns: + Transformed data + """ + axes = tuple(-i - 1 for i in range(self.ndim)[::-1]) if axes is None else axes + padding = ( + [ + 1, + ] + * self.ndim + if padding is None + else padding + ) + + result = [ + None, + ] * self.ncomponents + for comp in self.components: + i = self.index(comp) + + result[i] = self.transform_single_component(u[i], axes=axes, padding=padding) + + return self.xp.stack(result) + + def _transform_ifft(self, u, axes, **kwargs): + # TODO: clean up and try putting more of this in the 1D bases + ifft = self.get_fft(axes, 'backward', **kwargs) + return ifft(u, axes=axes) + + def _transform_idct(self, u, axes, padding=None, **kwargs): + ''' + This will only ever return real values! + ''' + # TODO: clean up and try putting more of this in the 1D bases + if self.debug: + assert self.xp.allclose(u.imag, 0), 'This function can only handle real input.' + + v = u.copy().astype(complex) + + if len(axes) > 1: + v = self._transform_idct(self._transform_idct(u, axes[1:]), (axes[0],)) + else: + axis = axes[0] + base = self.axes[axis] + + if padding is not None: + if padding[axis] != 1: + N_pad = int(np.ceil(v.shape[axis] * padding[axis])) + _pad = [[0, 0] for _ in v.shape] + _pad[axis] = [0, N_pad - base.N] + v = self.xp.pad(v, _pad, 'constant') + + shift = self.xp.exp(1j * np.pi * self.xp.arange(N_pad) / (2 * N_pad)) * base.N + else: + shift = base.fft_utils['bck']['shift'] + else: + shift = base.fft_utils['bck']['shift'] + + expansion = [np.newaxis for _ in u.shape] + expansion[axis] = slice(0, v.shape[axis], 1) + + v *= shift[(*expansion,)] + + if padding is not None: + if padding[axis] != 1: + shape = list(v.shape) + if self.comm: + shape[0] = self.comm.allreduce(v.shape[0]) + ifft = self.get_fft(axes, 'backward', shape=shape) + else: + ifft = self.get_fft(axes, 'backward', padding=padding, **kwargs) + else: + ifft = self.get_fft(axes, 'backward', padding=padding, **kwargs) + v = ifft(v, axes=axes) + + shuffle = [slice(0, s, 1) for s in v.shape] + shuffle[axis] = base.get_fft_shuffle(False, N=v.shape[axis]) + v = v[(*shuffle,)] + + return v.real + + def itransform_single_component(self, u, axes=None, padding=None): + """ + Inverse transform over single component of the solution + + Args: + u data to transform: + axes (tuple): Axes over which to transform + padding (list): Padding factor for transform. E.g. a padding factor of 2 will add as many zeros as there were modes before before transforming + + Returns: + Transformed data + """ + # TODO: clean up and try putting more of this in the 1D bases + trfs = { + FFTHelper: self._transform_ifft, + ChebychevHelper: self._transform_idct, + UltrasphericalHelper: self._transform_idct, + } + + axes = tuple(-i - 1 for i in range(self.ndim)[::-1]) if axes is None else axes + padding = ( + [ + 1, + ] + * self.ndim + if padding is None + else padding + ) + + result = u.copy().astype(complex) + alignment = self.ndim - 1 + + axes_collapsed = [tuple(sorted(me for me in axes if type(self.axes[me]) == base)) for base in trfs.keys()] + bases = [list(trfs.keys())[i] for i in range(len(axes_collapsed)) if len(axes_collapsed[i]) > 0] + axes_collapsed = [me for me in axes_collapsed if len(me) > 0] + shape = list(self.global_shape[1:]) + + for trf in range(len(axes_collapsed)): + _axes = axes_collapsed[trf] + base = bases[trf] + + if len(_axes) == 0: + continue + + fft = self.get_fft(_axes, 'object', padding=padding, shape=shape) + + _in = self.get_aligned( + result, axis_in=alignment, axis_out=self.ndim + _axes[0], forward=True, fft=fft, shape=shape + ) + if self.comm is not None: + _in /= np.prod([self.axes[i].N for i in _axes]) + + alignment = self.ndim + _axes[0] + + _out = trfs[base](_in, axes=_axes, padding=padding, shape=shape) + + for _ax in _axes: + if fft: + shape[_ax] = fft._input_shape[_ax] + else: + shape[_ax] = _out.shape[_ax] + + axes_next_base = axes_collapsed[(trf + 1) % len(axes_collapsed)] + alignment = alignment if len(axes_next_base) == 0 else self.ndim + axes_next_base[0] + result = self.get_aligned( + _out, axis_in=self.ndim + _axes[-1], axis_out=alignment, fft=fft, forward=False, shape=shape + ) + + fft = self.get_fft(axes=axes, padding=padding) + return self.get_aligned(result, axis_in=alignment, axis_out=self.ndim - 1, fft=fft, shape=shape) + + def get_aligned(self, u, axis_in, axis_out, fft=None, forward=False, **kwargs): + """ + Realign the data along the axis when using distributed FFTs. `kwargs` will be used to get the correct PFFT + object from `mpi4py-fft`, which has suitable transfer classes for the shape of data. Hence, they should include + shape especially, if applicable. + + Args: + u: The solution + axis_in (int): Current alignment + axis_out (int): New alignment + fft (mpi4py_fft.PFFT), optional: parallel FFT object + forward (bool): Whether the input is in spectral space or not + + Returns: + solution aligned on `axis_in` + """ + if self.comm is None or axis_in == axis_out: + return u.copy() + if self.comm.size == 1: + return u.copy() + + fft = self.get_fft(**kwargs) if fft is None else fft + + global_fft = self.get_fft(**kwargs) + axisA = [me.axisA for me in global_fft.transfer] + axisB = [me.axisB for me in global_fft.transfer] + + current_axis = axis_in + + if axis_in in axisA and axis_out in axisB: + while current_axis != axis_out: + transfer = global_fft.transfer[axisA.index(current_axis)] + + arrayB = self.xp.empty(shape=transfer.subshapeB, dtype=transfer.dtype) + arrayA = self.xp.empty(shape=transfer.subshapeA, dtype=transfer.dtype) + arrayA[:] = u[:] + + transfer.forward(arrayA=arrayA, arrayB=arrayB) + + current_axis = transfer.axisB + u = arrayB + + return u + elif axis_in in axisB and axis_out in axisA: + while current_axis != axis_out: + transfer = global_fft.transfer[axisB.index(current_axis)] + + arrayB = self.xp.empty(shape=transfer.subshapeB, dtype=transfer.dtype) + arrayA = self.xp.empty(shape=transfer.subshapeA, dtype=transfer.dtype) + arrayB[:] = u[:] + + transfer.backward(arrayA=arrayA, arrayB=arrayB) + + current_axis = transfer.axisA + u = arrayA + + return u + else: # go the potentially slower route of not reusing transfer classes + from mpi4py_fft import newDistArray + + _in = newDistArray(fft, forward).redistribute(axis_in) + _in[...] = u + + return _in.redistribute(axis_out) + + def itransform(self, u, axes=None, padding=None): + """ + Inverse transform over all components of the solution + + Args: + u data to transform: + axes (tuple): Axes over which to transform + padding (list): Padding factor for transform. E.g. a padding factor of 2 will add as many zeros as there were modes before before transforming + + Returns: + Transformed data + """ + axes = tuple(-i - 1 for i in range(self.ndim)[::-1]) if axes is None else axes + padding = ( + [ + 1, + ] + * self.ndim + if padding is None + else padding + ) + + result = [ + None, + ] * self.ncomponents + for comp in self.components: + i = self.index(comp) + + result[i] = self.itransform_single_component(u[i], axes=axes, padding=padding) + + return self.xp.stack(result) + + def get_local_slice_of_1D_matrix(self, M, axis): + """ + Get the local version of a 1D matrix. When using distributed FFTs, each rank will carry only a subset of modes, + which you can sort out via the `SpectralHelper.local_slice` attribute. When constructing a 1D matrix, you can + use this method to get the part corresponding to the modes carried by this rank. + + Args: + M (sparse matrix): Global 1D matrix you want to get the local version of + axis (int): Direction in which you want the local version. You will get the global matrix in other directions. This means slab decomposition only. + + Returns: + sparse local matrix + """ + return M.tocsc()[self.local_slice[axis], self.local_slice[axis]] + + def get_filter_matrix(self, axis, **kwargs): + """ + Get bandpass filter along `axis`. See the documentation `get_filter_matrix` in the 1D bases for what kwargs are + admissible. + + Returns: + sparse bandpass matrix + """ + if self.ndim == 1: + return self.axes[0].get_filter_matrix(**kwargs) + + mats = [base.get_Id() for base in self.axes] + mats[axis] = self.axes[axis].get_filter_matrix(**kwargs) + return self.sparse_lib.kron(*mats) + + def get_differentiation_matrix(self, axes, **kwargs): + """ + Get differentiation matrix along specified axis. `kwargs` are forwarded to the 1D base implementation. + + Args: + axes (tuple): Axes along which to differentiate. + + Returns: + sparse differentiation matrix + """ + sp = self.sparse_lib + ndim = self.ndim + + if ndim == 1: + D = self.axes[0].get_differentiation_matrix(**kwargs) + elif ndim == 2: + for axis in axes: + axis2 = (axis + 1) % ndim + D1D = self.axes[axis].get_differentiation_matrix(**kwargs) + + if len(axes) > 1: + I1D = sp.eye(self.axes[axis2].N) + else: + I1D = self.axes[axis2].get_Id() + + mats = [None] * ndim + mats[axis] = self.get_local_slice_of_1D_matrix(D1D, axis) + mats[axis2] = self.get_local_slice_of_1D_matrix(I1D, axis2) + + if axis == axes[0]: + D = sp.kron(*mats) + else: + D = D @ sp.kron(*mats) + else: + raise NotImplementedError(f'Differentiation matrix not implemented for {ndim} dimension!') + + return D + + def get_integration_matrix(self, axes): + """ + Get integration matrix to integrate along specified axis. + + Args: + axes (tuple): Axes along which to integrate over. + + Returns: + sparse integration matrix + """ + sp = self.sparse_lib + ndim = len(self.axes) + + if ndim == 1: + S = self.axes[0].get_integration_matrix() + elif ndim == 2: + for axis in axes: + axis2 = (axis + 1) % ndim + S1D = self.axes[axis].get_integration_matrix() + + if len(axes) > 1: + I1D = sp.eye(self.axes[axis2].N) + else: + I1D = self.axes[axis2].get_Id() + + mats = [None] * ndim + mats[axis] = self.get_local_slice_of_1D_matrix(S1D, axis) + mats[axis2] = self.get_local_slice_of_1D_matrix(I1D, axis2) + + if axis == axes[0]: + S = sp.kron(*mats) + else: + S = S @ sp.kron(*mats) + else: + raise NotImplementedError(f'Integration matrix not implemented for {ndim} dimension!') + + return S + + def get_Id(self): + """ + Get identity matrix + + Returns: + sparse identity matrix + """ + sp = self.sparse_lib + ndim = self.ndim + I = sp.eye(np.prod(self.init[0][1:]), dtype=complex) + + if ndim == 1: + I = self.axes[0].get_Id() + elif ndim == 2: + for axis in range(ndim): + axis2 = (axis + 1) % ndim + I1D = self.axes[axis].get_Id() + + I1D2 = sp.eye(self.axes[axis2].N) + + mats = [None] * ndim + mats[axis] = self.get_local_slice_of_1D_matrix(I1D, axis) + mats[axis2] = self.get_local_slice_of_1D_matrix(I1D2, axis2) + + I = I @ sp.kron(*mats) + else: + raise NotImplementedError(f'Identity matrix not implemented for {ndim} dimension!') + + return I + + def get_Dirichlet_recombination_matrix(self, axis=-1): + """ + Get Dirichlet recombination matrix along axis. Not that it only makes sense in directions discretized with variations of Chebychev bases. + + Args: + axis (int): Axis you discretized with Chebychev + + Returns: + sparse matrix + """ + sp = self.sparse_lib + ndim = len(self.axes) + + if ndim == 1: + C = self.axes[0].get_Dirichlet_recombination_matrix() + elif ndim == 2: + axis2 = (axis + 1) % ndim + C1D = self.axes[axis].get_Dirichlet_recombination_matrix() + + I1D = self.axes[axis2].get_Id() + + mats = [None] * ndim + mats[axis] = self.get_local_slice_of_1D_matrix(C1D, axis) + mats[axis2] = self.get_local_slice_of_1D_matrix(I1D, axis2) + + C = sp.kron(*mats) + else: + raise NotImplementedError(f'Basis change matrix not implemented for {ndim} dimension!') + + return C + + def get_basis_change_matrix(self, axes=None, **kwargs): + """ + Some spectral bases do a change between bases while differentiating. This method returns matrices that changes the basis to whatever you want. + Refer to the methods of the same name of the 1D bases to learn what parameters you need to pass here as `kwargs`. + + Args: + axes (tuple): Axes along which to change basis. + + Returns: + sparse basis change matrix + """ + axes = tuple(-i - 1 for i in range(self.ndim)) if axes is None else axes + + sp = self.sparse_lib + ndim = len(self.axes) + + if ndim == 1: + C = self.axes[0].get_basis_change_matrix(**kwargs) + elif ndim == 2: + for axis in axes: + axis2 = (axis + 1) % ndim + C1D = self.axes[axis].get_basis_change_matrix(**kwargs) + + if len(axes) > 1: + I1D = sp.eye(self.axes[axis2].N) + else: + I1D = self.axes[axis2].get_Id() + + mats = [None] * ndim + mats[axis] = self.get_local_slice_of_1D_matrix(C1D, axis) + mats[axis2] = self.get_local_slice_of_1D_matrix(I1D, axis2) + + if axis == axes[0]: + C = sp.kron(*mats) + else: + C = C @ sp.kron(*mats) + else: + raise NotImplementedError(f'Basis change matrix not implemented for {ndim} dimension!') + + return C diff --git a/pySDC/implementations/problem_classes/Burgers.py b/pySDC/implementations/problem_classes/Burgers.py new file mode 100644 index 0000000000..ef957f54d0 --- /dev/null +++ b/pySDC/implementations/problem_classes/Burgers.py @@ -0,0 +1,337 @@ +import numpy as np + +from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh +from pySDC.implementations.problem_classes.generic_spectral import GenericSpectralLinear + + +class Burgers1D(GenericSpectralLinear): + """ + See https://en.wikipedia.org/wiki/Burgers'_equation for the equation that is solved. + Discretization is done with a Chebychev method, which requires a first order derivative formulation. + Feel free to do a more efficient implementation using an ultraspherical method to avoid the first order business. + + Parameters: + N (int): Spatial resolution + epsilon (float): viscosity + BCl (float): Value at left boundary + BCr (float): Value at right boundary + f (int): Frequency of the initial conditions + mode (str): 'T2U' or 'T2T'. Use 'T2U' to get sparse differentiation matrices + """ + + dtype_u = mesh + dtype_f = imex_mesh + + def __init__(self, N=64, epsilon=0.1, BCl=1, BCr=-1, f=0, mode='T2U', **kwargs): + """ + Constructor. `kwargs` are forwarded to parent class constructor. + + Args: + N (int): Spatial resolution + epsilon (float): viscosity + BCl (float): Value at left boundary + BCr (float): Value at right boundary + f (int): Frequency of the initial conditions + mode (str): 'T2U' or 'T2T'. Use 'T2U' to get sparse differentiation matrices + """ + self._makeAttributeAndRegister('N', 'epsilon', 'BCl', 'BCr', 'f', 'mode', localVars=locals(), readOnly=True) + + bases = [{'base': 'cheby', 'N': N}] + components = ['u', 'ux'] + + super().__init__(bases=bases, components=components, spectral_space=False, **kwargs) + + self.x = self.get_grid()[0] + + # prepare matrices + Dx = self.get_differentiation_matrix(axes=(0,)) + I = self.get_Id() + + T2U = self.get_basis_change_matrix(conv=mode) + + self.Dx = Dx + + # construct linear operator + L_lhs = {'u': {'ux': -epsilon * (T2U @ Dx)}, 'ux': {'u': -T2U @ Dx, 'ux': T2U @ I}} + self.setup_L(L_lhs) + + # construct mass matrix + M_lhs = {'u': {'u': T2U @ I}} + self.setup_M(M_lhs) + + # boundary conditions + self.add_BC(component='u', equation='u', axis=0, x=1, v=BCr, kind='Dirichlet') + self.add_BC(component='u', equation='ux', axis=0, x=-1, v=BCl, kind='Dirichlet') + self.setup_BCs() + + def u_exact(self, t=0, *args, **kwargs): + me = self.u_init + + # x = (self.x + 1) / 2 + # g = 4 * (1 + np.exp(-(4 * x + t)/self.epsilon/32)) + # g_x = 4 * np.exp(-(4 * x + t)/self.epsilon/32) * (-4/self.epsilon/32) + + # me[0] = 3./4. - 1./g + # me[1] = 1/g**2 * g_x + + # return me + + if t == 0: + me[self.index('u')][:] = ((self.BCr + self.BCl) / 2 + (self.BCr - self.BCl) / 2 * self.x) * np.cos( + self.x * np.pi * self.f + ) + me[self.index('ux')][:] = (self.BCr - self.BCl) / 2 * np.cos(self.x * np.pi * self.f) + ( + (self.BCr + self.BCl) / 2 + (self.BCr - self.BCl) / 2 * self.x + ) * self.f * np.pi * -np.sin(self.x * np.pi * self.f) + elif t == np.inf and self.f == 0 and self.BCl == -self.BCr: + me[0] = (self.BCl * np.exp((self.BCr - self.BCl) / (2 * self.epsilon) * self.x) + self.BCr) / ( + np.exp((self.BCr - self.BCl) / (2 * self.epsilon) * self.x) + 1 + ) + else: + raise NotImplementedError + + return me + + def eval_f(self, u, *args, **kwargs): + f = self.f_init + iu, iux = self.index('u'), self.index('ux') + + u_hat = self.transform(u) + + Dx_u_hat = self.u_init_forward + Dx_u_hat[iux] = (self.Dx @ u_hat[iux].flatten()).reshape(u_hat[iu].shape) + + f.impl[iu] = self.epsilon * self.itransform(Dx_u_hat)[iux].real + f.expl[iu] = -u[iu] * u[iux] + return f + + def get_fig(self): # pragma: no cover + """ + Get a figure suitable to plot the solution of this problem. + + Returns + ------- + self.fig : matplotlib.pyplot.figure.Figure + """ + import matplotlib.pyplot as plt + + plt.rcParams['figure.constrained_layout.use'] = True + self.fig, axs = plt.subplots() + return self.fig + + def plot(self, u, t=None, fig=None, comp='u'): # pragma: no cover + r""" + Plot the solution. + + Parameters + ---------- + u : dtype_u + Solution to be plotted + t : float + Time to display at the top of the figure + fig : matplotlib.pyplot.figure.Figure, optional + Figure with the same structure as a figure generated by `self.get_fig`. If none is supplied, a new figure will be generated. + + Returns + ------- + None + """ + fig = self.get_fig() if fig is None else fig + ax = fig.axes[0] + + ax.plot(self.x, u[self.index(comp)]) + + if t is not None: + fig.suptitle(f't = {t:.2e}') + + ax.set_xlabel(r'$x$') + ax.set_ylabel(r'$u$') + + +class Burgers2D(GenericSpectralLinear): + """ + See https://en.wikipedia.org/wiki/Burgers'_equation for the equation that is solved. + This implementation is discretized with FFTs in x and Chebychev in z. + + Parameters: + nx (int): Spatial resolution in x direction + nz (int): Spatial resolution in z direction + epsilon (float): viscosity + BCl (float): Value at left boundary + BCr (float): Value at right boundary + fux (int): Frequency of the initial conditions in x-direction + fuz (int): Frequency of the initial conditions in z-direction + mode (str): 'T2U' or 'T2T'. Use 'T2U' to get sparse differentiation matrices + """ + + dtype_u = mesh + dtype_f = imex_mesh + + def __init__(self, nx=64, nz=64, epsilon=0.1, fux=2, fuz=1, mode='T2U', **kwargs): + """ + Constructor. `kwargs` are forwarded to parent class constructor. + + Args: + nx (int): Spatial resolution in x direction + nz (int): Spatial resolution in z direction + epsilon (float): viscosity + fux (int): Frequency of the initial conditions in x-direction + fuz (int): Frequency of the initial conditions in z-direction + mode (str): 'T2U' or 'T2T'. Use 'T2U' to get sparse differentiation matrices + """ + self._makeAttributeAndRegister('nx', 'nz', 'epsilon', 'fux', 'fuz', 'mode', localVars=locals(), readOnly=True) + + bases = [ + {'base': 'fft', 'N': nx}, + {'base': 'cheby', 'N': nz}, + ] + components = ['u', 'v', 'ux', 'uz', 'vx', 'vz'] + super().__init__(bases=bases, components=components, spectral_space=False, **kwargs) + + self.Z, self.X = self.get_grid() + + # prepare matrices + Dx = self.get_differentiation_matrix(axes=(0,)) + Dz = self.get_differentiation_matrix(axes=(1,)) + I = self.get_Id() + + T2U = self.get_basis_change_matrix(axes=(1,), conv=mode) + + self.Dx = Dx + self.Dz = Dz + + # construct linear operator + L_lhs = { + 'u': {'ux': -epsilon * T2U @ Dx, 'uz': -epsilon * T2U @ Dz}, + 'v': {'vx': -epsilon * T2U @ Dx, 'vz': -epsilon * T2U @ Dz}, + 'ux': {'u': -T2U @ Dx, 'ux': T2U @ I}, + 'uz': {'u': -T2U @ Dz, 'uz': T2U @ I}, + 'vx': {'v': -T2U @ Dx, 'vx': T2U @ I}, + 'vz': {'v': -T2U @ Dz, 'vz': T2U @ I}, + } + self.setup_L(L_lhs) + + # construct mass matrix + M_lhs = { + 'u': {'u': T2U @ I}, + 'v': {'v': T2U @ I}, + } + self.setup_M(M_lhs) + + # boundary conditions + self.BCtop = 1 + self.BCbottom = -self.BCtop + self.BCtopu = 0 + self.add_BC(component='v', equation='v', axis=1, v=self.BCtop, x=1, kind='Dirichlet') + self.add_BC(component='v', equation='vz', axis=1, v=self.BCbottom, x=-1, kind='Dirichlet') + self.add_BC(component='u', equation='uz', axis=1, v=self.BCtopu, x=1, kind='Dirichlet') + self.add_BC(component='u', equation='u', axis=1, v=self.BCtopu, x=-1, kind='Dirichlet') + self.setup_BCs() + + def u_exact(self, t=0, *args, noise_level=0, **kwargs): + me = self.u_init + + iu, iv, iux, iuz, ivx, ivz = self.index(self.components) + if t == 0: + me[iu] = self.xp.cos(self.X * self.fux) * self.xp.sin(self.Z * np.pi * self.fuz) + self.BCtopu + me[iux] = -self.xp.sin(self.X * self.fux) * self.fux * self.xp.sin(self.Z * np.pi * self.fuz) + me[iuz] = self.xp.cos(self.X * self.fux) * self.xp.cos(self.Z * np.pi * self.fuz) * np.pi * self.fuz + + me[iv] = (self.BCtop + self.BCbottom) / 2 + (self.BCtop - self.BCbottom) / 2 * self.Z + me[ivz][:] = (self.BCtop - self.BCbottom) / 2 + + # add noise + rng = self.xp.random.default_rng(seed=99) + me[iv].real += rng.normal(size=me[iv].shape) * (self.Z - 1) * (self.Z + 1) * noise_level + + else: + raise NotImplementedError + + return me + + def eval_f(self, u, *args, **kwargs): + f = self.f_init + iu, iv, iux, iuz, ivx, ivz = self.index(self.components) + + u_hat = self.transform(u) + f_hat = self.u_init_forward + f_hat[iu] = self.epsilon * ((self.Dx @ u_hat[iux].flatten() + self.Dz @ u_hat[iuz].flatten())).reshape( + u_hat[iux].shape + ) + f_hat[iv] = self.epsilon * ((self.Dx @ u_hat[ivx].flatten() + self.Dz @ u_hat[ivz].flatten())).reshape( + u_hat[iux].shape + ) + f.impl[...] = self.itransform(f_hat).real + + f.expl[iu] = -(u[iu] * u[iux] + u[iv] * u[iuz]) + f.expl[iv] = -(u[iu] * u[ivx] + u[iv] * u[ivz]) + return f + + def compute_vorticity(self, u): + me = self.u_init_forward + + u_hat = self.transform(u) + iu, iv = self.index(['u', 'v']) + + me[iu] = (self.Dx @ u_hat[iv].flatten() + self.Dz @ u_hat[iu].flatten()).reshape(u[iu].shape) + return self.itransform(me)[iu].real + + def get_fig(self): # pragma: no cover + """ + Get a figure suitable to plot the solution of this problem + + Returns + ------- + self.fig : matplotlib.pyplot.figure.Figure + """ + import matplotlib.pyplot as plt + from mpl_toolkits.axes_grid1 import make_axes_locatable + + plt.rcParams['figure.constrained_layout.use'] = True + self.fig, axs = plt.subplots(3, 1, sharex=True, sharey=True, figsize=((8, 7))) + self.cax = [] + divider = make_axes_locatable(axs[0]) + self.cax += [divider.append_axes('right', size='3%', pad=0.03)] + divider2 = make_axes_locatable(axs[1]) + self.cax += [divider2.append_axes('right', size='3%', pad=0.03)] + divider3 = make_axes_locatable(axs[2]) + self.cax += [divider3.append_axes('right', size='3%', pad=0.03)] + return self.fig + + def plot(self, u, t=None, fig=None, vmin=None, vmax=None): # pragma: no cover + r""" + Plot the solution. Please supply a figure with the same structure as returned by ``self.get_fig``. + + Parameters + ---------- + u : dtype_u + Solution to be plotted + t : float + Time to display at the top of the figure + fig : matplotlib.pyplot.figure.Figure + Figure with the correct structure + + Returns + ------- + None + """ + fig = self.get_fig() if fig is None else fig + axs = fig.axes + + iu, iv = self.index(['u', 'v']) + + imU = axs[0].pcolormesh(self.X, self.Z, u[iu].real, vmin=vmin, vmax=vmax) + imV = axs[1].pcolormesh(self.X, self.Z, u[iv].real, vmin=vmin, vmax=vmax) + imVort = axs[2].pcolormesh(self.X, self.Z, self.compute_vorticity(u).real) + + for i, label in zip([0, 1, 2], [r'$u$', '$v$', 'vorticity']): + axs[i].set_aspect(1) + axs[i].set_title(label) + + if t is not None: + fig.suptitle(f't = {t:.2e}') + axs[-1].set_xlabel(r'$x$') + axs[-1].set_ylabel(r'$z$') + fig.colorbar(imU, self.cax[0]) + fig.colorbar(imV, self.cax[1]) + fig.colorbar(imVort, self.cax[2]) diff --git a/pySDC/implementations/problem_classes/HeatEquation_Chebychev.py b/pySDC/implementations/problem_classes/HeatEquation_Chebychev.py new file mode 100644 index 0000000000..7a3423ec97 --- /dev/null +++ b/pySDC/implementations/problem_classes/HeatEquation_Chebychev.py @@ -0,0 +1,444 @@ +import numpy as np +from scipy import sparse as sp + +from pySDC.core.problem import Problem +from pySDC.implementations.datatype_classes.mesh import mesh +from pySDC.implementations.problem_classes.generic_spectral import GenericSpectralLinear + + +class Heat1DChebychev(GenericSpectralLinear): + """ + 1D Heat equation with Dirichlet Boundary conditions discretized on (-1, 1) using a Chebychev spectral method. + """ + + dtype_u = mesh + dtype_f = mesh + + def __init__(self, nvars=128, a=0, b=0, f=1, nu=1.0, mode='T2U', **kwargs): + """ + Constructor. `kwargs` are forwarded to parent class constructor. + + Args: + nvars (int): Resolution + a (float): Left BC value + b (float): Right BC value + f (int): Frequency of the solution + nu (float): Diffusion parameter + mode ('T2T' or 'T2U'): Mode for Chebychev method. + + """ + self._makeAttributeAndRegister('nvars', 'a', 'b', 'f', 'nu', 'mode', localVars=locals(), readOnly=True) + + bases = [{'base': 'chebychev', 'N': nvars}] + components = ['u', 'ux'] + + super().__init__(bases, components, real_spectral_coefficients=True, **kwargs) + + self.x = self.get_grid()[0] + + I = self.get_Id() + Dx = self.get_differentiation_matrix(axes=(0,)) + self.Dx = Dx + + self.T2U = self.get_basis_change_matrix(conv=mode) + + L_lhs = { + 'ux': {'u': -self.T2U @ Dx, 'ux': self.T2U @ I}, + 'u': {'ux': -nu * (self.T2U @ Dx)}, + } + self.setup_L(L_lhs) + + M_lhs = {'u': {'u': self.T2U @ I}} + self.setup_M(M_lhs) + + self.add_BC(component='u', equation='u', axis=0, x=-1, v=a, kind="Dirichlet") + self.add_BC(component='u', equation='ux', axis=0, x=1, v=b, kind="Dirichlet") + self.setup_BCs() + + def eval_f(self, u, *args, **kwargs): + f = self.f_init + iu, iux = self.index(self.components) + + if self.spectral_space: + u_hat = u.copy() + else: + u_hat = self.transform(u) + + u_hat[iu] = (self.nu * self.Dx @ u_hat[iux].flatten()).reshape(u_hat[iu].shape) + + if self.spectral_space: + me = u_hat + else: + me = self.itransform(u_hat).real + + f[iu] = me[iu] + return f + + def u_exact(self, t, noise=0, seed=666): + """ + Get exact solution at time `t` + + Args: + t (float): When you want the exact solution + noise (float): Add noise of this level + seed (int): Random seed for the noise + + Returns: + Heat1DChebychev.dtype_u: Exact solution + """ + xp = self.xp + iu, iux = self.index(self.components) + u = self.spectral.u_init + + u[iu] = ( + xp.sin(self.f * np.pi * self.x) * xp.exp(-self.nu * (self.f * np.pi) ** 2 * t) + + (self.b - self.a) / 2 * self.x + + (self.b + self.a) / 2 + ) + u[iux] = ( + self.f * np.pi * xp.cos(self.f * np.pi * self.x) * xp.exp(-self.nu * (self.f * np.pi) ** 2 * t) + + (self.b - self.a) / 2 + ) + + if noise > 0: + assert t == 0 + _noise = self.u_init + rng = self.xp.random.default_rng(seed=seed) + _noise[iu] = rng.normal(size=u[iu].shape) + noise_hat = self.transform(_noise) + low_pass = self.get_filter_matrix(axis=0, kmax=self.nvars - 2) + noise_hat[iu] = (low_pass @ noise_hat[iu].flatten()).reshape(noise_hat[iu].shape) + _noise[:] = self.itransform(noise_hat) + u += _noise * noise * (self.x - 1) * (self.x + 1) + + self.check_BCs(u) + + if self.spectral_space: + u_hat = self.u_init + u_hat[...] = self.transform(u) + return u_hat + else: + return u + + +class Heat1DUltraspherical(GenericSpectralLinear): + """ + 1D Heat equation with Dirichlet Boundary conditions discretized on (-1, 1) using an ultraspherical spectral method. + """ + + dtype_u = mesh + dtype_f = mesh + + def __init__(self, nvars=128, a=0, b=0, f=1, nu=1.0, **kwargs): + """ + Constructor. `kwargs` are forwarded to parent class constructor. + + Args: + nvars (int): Resolution + a (float): Left BC value + b (float): Right BC value + f (int): Frequency of the solution + nu (float): Diffusion parameter + """ + self._makeAttributeAndRegister('nvars', 'a', 'b', 'f', 'nu', localVars=locals(), readOnly=True) + + bases = [{'base': 'ultraspherical', 'N': nvars}] + components = ['u'] + + GenericSpectralLinear.__init__(self, bases, components, real_spectral_coefficients=True, **kwargs) + + self.x = self.get_grid()[0] + + I = self.get_Id() + Dxx = self.get_differentiation_matrix(axes=(0,), p=2) + + S2 = self.get_basis_change_matrix(p_in=2, p_out=0) + U2 = self.get_basis_change_matrix(p_in=0, p_out=2) + + self.Dxx = S2 @ Dxx + + L_lhs = { + 'u': {'u': -nu * Dxx}, + } + self.setup_L(L_lhs) + + M_lhs = {'u': {'u': U2 @ I}} + self.setup_M(M_lhs) + + self.add_BC(component='u', equation='u', axis=0, x=-1, v=a, kind="Dirichlet", line=-1) + self.add_BC(component='u', equation='u', axis=0, x=1, v=b, kind="Dirichlet", line=-2) + self.setup_BCs() + + def eval_f(self, u, *args, **kwargs): + f = self.f_init + iu = self.index('u') + + if self.spectral_space: + u_hat = u.copy() + else: + u_hat = self.transform(u) + + u_hat[iu] = (self.nu * (self.Dxx @ u_hat[iu].flatten())).reshape(u_hat[iu].shape) + + if self.spectral_space: + me = u_hat + else: + me = self.itransform(u_hat).real + + f[iu][...] = me[iu] + return f + + def u_exact(self, t, noise=0, seed=666): + """ + Get exact solution at time `t` + + Args: + t (float): When you want the exact solution + noise (float): Add noise of this level + seed (int): Random seed for the noise + + Returns: + Heat1DUltraspherical.dtype_u: Exact solution + """ + xp = self.xp + iu = self.index('u') + u = self.spectral.u_init + + u[iu] = ( + xp.sin(self.f * np.pi * self.x) * xp.exp(-self.nu * (self.f * np.pi) ** 2 * t) + + (self.b - self.a) / 2 * self.x + + (self.b + self.a) / 2 + ) + + if noise > 0: + assert t == 0 + _noise = self.u_init + rng = self.xp.random.default_rng(seed=seed) + _noise[iu] = rng.normal(size=u[iu].shape) + noise_hat = self.transform(_noise) + low_pass = self.get_filter_matrix(axis=0, kmax=self.nvars - 2) + noise_hat[iu] = (low_pass @ noise_hat[iu].flatten()).reshape(noise_hat[iu].shape) + _noise[:] = self.itransform(noise_hat) + u += _noise * noise * (self.x - 1) * (self.x + 1) + + self.check_BCs(u) + + if self.spectral_space: + return self.transform(u) + else: + return u + + +class Heat2DChebychev(GenericSpectralLinear): + """ + 2D Heat equation with Dirichlet Boundary conditions discretized on (-1, 1)x(-1,1) using spectral methods based on FFT and Chebychev. + """ + + dtype_u = mesh + dtype_f = mesh + + def __init__(self, nx=128, ny=128, base_x='fft', base_y='chebychev', a=0, b=0, c=0, fx=1, fy=1, nu=1.0, **kwargs): + """ + Constructor. `kwargs` are forwarded to parent class constructor. + + Args: + nx (int): Resolution in x-direction + ny (int): Resolution in y-direction + base_x (str): Spectral base in x-direction + base_y (str): Spectral base in y-direction + a (float): BC at y=0 and x=-1 + b (float): BC at y=0 and x=+1 + c (float): BC at y=1 and x = -1 + fx (int): Horizontal frequency of initial conditions + fy (int): Vertical frequency of initial conditions + nu (float): Diffusion parameter + """ + assert nx % 2 == 1 or base_x == 'fft' + assert ny % 2 == 1 or base_y == 'fft' + self._makeAttributeAndRegister( + 'nx', 'ny', 'base_x', 'base_y', 'a', 'b', 'c', 'fx', 'fy', 'nu', localVars=locals(), readOnly=True + ) + + bases = [{'base': base_x, 'N': nx}, {'base': base_y, 'N': ny}] + components = ['u', 'ux', 'uy'] + + super().__init__(bases, components, Dirichlet_recombination=False, spectral_space=False, **kwargs) + + self.Y, self.X = self.get_grid() + + I = self.get_Id() + self.Dx = self.get_differentiation_matrix(axes=(0,)) + self.Dy = self.get_differentiation_matrix(axes=(1,)) + + L_lhs = { + 'ux': {'u': -self.Dx, 'ux': I}, + 'uy': {'u': -self.Dy, 'uy': I}, + 'u': {'ux': -nu * self.Dx, 'uy': -nu * self.Dy}, + } + self.setup_L(L_lhs) + + M_lhs = {'u': {'u': I}} + self.setup_M(M_lhs) + + for base in [base_x, base_y]: + assert base in ['chebychev', 'fft'] + + alpha = (self.b - self.a) / 2.0 + beta = (self.c - self.b) / 2.0 + gamma = (self.c + self.a) / 2.0 + + if base_x == 'chebychev': + y = self.Y[0, :] + if self.base_y == 'fft': + self.add_BC(component='u', equation='u', axis=0, x=-1, v=beta * y - alpha + gamma, kind='Dirichlet') + self.add_BC(component='ux', equation='ux', axis=0, v=2 * alpha, kind='integral') + else: + assert a == b, f'Need periodic boundary conditions in x for {base_x} method!' + if base_y == 'chebychev': + x = self.X[:, 0] + self.add_BC(component='u', equation='u', axis=1, x=-1, v=alpha * x - beta + gamma, kind='Dirichlet') + self.add_BC(component='uy', equation='uy', axis=1, v=2 * beta, kind='integral') + else: + assert c == b, f'Need periodic boundary conditions in y for {base_y} method!' + self.setup_BCs() + + def eval_f(self, u, *args, **kwargs): + f = self.f_init + iu, iux, iuy = self.index(self.components) + + me_hat = self.u_init_forward + me_hat[:] = self.transform(u) + me_hat[iu] = self.nu * (self.Dx @ me_hat[iux].flatten() + self.Dy @ me_hat[iuy].flatten()).reshape( + me_hat[iu].shape + ) + me = self.itransform(me_hat) + + f[self.index("u")] = me[iu].real + return f + + def u_exact(self, t): + xp = self.xp + iu, iux, iuy = self.index(self.components) + u = self.u_init + + fx = self.fx if self.base_x == 'fft' else np.pi * self.fx + fy = self.fy if self.base_y == 'fft' else np.pi * self.fy + + time_dep = xp.exp(-self.nu * (fx**2 + fy**2) * t) + + alpha = (self.b - self.a) / 2.0 + beta = (self.c - self.b) / 2.0 + gamma = (self.c + self.a) / 2.0 + + u[iu] = xp.sin(fx * self.X) * xp.sin(fy * self.Y) * time_dep + alpha * self.X + beta * self.Y + gamma + u[iux] = fx * xp.cos(fx * self.X) * xp.sin(fy * self.Y) * time_dep + alpha + u[iuy] = fy * xp.sin(fx * self.X) * xp.cos(fy * self.Y) * time_dep + beta + + return u + + +class Heat2DUltraspherical(GenericSpectralLinear): + """ + 2D Heat equation with Dirichlet Boundary conditions discretized on (-1, 1)x(-1,1) using spectral methods based on FFT and Gegenbauer. + """ + + dtype_u = mesh + dtype_f = mesh + + def __init__( + self, nx=128, ny=128, base_x='fft', base_y='ultraspherical', a=0, b=0, c=0, fx=1, fy=1, nu=1.0, **kwargs + ): + """ + Constructor. `kwargs` are forwarded to parent class constructor. + + Args: + nx (int): Resolution in x-direction + ny (int): Resolution in y-direction + base_x (str): Spectral base in x-direction + base_y (str): Spectral base in y-direction + a (float): BC at y=0 and x=-1 + b (float): BC at y=0 and x=+1 + c (float): BC at y=1 and x = -1 + fx (int): Horizontal frequency of initial conditions + fy (int): Vertical frequency of initial conditions + nu (float): Diffusion parameter + """ + self._makeAttributeAndRegister( + 'nx', 'ny', 'base_x', 'base_y', 'a', 'b', 'c', 'fx', 'fy', 'nu', localVars=locals(), readOnly=True + ) + + bases = [{'base': base_x, 'N': nx}, {'base': base_y, 'N': ny}] + components = ['u'] + + super().__init__(bases, components, Dirichlet_recombination=False, spectral_space=False, **kwargs) + + self.Y, self.X = self.get_grid() + + self.Dxx = self.get_differentiation_matrix(axes=(0,), p=2) + self.Dyy = self.get_differentiation_matrix(axes=(1,), p=2) + self.S2 = self.get_basis_change_matrix(p=2) + I = self.get_Id() + + L_lhs = { + 'u': {'u': -nu * self.Dxx - nu * self.Dyy}, + } + self.setup_L(L_lhs) + + M_lhs = {'u': {'u': I}} + self.setup_M(M_lhs) + + for base in [base_x, base_y]: + assert base in ['ultraspherical', 'fft'] + + alpha = (self.b - self.a) / 2.0 + beta = (self.c - self.b) / 2.0 + gamma = (self.c + self.a) / 2.0 + + if base_x == 'ultraspherical': + y = self.Y[0, :] + if self.base_y == 'fft': + self.add_BC(component='u', equation='u', axis=0, x=-1, v=beta * y - alpha + gamma, kind='Dirichlet') + self.add_BC(component='u', equation='u', axis=0, v=beta * y + alpha + gamma, x=1, line=-2, kind='Dirichlet') + else: + assert a == b, f'Need periodic boundary conditions in x for {base_x} method!' + if base_y == 'ultraspherical': + x = self.X[:, 0] + self.add_BC( + component='u', equation='u', axis=1, x=-1, v=alpha * x - beta + gamma, kind='Dirichlet', line=-1 + ) + self.add_BC( + component='u', equation='u', axis=1, x=+1, v=alpha * x + beta + gamma, kind='Dirichlet', line=-2 + ) + else: + assert c == b, f'Need periodic boundary conditions in y for {base_y} method!' + self.setup_BCs() + + def eval_f(self, u, *args, **kwargs): + f = self.f_init + iu = self.index('u') + + me_hat = self.u_init_forward + me_hat[:] = self.transform(u) + me_hat[iu] = self.nu * (self.S2 @ (self.Dxx + self.Dyy) @ me_hat[iu].flatten()).reshape(me_hat[iu].shape) + me = self.itransform(me_hat) + + f[iu] = me[iu].real + return f + + def u_exact(self, t): + xp = self.xp + iu = self.index('u') + u = self.u_init + + fx = self.fx if self.base_x == 'fft' else np.pi * self.fx + fy = self.fy if self.base_y == 'fft' else np.pi * self.fy + + time_dep = xp.exp(-self.nu * (fx**2 + fy**2) * t) + + alpha = (self.b - self.a) / 2.0 + beta = (self.c - self.b) / 2.0 + gamma = (self.c + self.a) / 2.0 + + u[iu] = xp.sin(fx * self.X) * xp.sin(fy * self.Y) * time_dep + alpha * self.X + beta * self.Y + gamma + + return u diff --git a/pySDC/implementations/problem_classes/RayleighBenard.py b/pySDC/implementations/problem_classes/RayleighBenard.py new file mode 100644 index 0000000000..3a16f2ad06 --- /dev/null +++ b/pySDC/implementations/problem_classes/RayleighBenard.py @@ -0,0 +1,560 @@ +import numpy as np +from mpi4py import MPI + +from pySDC.implementations.problem_classes.generic_spectral import GenericSpectralLinear +from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh +from pySDC.core.convergence_controller import ConvergenceController +from pySDC.core.hooks import Hooks +from pySDC.implementations.convergence_controller_classes.check_convergence import CheckConvergence + + +class RayleighBenard(GenericSpectralLinear): + """ + Rayleigh-Benard Convection is a variation of incompressible Navier-Stokes. + + The equations we solve are + + u_x + v_z = 0 + T_t - kappa (T_xx + T_zz) = -uT_x - vT_z + u_t - nu (u_xx + u_zz) + p_x = -uu_x - vu_z + v_t - nu (v_xx + v_zz) + p_z - T = -uv_x - vv_z + + with u the horizontal velocity, v the vertical velocity (in z-direction), T the temperature, p the pressure, indices + denoting derivatives, kappa=(Rayleigh * Prandl)**(-1/2) and nu = (Rayleigh / Prandl)**(-1/2). Everything on the left + hand side, that is the viscous part, the pressure gradient and the buoyancy due to temperature are treated + implicitly, while the non-linear convection part on the right hand side is integrated explicitly. + + The domain, vertical boundary conditions and pressure gauge are + + Omega = [0, 8) x (-1, 1) + T(z=+1) = 0 + T(z=-1) = 2 + u(z=+-1) = v(z=+-1) = 0 + integral over p = 0 + + The spectral discretization uses FFT horizontally, implying periodic BCs, and an ultraspherical method vertically to + facilitate the Dirichlet BCs. + + Parameters: + Prandl (float): Prandl number + Rayleigh (float): Rayleigh number + nx (int): Horizontal resolution + nz (int): Vertical resolution + BCs (dict): Can specify boundary conditions here + dealiasing (float): Dealiasing factor for evaluating the non-linear part + comm (mpi4py.Intracomm): Space communicator + """ + + dtype_u = mesh + dtype_f = imex_mesh + + def __init__( + self, + Prandl=1, + Rayleigh=2e6, + nx=256, + nz=64, + BCs=None, + dealiasing=3 / 2, + comm=None, + **kwargs, + ): + """ + Constructor. `kwargs` are forwarded to parent class constructor. + + Args: + Prandl (float): Prandtl number + Rayleigh (float): Rayleigh number + nx (int): Resolution in x-direction + nz (int): Resolution in z direction + BCs (dict): Vertical boundary conditions + dealiasing (float): Dealiasing for evaluating the non-linear part in real space + comm (mpi4py.Intracomm): Space communicator + """ + BCs = {} if BCs is None else BCs + BCs = { + 'T_top': 0, + 'T_bottom': 2, + 'v_top': 0, + 'v_bottom': 0, + 'u_top': 0, + 'u_bottom': 0, + 'p_integral': 0, + **BCs, + } + if comm is None: + try: + from mpi4py import MPI + + comm = MPI.COMM_WORLD + except ModuleNotFoundError: + pass + self._makeAttributeAndRegister( + 'Prandl', + 'Rayleigh', + 'nx', + 'nz', + 'BCs', + 'dealiasing', + 'comm', + localVars=locals(), + readOnly=True, + ) + + bases = [{'base': 'fft', 'N': nx, 'x0': 0, 'x1': 8}, {'base': 'ultraspherical', 'N': nz}] + components = ['u', 'v', 'T', 'p'] + super().__init__(bases, components, comm=comm, **kwargs) + + self.Z, self.X = self.get_grid() + self.Kz, self.Kx = self.get_wavenumbers() + + # construct 2D matrices + Dzz = self.get_differentiation_matrix(axes=(1,), p=2) + Dz = self.get_differentiation_matrix(axes=(1,)) + Dx = self.get_differentiation_matrix(axes=(0,)) + Dxx = self.get_differentiation_matrix(axes=(0,), p=2) + Id = self.get_Id() + + S1 = self.get_basis_change_matrix(p_out=0, p_in=1) + S2 = self.get_basis_change_matrix(p_out=0, p_in=2) + + U01 = self.get_basis_change_matrix(p_in=0, p_out=1) + U12 = self.get_basis_change_matrix(p_in=1, p_out=2) + U02 = self.get_basis_change_matrix(p_in=0, p_out=2) + + self.Dx = Dx + self.Dxx = Dxx + self.Dz = S1 @ Dz + self.Dzz = S2 @ Dzz + + kappa = (Rayleigh * Prandl) ** (-1 / 2.0) + nu = (Rayleigh / Prandl) ** (-1 / 2.0) + + # construct operators + L_lhs = { + 'p': {'u': U01 @ Dx, 'v': Dz}, # divergence free constraint + 'u': {'p': U02 @ Dx, 'u': -nu * (U02 @ Dxx + Dzz)}, + 'v': {'p': U12 @ Dz, 'v': -nu * (U02 @ Dxx + Dzz), 'T': -U02 @ Id}, + 'T': {'T': -kappa * (U02 @ Dxx + Dzz)}, + } + self.setup_L(L_lhs) + + # mass matrix + M_lhs = {i: {i: U02 @ Id} for i in ['u', 'v', 'T']} + self.setup_M(M_lhs) + + # Prepare going from second (first for divergence free equation) derivative basis back to Chebychov-T + self.base_change = self._setup_operator({**{comp: {comp: S2} for comp in ['u', 'v', 'T']}, 'p': {'p': S1}}) + + # BCs + self.add_BC( + component='p', equation='p', axis=1, v=self.BCs['p_integral'], kind='integral', line=-1, scalar=True + ) + self.add_BC(component='T', equation='T', axis=1, x=-1, v=self.BCs['T_bottom'], kind='Dirichlet', line=-1) + self.add_BC(component='T', equation='T', axis=1, x=1, v=self.BCs['T_top'], kind='Dirichlet', line=-2) + self.add_BC(component='v', equation='v', axis=1, x=1, v=self.BCs['v_bottom'], kind='Dirichlet', line=-1) + self.add_BC(component='v', equation='v', axis=1, x=-1, v=self.BCs['v_bottom'], kind='Dirichlet', line=-2) + self.remove_BC(component='v', equation='v', axis=1, x=-1, kind='Dirichlet', line=-2, scalar=True) + self.add_BC(component='u', equation='u', axis=1, v=self.BCs['u_top'], x=1, kind='Dirichlet', line=-2) + self.add_BC( + component='u', + equation='u', + axis=1, + v=self.BCs['u_bottom'], + x=-1, + kind='Dirichlet', + line=-1, + ) + + # eliminate Nyquist mode if needed + if nx % 2 == 0: + Nyquist_mode_index = self.axes[0].get_Nyquist_mode_index() + for component in self.components: + self.add_BC( + component=component, equation=component, axis=0, kind='Nyquist', line=int(Nyquist_mode_index), v=0 + ) + self.setup_BCs() + + def eval_f(self, u, *args, **kwargs): + f = self.f_init + + if self.spectral_space: + u_hat = u.copy() + else: + u_hat = self.transform(u) + + f_impl_hat = self.u_init_forward + + Dz = self.Dz + Dx = self.Dx + + iu, iv, iT, ip = self.index(['u', 'v', 'T', 'p']) + + # evaluate implicit terms + f_impl_hat = -(self.base_change @ self.L @ u_hat.flatten()).reshape(u_hat.shape) + + if self.spectral_space: + f.impl[:] = f_impl_hat + else: + f.impl[:] = self.itransform(f_impl_hat).real + + # treat convection explicitly with dealiasing + Dx_u_hat = self.u_init_forward + for i in [iu, iv, iT]: + Dx_u_hat[i][:] = (Dx @ u_hat[i].flatten()).reshape(Dx_u_hat[i].shape) + Dz_u_hat = self.u_init_forward + for i in [iu, iv, iT]: + Dz_u_hat[i][:] = (Dz @ u_hat[i].flatten()).reshape(Dz_u_hat[i].shape) + + padding = [self.dealiasing, self.dealiasing] + Dx_u_pad = self.itransform(Dx_u_hat, padding=padding).real + Dz_u_pad = self.itransform(Dz_u_hat, padding=padding).real + u_pad = self.itransform(u_hat, padding=padding).real + + fexpl_pad = self.xp.zeros_like(u_pad) + fexpl_pad[iu][:] = -(u_pad[iu] * Dx_u_pad[iu] + u_pad[iv] * Dz_u_pad[iu]) + fexpl_pad[iv][:] = -(u_pad[iu] * Dx_u_pad[iv] + u_pad[iv] * Dz_u_pad[iv]) + fexpl_pad[iT][:] = -(u_pad[iu] * Dx_u_pad[iT] + u_pad[iv] * Dz_u_pad[iT]) + + if self.spectral_space: + f.expl[:] = self.transform(fexpl_pad, padding=padding) + else: + f.expl[:] = self.itransform(self.transform(fexpl_pad, padding=padding)).real + + return f + + def u_exact(self, t=0, noise_level=1e-3, seed=99): + assert t == 0 + assert ( + self.BCs['v_top'] == self.BCs['v_bottom'] + ), 'Initial conditions are only implemented for zero velocity gradient' + + me = self.spectral.u_init + iu, iv, iT, ip = self.index(['u', 'v', 'T', 'p']) + + # linear temperature gradient + for comp in ['T', 'v', 'u']: + a = (self.BCs[f'{comp}_top'] - self.BCs[f'{comp}_bottom']) / 2 + b = (self.BCs[f'{comp}_top'] + self.BCs[f'{comp}_bottom']) / 2 + me[self.index(comp)] = a * self.Z + b + + # perturb slightly + rng = self.xp.random.default_rng(seed=seed) + + noise = self.spectral.u_init + noise[iT] = rng.random(size=me[iT].shape) + + me[iT] += noise[iT].real * noise_level * (self.Z - 1) * (self.Z + 1) + + if self.spectral_space: + me_hat = self.spectral.u_init_forward + me_hat[:] = self.transform(me) + return me_hat + else: + return me + + def get_fig(self): # pragma: no cover + """ + Get a figure suitable to plot the solution of this problem + + Returns + ------- + self.fig : matplotlib.pyplot.figure.Figure + """ + import matplotlib.pyplot as plt + from mpl_toolkits.axes_grid1 import make_axes_locatable + + plt.rcParams['figure.constrained_layout.use'] = True + self.fig, axs = plt.subplots(2, 1, sharex=True, sharey=True, figsize=((10, 5))) + self.cax = [] + divider = make_axes_locatable(axs[0]) + self.cax += [divider.append_axes('right', size='3%', pad=0.03)] + divider2 = make_axes_locatable(axs[1]) + self.cax += [divider2.append_axes('right', size='3%', pad=0.03)] + return self.fig + + def plot(self, u, t=None, fig=None, quantity='T'): # pragma: no cover + r""" + Plot the solution. + + Parameters + ---------- + u : dtype_u + Solution to be plotted + t : float + Time to display at the top of the figure + fig : matplotlib.pyplot.figure.Figure + Figure with the same structure as a figure generated by `self.get_fig`. If none is supplied, a new figure will be generated. + quantity : (str) + quantity you want to plot + + Returns + ------- + None + """ + fig = self.get_fig() if fig is None else fig + axs = fig.axes + + imV = axs[1].pcolormesh(self.X, self.Z, self.compute_vorticity(u).real) + + if self.spectral_space: + u = self.itransform(u) + + imT = axs[0].pcolormesh(self.X, self.Z, u[self.index(quantity)].real) + + for i, label in zip([0, 1], [rf'${quantity}$', 'vorticity']): + axs[i].set_aspect(1) + axs[i].set_title(label) + + if t is not None: + fig.suptitle(f't = {t:.2f}') + axs[1].set_xlabel(r'$x$') + axs[1].set_ylabel(r'$z$') + fig.colorbar(imT, self.cax[0]) + fig.colorbar(imV, self.cax[1]) + + def compute_vorticity(self, u): + if self.spectral_space: + u_hat = u.copy() + else: + u_hat = self.transform(u) + Dz = self.Dz + Dx = self.Dx + iu, iv = self.index(['u', 'v']) + + vorticity_hat = self.spectral.u_init_forward + vorticity_hat[0] = (Dx * u_hat[iv].flatten() + Dz @ u_hat[iu].flatten()).reshape(u[iu].shape) + return self.itransform(vorticity_hat)[0].real + + def compute_Nusselt_numbers(self, u): + """ + Compute the various versions of the Nusselt number. This reflects the type of heat transport. + If the Nusselt number is equal to one, it indicates heat transport due to conduction. If it is larger, + advection is present. + Computing the Nusselt number at various places can be used to check the code. + + Args: + u: The solution you want to compute the Nusselt numbers of + + Returns: + dict: Nusselt number averaged over the entire volume and horizontally averaged at the top and bottom. + """ + iv, iT = self.index(['v', 'T']) + + DzT_hat = self.spectral.u_init_forward + + if self.spectral_space: + u_hat = u.copy() + else: + u_hat = self.transform(u) + + DzT_hat[iT] = (self.Dz @ u_hat[iT].flatten()).reshape(DzT_hat[iT].shape) + + # compute vT with dealiasing + padding = [self.dealiasing, self.dealiasing] + u_pad = self.itransform(u_hat, padding=padding).real + _me = self.xp.zeros_like(u_pad) + _me[0] = u_pad[iv] * u_pad[iT] + vT_hat = self.transform(_me, padding=padding) + + nusselt_hat = (vT_hat[0] - DzT_hat[iT]) / self.nx + nusselt_no_v_hat = (-DzT_hat[iT]) / self.nx + + integral_z = self.xp.sum(nusselt_hat * self.spectral.axes[1].get_BC(kind='integral'), axis=-1).real + integral_V = ( + integral_z[0] * self.axes[0].L + ) # only the first Fourier mode has non-zero integral with periodic BCs + Nusselt_V = self.comm.bcast(integral_V / self.spectral.V, root=0) + + Nusselt_t = self.comm.bcast( + self.xp.sum(nusselt_hat * self.spectral.axes[1].get_BC(kind='Dirichlet', x=1), axis=-1).real[0], root=0 + ) + Nusselt_b = self.comm.bcast( + self.xp.sum(nusselt_hat * self.spectral.axes[1].get_BC(kind='Dirichlet', x=-1), axis=-1).real[0], root=0 + ) + Nusselt_no_v_t = self.comm.bcast( + self.xp.sum(nusselt_no_v_hat * self.spectral.axes[1].get_BC(kind='Dirichlet', x=1), axis=-1).real[0], root=0 + ) + Nusselt_no_v_b = self.comm.bcast( + self.xp.sum(nusselt_no_v_hat * self.spectral.axes[1].get_BC(kind='Dirichlet', x=-1), axis=-1).real[0], + root=0, + ) + + return { + 'V': Nusselt_V, + 't': Nusselt_t, + 'b': Nusselt_b, + 't_no_v': Nusselt_no_v_t, + 'b_no_v': Nusselt_no_v_b, + } + + def compute_viscous_dissipation(self, u): + iu, iv = self.index(['u', 'v']) + + Lap_u_hat = self.spectral.u_init_forward + + if self.spectral_space: + u_hat = u.copy() + else: + u_hat = self.transform(u) + Lap_u_hat[iu] = ((self.Dzz + self.Dxx) @ u_hat[iu].flatten()).reshape(u_hat[iu].shape) + Lap_u_hat[iv] = ((self.Dzz + self.Dxx) @ u_hat[iv].flatten()).reshape(u_hat[iu].shape) + Lap_u = self.itransform(Lap_u_hat) + + return abs(u[iu] * Lap_u[iu] + u[iv] * Lap_u[iv]) + + def compute_buoyancy_generation(self, u): + if self.spectral_space: + u = self.itransform(u) + iv, iT = self.index(['v', 'T']) + return abs(u[iv] * self.Rayleigh * u[iT]) + + +class CFLLimit(ConvergenceController): + + def dependencies(self, controller, *args, **kwargs): + from pySDC.implementations.hooks.log_step_size import LogStepSize + + controller.add_hook(LogCFL) + controller.add_hook(LogStepSize) + + def setup_status_variables(self, controller, **kwargs): + """ + Add the embedded error variable to the error function. + + Args: + controller (pySDC.Controller): The controller + """ + self.add_status_variable_to_level('CFL_limit') + + def setup(self, controller, params, description, **kwargs): + """ + Define default parameters here. + + Default parameters are: + - control_order (int): The order relative to other convergence controllers + - dt_max (float): maximal step size + - dt_min (float): minimal step size + + Args: + controller (pySDC.Controller): The controller + params (dict): The params passed for this specific convergence controller + description (dict): The description object used to instantiate the controller + + Returns: + (dict): The updated params dictionary + """ + defaults = { + "control_order": -50, + "dt_max": np.inf, + "dt_min": 0, + "cfl": 0.4, + } + return {**defaults, **super().setup(controller, params, description, **kwargs)} + + @staticmethod + def compute_max_step_size(P, u): + grid_spacing_x = P.X[1, 0] - P.X[0, 0] + + cell_wallz = P.xp.zeros(P.nz + 1) + cell_wallz[0] = 1 + cell_wallz[-1] = -1 + cell_wallz[1:-1] = (P.Z[0, :-1] + P.Z[0, 1:]) / 2 + grid_spacing_z = cell_wallz[:-1] - cell_wallz[1:] + + iu, iv = P.index(['u', 'v']) + + if P.spectral_space: + u = P.itransform(u) + + max_step_size_x = P.xp.min(grid_spacing_x / P.xp.abs(u[iu])) + max_step_size_z = P.xp.min(grid_spacing_z / P.xp.abs(u[iv])) + max_step_size = min([max_step_size_x, max_step_size_z]) + + if hasattr(P, 'comm'): + max_step_size = P.comm.allreduce(max_step_size, op=MPI.MIN) + return float(max_step_size) + + def get_new_step_size(self, controller, step, **kwargs): + if not CheckConvergence.check_convergence(step): + return None + + L = step.levels[0] + P = step.levels[0].prob + + L.sweep.compute_end_point() + max_step_size = self.compute_max_step_size(P, L.uend) + + L.status.CFL_limit = self.params.cfl * max_step_size + + dt_new = L.status.dt_new if L.status.dt_new else max([self.params.dt_max, L.params.dt]) + L.status.dt_new = min([dt_new, self.params.cfl * max_step_size]) + L.status.dt_new = max([self.params.dt_min, L.status.dt_new]) + + self.log(f'dt max: {max_step_size:.2e} -> New step size: {L.status.dt_new:.2e}', step) + + +class LogCFL(Hooks): + + def post_step(self, step, level_number): + """ + Record CFL limit. + + Args: + step (pySDC.Step.step): the current step + level_number (int): the current level number + + Returns: + None + """ + super().post_step(step, level_number) + + L = step.levels[level_number] + + self.add_to_stats( + process=step.status.slot, + time=L.time + L.dt, + level=L.level_index, + iter=step.status.iter, + sweep=L.status.sweep, + type='CFL_limit', + value=L.status.CFL_limit, + ) + + +class LogAnalysisVariables(Hooks): + + def post_step(self, step, level_number): + """ + Record Nusselt numbers. + + Args: + step (pySDC.Step.step): the current step + level_number (int): the current level number + + Returns: + None + """ + super().post_step(step, level_number) + + L = step.levels[level_number] + P = L.prob + + L.sweep.compute_end_point() + Nusselt = P.compute_Nusselt_numbers(L.uend) + buoyancy_production = P.compute_buoyancy_generation(L.uend) + viscous_dissipation = P.compute_viscous_dissipation(L.uend) + + for key, value in zip( + ['Nusselt', 'buoyancy_production', 'viscous_dissipation'], + [Nusselt, buoyancy_production, viscous_dissipation], + ): + self.add_to_stats( + process=step.status.slot, + time=L.time + L.dt, + level=L.level_index, + iter=step.status.iter, + sweep=L.status.sweep, + type=key, + value=value, + ) diff --git a/pySDC/implementations/problem_classes/generic_spectral.py b/pySDC/implementations/problem_classes/generic_spectral.py new file mode 100644 index 0000000000..04f88823fc --- /dev/null +++ b/pySDC/implementations/problem_classes/generic_spectral.py @@ -0,0 +1,389 @@ +from pySDC.core.problem import Problem, WorkCounter +from pySDC.helpers.spectral_helper import SpectralHelper +import numpy as np +from pySDC.core.errors import ParameterError + + +class GenericSpectralLinear(Problem): + """ + Generic class to solve problems of the form M u_t + L u = y, with mass matrix M, linear operator L and some right + hand side y using spectral methods. + L may contain algebraic conditions, as long as (M + dt L) is invertible. + + Note that the `__getattr__` method is overloaded to pass requests on to the spectral helper if they are not + attributes of this class itself. For instance, you can add a BC by calling `self.spectral.add_BC` or equivalently + `self.add_BC`. + + You can port problems derived from this more or less seamlessly to GPU by using the numerical libraries that are + class attributes of the spectral helper. This class will automatically switch the datatype using the `setup_GPU` class method. + + Attributes: + spectral (pySDC.helpers.spectral_helper.SpectralHelper): Spectral helper + work_counters (dict): Dictionary for counting work + cached_factorizations (dict): Dictionary of cached matrix factorizations for solving + L (sparse matrix): Linear operator + M (sparse matrix): Mass matrix + diff_mask (list): Mask for separating differential and algebraic terms + Pl (sparse matrix): Left preconditioner + Pr (sparse matrix): Right preconditioner + """ + + @classmethod + def setup_GPU(cls): + """switch to GPU modules""" + import cupy as cp + from pySDC.implementations.datatype_classes.cupy_mesh import cupy_mesh, imex_cupy_mesh + from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh + + cls.dtype_u = cupy_mesh + + GPU_versions = { + mesh: cupy_mesh, + imex_mesh: imex_cupy_mesh, + } + + cls.dtype_f = GPU_versions[cls.dtype_f] + + def __init__( + self, + bases, + components, + comm=None, + Dirichlet_recombination=True, + left_preconditioner=True, + solver_type='cached_direct', + solver_args=None, + useGPU=False, + max_cached_factorizations=12, + spectral_space=True, + real_spectral_coefficients=False, + debug=False, + ): + """ + Base class for problems discretized with spectral methods. + + Args: + bases (list of dictionaries): 1D Bases + components (list of strings): Components of the equations + comm (mpi4py.Intracomm or None): MPI communicator + Dirichlet_recombination (bool): Use Dirichlet recombination in the last axis as right preconditioner + left_preconditioner (bool): Reverse the Kronecker product if yes + solver_type (str): Solver for linear systems + solver_args (dict): Arguments for linear solver + useGPU (bool): Run on GPU or CPU + max_cached_factorizations (int): Number of matrix decompositions to cache before starting eviction + spectral_space (bool): If yes, the solution will not be transformed back after solving and evaluating the RHS, and is expected as input in spectral space to these functions + real_spectral_coefficients (bool): If yes, allow only real values in spectral space, otherwise, allow complex. + debug (bool): Make additional tests at extra computational cost + """ + solver_args = {} if solver_args is None else solver_args + self._makeAttributeAndRegister( + 'max_cached_factorizations', + 'useGPU', + 'solver_type', + 'solver_args', + 'left_preconditioner', + 'Dirichlet_recombination', + 'comm', + 'spectral_space', + 'real_spectral_coefficients', + 'debug', + localVars=locals(), + ) + self.spectral = SpectralHelper(comm=comm, useGPU=useGPU, debug=debug) + + if useGPU: + self.setup_GPU() + + for base in bases: + self.spectral.add_axis(**base) + self.spectral.add_component(components) + + self.spectral.setup_fft(real_spectral_coefficients) + + super().__init__(init=self.spectral.init_forward if spectral_space else self.spectral.init) + + self.work_counters[solver_type] = WorkCounter() + self.work_counters['factorizations'] = WorkCounter() + + self.setup_preconditioner(Dirichlet_recombination, left_preconditioner) + + self.cached_factorizations = {} + + def __getattr__(self, name): + """ + Pass requests on to the helper if they are not directly attributes of this class for convenience. + + Args: + name (str): Name of the attribute you want + + Returns: + request + """ + return getattr(self.spectral, name) + + def _setup_operator(self, LHS): + """ + Setup a sparse linear operator by adding relationships. See documentation for ``GenericSpectralLinear.setup_L`` to learn more. + + Args: + LHS (dict): Equations to be added to the operator + + Returns: + sparse linear operator + """ + operator = self.spectral.get_empty_operator_matrix() + for line, equation in LHS.items(): + self.spectral.add_equation_lhs(operator, line, equation) + return self.spectral.convert_operator_matrix_to_operator(operator) + + def setup_L(self, LHS): + """ + Setup the left hand side of the linear operator L and store it in ``self.L``. + + The argument is meant to be a dictionary with the line you want to write the equation in as the key and the relationship between components as another dictionary. For instance, you can add an algebraic condition capturing a first derivative relationship between u and ux as follows: + + ``` + Dx = self.get_differentiation_matrix(axes=(0,)) + I = self.get_Id() + LHS = {'ux': {'u': Dx, 'ux': -I}} + self.setup_L(LHS) + ``` + + If you put zero as right hand side for the solver in the line for ux, ux will contain the x-derivative of u afterwards. + + Args: + LHS (dict): Dictionary containing the equations. + """ + self.L = self._setup_operator(LHS) + + def setup_M(self, LHS): + ''' + Setup mass matrix, see documentation of ``GenericSpectralLinear.setup_L``. + ''' + diff_index = list(LHS.keys()) + self.diff_mask = [me in diff_index for me in self.components] + self.M = self._setup_operator(LHS) + + def setup_preconditioner(self, Dirichlet_recombination=True, left_preconditioner=True): + """ + Get left and right preconditioners. + + Args: + Dirichlet_recombination (bool): Basis conversion for right preconditioner. Useful for Chebychev and Ultraspherical methods. 10/10 would recommend. + left_preconditioner (bool): If True, it will interleave the variables and reverse the Kronecker product + """ + sp = self.spectral.sparse_lib + N = np.prod(self.init[0][1:]) + + Id = sp.eye(N) + Pl_lhs = {comp: {comp: Id} for comp in self.components} + self.Pl = self._setup_operator(Pl_lhs) + + if left_preconditioner: + # reverse Kronecker product + + if self.spectral.useGPU: + R = self.Pl.get().tolil() * 0 + else: + R = self.Pl.tolil() * 0 + + for j in range(self.ncomponents): + for i in range(N): + R[i * self.ncomponents + j, j * N + i] = 1.0 + + self.Pl = self.spectral.sparse_lib.csc_matrix(R) + + if Dirichlet_recombination and type(self.axes[-1]).__name__ in ['ChebychevHelper, Ultraspherical']: + _Pr = self.spectral.get_Dirichlet_recombination_matrix(axis=-1) + else: + _Pr = Id + + Pr_lhs = {comp: {comp: _Pr} for comp in self.components} + self.Pr = self._setup_operator(Pr_lhs) @ self.Pl.T + + def solve_system(self, rhs, dt, u0=None, *args, skip_itransform=False, **kwargs): + """ + Do an implicit Euler step to solve M u_t + Lu = rhs, with M the mass matrix and L the linear operator as setup by + ``GenericSpectralLinear.setup_L`` and ``GenericSpectralLinear.setup_M``. + + The implicit Euler step is (M - dt L) u = M rhs. Note that M need not be invertible as long as (M + dt*L) is. + This means solving with dt=0 to mimic explicit methods does not work for all problems, in particular simple DAEs. + + Note that by putting M rhs on the right hand side, this function can only solve algebraic conditions equal to + zero. If you want something else, it should be easy to overload this function. + """ + + sp = self.spectral.sparse_lib + + if self.spectral_space: + rhs_hat = rhs.copy() + else: + rhs_hat = self.spectral.transform(rhs) + + rhs_hat = (self.M @ rhs_hat.flatten()).reshape(rhs_hat.shape) + rhs_hat = self.spectral.put_BCs_in_rhs_hat(rhs_hat) + rhs_hat = self.Pl @ rhs_hat.flatten() + + if dt not in self.cached_factorizations.keys(): + A = self.M + dt * self.L + A = self.Pl @ self.spectral.put_BCs_in_matrix(A) @ self.Pr + + # import numpy as np + # if A.shape[0] < 200: + # import matplotlib.pyplot as plt + + # # M = self.spectral.put_BCs_in_matrix(self.L.copy()) + # M = A # self.L + # im = plt.imshow((M / abs(M)).real) + # # im = plt.imshow(np.log10(abs(A.toarray())).real) + # # im = plt.imshow(((A.toarray())).real) + # plt.colorbar(im) + # plt.show() + + if self.solver_type.lower() == 'cached_direct': + if dt not in self.cached_factorizations.keys(): + if len(self.cached_factorizations) >= self.max_cached_factorizations: + self.cached_factorizations.pop(list(self.cached_factorizations.keys())[0]) + self.logger.debug(f'Evicted matrix factorization for {dt=:.6f} from cache') + self.cached_factorizations[dt] = self.spectral.linalg.factorized(A) + self.logger.debug(f'Cached matrix factorization for {dt=:.6f}') + self.work_counters['factorizations']() + + _sol_hat = self.cached_factorizations[dt](rhs_hat) + self.logger.debug(f'Used cached matrix factorization for {dt=:.6f}') + + elif self.solver_type.lower() == 'direct': + _sol_hat = sp.linalg.spsolve(A, rhs_hat) + elif self.solver_type.lower() == 'gmres': + _sol_hat, _ = sp.linalg.gmres( + A, + rhs_hat, + x0=u0.flatten(), + **self.solver_args, + callback=self.work_counters[self.solver_type], + callback_type='legacy', + ) + elif self.solver_type.lower() == 'cg': + _sol_hat, _ = sp.linalg.cg( + A, rhs_hat, x0=u0.flatten(), **self.solver_args, callback=self.work_counters[self.solver_type] + ) + else: + raise NotImplementedError(f'Solver {self.solver_type:!} not implemented in {type(self).__name__}!') + + sol_hat = self.spectral.u_init_forward + sol_hat[...] = (self.Pr @ _sol_hat).reshape(sol_hat.shape) + + if self.spectral_space: + return sol_hat + else: + sol = self.spectral.u_init + sol[:] = self.spectral.itransform(sol_hat).real + + if self.spectral.debug: + self.spectral.check_BCs(sol) + + return sol + + +def compute_residual_DAE(self, stage=''): + """ + Computation of the residual that does not add u_0 - u_m in algebraic equations. + + Args: + stage (str): The current stage of the step the level belongs to + """ + + # get current level and problem description + L = self.level + + # Check if we want to skip the residual computation to gain performance + # Keep in mind that skipping any residual computation is likely to give incorrect outputs of the residual! + if stage in self.params.skip_residual_computation: + L.status.residual = 0.0 if L.status.residual is None else L.status.residual + return None + + # check if there are new values (e.g. from a sweep) + # assert L.status.updated + + # compute the residual for each node + + # build QF(u) + res_norm = [] + res = self.integrate() + mask = L.prob.diff_mask + for m in range(self.coll.num_nodes): + res[m][mask] += L.u[0][mask] - L.u[m + 1][mask] + # add tau if associated + if L.tau[m] is not None: + res[m] += L.tau[m] + # use abs function from data type here + res_norm.append(abs(res[m])) + # print(m, [abs(me) for me in res[m]], [abs(me) for me in L.u[0] - L.u[m + 1]]) + + # find maximal residual over the nodes + if L.params.residual_type == 'full_abs': + L.status.residual = max(res_norm) + elif L.params.residual_type == 'last_abs': + L.status.residual = res_norm[-1] + elif L.params.residual_type == 'full_rel': + L.status.residual = max(res_norm) / abs(L.u[0]) + elif L.params.residual_type == 'last_rel': + L.status.residual = res_norm[-1] / abs(L.u[0]) + else: + raise ParameterError( + f'residual_type = {L.params.residual_type} not implemented, choose ' + f'full_abs, last_abs, full_rel or last_rel instead' + ) + + # indicate that the residual has seen the new values + L.status.updated = False + + return None + + +def compute_residual_DAE_MPI(self, stage=None): + """ + Computation of the residual using the collocation matrix Q + + Args: + stage (str): The current stage of the step the level belongs to + """ + from mpi4py import MPI + + L = self.level + + # Check if we want to skip the residual computation to gain performance + # Keep in mind that skipping any residual computation is likely to give incorrect outputs of the residual! + if stage in self.params.skip_residual_computation: + L.status.residual = 0.0 if L.status.residual is None else L.status.residual + return None + + # compute the residual for each node + + # build QF(u) + res = self.integrate(last_only=L.params.residual_type[:4] == 'last') + mask = L.prob.diff_mask + res[mask] += L.u[0][mask] - L.u[self.rank + 1][mask] + # add tau if associated + if L.tau[self.rank] is not None: + res += L.tau[self.rank] + # use abs function from data type here + res_norm = abs(res) + + # find maximal residual over the nodes + if L.params.residual_type == 'full_abs': + L.status.residual = self.comm.allreduce(res_norm, op=MPI.MAX) + elif L.params.residual_type == 'last_abs': + L.status.residual = self.comm.bcast(res_norm, root=self.comm.size - 1) + elif L.params.residual_type == 'full_rel': + L.status.residual = self.comm.allreduce(res_norm / abs(L.u[0]), op=MPI.MAX) + elif L.params.residual_type == 'last_rel': + L.status.residual = self.comm.bcast(res_norm / abs(L.u[0]), root=self.comm.size - 1) + else: + raise NotImplementedError(f'residual type \"{L.params.residual_type}\" not implemented!') + + # indicate that the residual has seen the new values + L.status.updated = False + + return None diff --git a/pySDC/projects/Resilience/venv/modules.sh b/pySDC/projects/Resilience/venv/modules.sh index f296134147..59522b0904 100644 --- a/pySDC/projects/Resilience/venv/modules.sh +++ b/pySDC/projects/Resilience/venv/modules.sh @@ -5,3 +5,4 @@ module load OpenMPI module load FFTW module load Python/3.11.3 module load mpi4py +module load FFmpeg/.6.0 diff --git a/pySDC/tests/test_helpers/test_spectral_helper.py b/pySDC/tests/test_helpers/test_spectral_helper.py new file mode 100644 index 0000000000..627632b14b --- /dev/null +++ b/pySDC/tests/test_helpers/test_spectral_helper.py @@ -0,0 +1,593 @@ +import pytest + + +@pytest.mark.base +@pytest.mark.parametrize('nx', [16]) +@pytest.mark.parametrize('nz', [16]) +@pytest.mark.parametrize('variant', ['T2U', 'T2T']) +@pytest.mark.parametrize('axes', [(-2,), (-1,), (-2, -1)]) +def test_integration_matrix2D(nx, nz, variant, axes, useMPI=False, **kwargs): + import numpy as np + from pySDC.helpers.spectral_helper import SpectralHelper + + if useMPI: + from mpi4py import MPI + + comm = MPI.COMM_WORLD + else: + comm = None + + helper = SpectralHelper(comm=comm, debug=True) + helper.add_axis(base='fft', N=nx) + helper.add_axis(base='cheby', N=nz) + helper.setup_fft() + + Z, X = helper.get_grid() + + conv = helper.get_basis_change_matrix() + S = helper.get_integration_matrix(axes=axes) + + u = helper.u_init + u[0, ...] = np.sin(X) * Z**2 + np.cos(X) * Z**3 + if axes == (-2,): + expect = -np.cos(X) * Z**2 + np.sin(X) * Z**3 + elif axes == (-1,): + expect = np.sin(X) * 1 / 3 * Z**3 + np.cos(X) * 1 / 4 * Z**4 + elif axes in [(-2, -1), (-1, -2)]: + expect = -np.cos(X) * 1 / 3 * Z**3 + np.sin(X) * 1 / 4 * Z**4 + else: + raise NotImplementedError + + u_hat = helper.transform(u, axes=(-2, -1)) + S_u_hat = (conv @ S @ u_hat.flatten()).reshape(u_hat.shape) + S_u = helper.itransform(S_u_hat, axes=(-1, -2)) + + assert np.allclose(S_u, expect, atol=1e-12) + + +@pytest.mark.base +@pytest.mark.parametrize('nx', [32]) +@pytest.mark.parametrize('nz', [16]) +@pytest.mark.parametrize('variant', ['T2U', 'T2T']) +@pytest.mark.parametrize('axes', [(-2,), (-1,), (-2, -1)]) +@pytest.mark.parametrize('bx', ['cheby', 'fft']) +@pytest.mark.parametrize('bz', ['cheby', 'fft']) +def test_differentiation_matrix2D(nx, nz, variant, axes, bx, bz, useMPI=False, **kwargs): + import numpy as np + from pySDC.helpers.spectral_helper import SpectralHelper + + if useMPI: + from mpi4py import MPI + + comm = MPI.COMM_WORLD + else: + comm = None + + helper = SpectralHelper(comm=comm, debug=True) + helper.add_axis(base=bx, N=nx) + helper.add_axis(base=bz, N=nz) + helper.setup_fft() + + Z, X = helper.get_grid() + conv = helper.get_basis_change_matrix() + D = helper.get_differentiation_matrix(axes) + + u = helper.u_init + + if bz == 'cheby': + u[0, ...] = np.sin(X) * Z**2 + Z**3 + np.cos(2 * X) + if axes == (-2,): + expect = np.cos(X) * Z**2 - 2 * np.sin(2 * X) + elif axes == (-1,): + expect = np.sin(X) * Z * 2 + Z**2 * 3 + elif axes in [(-2, -1), (-1, -2)]: + expect = np.cos(X) * 2 * Z + else: + raise NotImplementedError + else: + u[0, ...] = np.sin(X) * np.cos(2 * Z) + np.cos(2 * X) + np.sin(Z) + if axes == (-2,): + expect = np.cos(X) * np.cos(2 * Z) - 2 * np.sin(2 * X) + elif axes == (-1,): + expect = np.sin(X) * (-2) * np.sin(2 * Z) + np.cos(Z) + elif axes in [(-2, -1), (-1, -2)]: + expect = -2 * np.cos(X) * np.sin(2 * Z) + else: + raise NotImplementedError + + u_hat = helper.transform(u, axes=(-2, -1)) + D_u_hat = (conv @ D @ u_hat.flatten()).reshape(u_hat.shape) + D_u = helper.itransform(D_u_hat, axes=(-1, -2)).real + + assert np.allclose(D_u, expect, atol=1e-12) + + +@pytest.mark.base +@pytest.mark.parametrize('nx', [32]) +@pytest.mark.parametrize('nz', [16]) +@pytest.mark.parametrize('variant', ['T2U', 'T2T']) +@pytest.mark.parametrize('bx', ['cheby', 'fft']) +def test_identity_matrix2D(nx, nz, variant, bx, useMPI=False, **kwargs): + import numpy as np + from pySDC.helpers.spectral_helper import SpectralHelper + + if useMPI: + from mpi4py import MPI + + comm = MPI.COMM_WORLD + else: + comm = None + + helper = SpectralHelper(comm=comm, debug=True) + helper.add_axis(base=bx, N=nx) + helper.add_axis(base='cheby', N=nz) + helper.setup_fft() + + Z, X = helper.get_grid() + conv = helper.get_basis_change_matrix() + I = helper.get_Id() + + u = helper.u_init + u[0, ...] = np.sin(X) * Z**2 + Z**3 + np.cos(2 * X) + + u_hat = helper.transform(u, axes=(-2, -1)) + I_u_hat = (conv @ I @ u_hat.flatten()).reshape(u_hat.shape) + I_u = helper.itransform(I_u_hat, axes=(-1, -2)) + + assert np.allclose(I_u, u, atol=1e-12) + + +@pytest.mark.base +@pytest.mark.parametrize('N', [4, 32]) +@pytest.mark.parametrize('base', ['cheby']) +@pytest.mark.parametrize('type', ['diff', 'int']) +def test_matrix1D(N, base, type): + import numpy as np + from pySDC.helpers.spectral_helper import SpectralHelper + + coeffs = np.random.random(N) + + helper = SpectralHelper(debug=True) + helper.add_axis(base=base, N=N) + helper.setup_fft() + + x = helper.get_grid() + + if type == 'diff': + D = helper.get_differentiation_matrix(axes=(-1,)) + elif type == 'int': + D = helper.get_integration_matrix(axes=(-1,)) + + C = helper.get_basis_change_matrix() + + u = helper.u_init + u[0] = C @ D @ coeffs + du = helper.itransform(u, axes=(-1,)) + + if type == 'diff': + exact = np.polynomial.Chebyshev(coeffs).deriv(1)(x) + elif type == 'int': + exact = np.polynomial.Chebyshev(coeffs).integ(1)(x) + + assert np.allclose(exact, du) + + +def _test_transform_dealias( + bx, + bz, + axis, + nx=2**4 + 1, + nz=2**2 + 1, + padding=3 / 2, + axes=( + -2, + -1, + ), + useMPI=True, + **kwargs, +): + import numpy as np + from pySDC.helpers.spectral_helper import SpectralHelper + + if useMPI: + from mpi4py import MPI + + comm = MPI.COMM_WORLD + else: + comm = None + + helper = SpectralHelper(comm=comm, debug=True) + helper.add_axis(base=bx, N=nx) + helper.add_axis(base=bz, N=nz) + helper.setup_fft() + xp = helper.xp + + _padding = [ + padding, + ] * helper.ndim + + helper_pad = SpectralHelper(comm=comm, debug=True) + helper_pad.add_axis(base=bx, N=int(_padding[0] * nx)) + helper_pad.add_axis(base=bz, N=int(_padding[1] * nz)) + helper_pad.setup_fft() + + u_hat = helper.u_init_forward + u2_hat_expect = helper.u_init_forward + u_expect = helper.u_init + u_expect_pad = helper_pad.u_init + Kz, Kx = helper.get_wavenumbers() + Z, X = helper.get_grid() + Z_pad, X_pad = helper_pad.get_grid() + + if axis == -2: + f = nx // 3 + u_hat[0][xp.logical_and(xp.abs(Kx) == f, Kz == 0)] += 1 + u2_hat_expect[0][xp.logical_and(xp.abs(Kx) == 2 * f, Kz == 0)] += 1 / nx + u2_hat_expect[0][xp.logical_and(xp.abs(Kx) == 0, Kz == 0)] += 2 / nx + u_expect[0] += np.cos(f * X) * 2 / nx + u_expect_pad[0] += np.cos(f * X_pad) * 2 / nx + elif axis == -1: + f = nz // 2 + 1 + u_hat[0][xp.logical_and(xp.abs(Kz) == f, Kx == 0)] += 1 + u2_hat_expect[0][xp.logical_and(Kz == 2 * f, Kx == 0)] += 1 / (2 * nx) + u2_hat_expect[0][xp.logical_and(Kz == 0, Kx == 0)] += 1 / (2 * nx) + + coef = np.zeros(nz) + coef[f] = 1 / nx + u_expect[0] = np.polynomial.Chebyshev(coef)(Z) + u_expect_pad[0] = np.polynomial.Chebyshev(coef)(Z_pad) + elif axis in [(-1, -2), (-2, -1)]: + fx = nx // 3 + fz = nz // 2 + 1 + + u_hat[0][xp.logical_and(xp.abs(Kx) == fx, Kz == 0)] += 1 + u_hat[0][xp.logical_and(xp.abs(Kz) == fz, Kx == 0)] += 1 + + u2_hat_expect[0][xp.logical_and(xp.abs(Kx) == 2 * fx, Kz == 0)] += 1 / nx + u2_hat_expect[0][xp.logical_and(xp.abs(Kx) == 0, Kz == 0)] += 2 / nx + u2_hat_expect[0][xp.logical_and(Kz == 2 * fz, Kx == 0)] += 1 / (2 * nx) + u2_hat_expect[0][xp.logical_and(Kz == 0, Kx == 0)] += 1 / (2 * nx) + u2_hat_expect[0][xp.logical_and(Kz == fz, xp.abs(Kx) == fx)] += 2 / nx + + coef = np.zeros(nz) + coef[fz] = 1 / nx + + u_expect[0] = np.cos(fx * X) * 2 / nx + np.polynomial.Chebyshev(coef)(Z) + u_expect_pad[0] = np.cos(fx * X_pad) * 2 / nx + np.polynomial.Chebyshev(coef)(Z_pad) + else: + raise NotImplementedError + + assert bx == 'fft' and bz == 'cheby', 'This test is not implemented for the bases you are looking for' + + u_pad = helper.itransform(u_hat, padding=_padding, axes=axes) + u = helper.itransform(u_hat, axes=axes).real + + assert not np.allclose(u_pad.shape, u.shape) + + u2 = u**2 + u2_pad = u_pad**2 + + assert np.allclose(u, u_expect) + assert np.allclose(u_pad, u_expect_pad) + + assert np.allclose(u2_hat_expect, helper.transform(u2_pad, padding=_padding)) + assert not np.allclose(u2_hat_expect, helper.transform(u2)), 'Test is too boring, no dealiasing needed' + + +@pytest.mark.base +@pytest.mark.parametrize('nx', [3, 8]) +@pytest.mark.parametrize('nz', [3, 8]) +@pytest.mark.parametrize('bz', ['fft', 'cheby']) +@pytest.mark.parametrize('bx', ['fft', 'cheby']) +@pytest.mark.parametrize('axes', [(-1,), (-2,), (-1, -2), (-2, -1)]) +def test_transform(nx, nz, bx, bz, axes, useMPI=False, **kwargs): + import numpy as np + from pySDC.helpers.spectral_helper import SpectralHelper + + if useMPI: + from mpi4py import MPI + + comm = MPI.COMM_WORLD + rank = comm.rank + else: + comm = None + + helper = SpectralHelper(comm=comm, debug=True) + helper.add_axis(base=bx, N=nx) + helper.add_axis(base=bz, N=nz) + helper.setup_fft() + + u = helper.u_init + u[...] = np.random.random(u.shape) + + u_all = np.empty(shape=(1, nx, nz), dtype=u.dtype) + + if useMPI: + rank = comm.rank + u_all[...] = (np.array(comm.allgather(u[0]))).reshape(u_all.shape) + if comm.size == 1: + assert np.allclose(u_all, u) + else: + rank = 0 + u_all[...] = u + + expect_trf = u_all.copy() + + if bx == 'fft' and bz == 'cheby': + axes_1d = sorted(axes)[::-1] + elif bx == 'cheby' and bz == 'fft': + axes_1d = sorted(axes) + else: + axes_1d = axes + + for i in axes_1d: + base = helper.axes[i] + expect_trf = base.transform(expect_trf, axis=i) + + trf = helper.transform(u, axes=axes) + itrf = helper.itransform(trf, axes=axes) + + expect_local = expect_trf[:, trf.shape[1] * rank : trf.shape[1] * (rank + 1), :] + + assert np.allclose(expect_local, trf), 'Forward transform is unexpected' + assert np.allclose(itrf, u), 'Backward transform is unexpected' + + +def run_MPI_test(num_procs, **kwargs): + import os + import subprocess + + # Set python path once + my_env = os.environ.copy() + my_env['PYTHONPATH'] = '../../..:.' + my_env['COVERAGE_PROCESS_START'] = 'pyproject.toml' + + cmd = f"mpirun -np {num_procs} python {__file__}" + + for key, value in kwargs.items(): + cmd += f' --{key}={value}' + p = subprocess.Popen(cmd.split(), env=my_env, cwd=".") + + p.wait() + assert p.returncode == 0, 'ERROR: did not get return code 0, got %s with %2i processes' % ( + p.returncode, + num_procs, + ) + + +@pytest.mark.mpi4py +@pytest.mark.parametrize('nx', [4, 8]) +@pytest.mark.parametrize('nz', [4, 8]) +@pytest.mark.parametrize('bz', ['fft', 'cheby']) +@pytest.mark.parametrize('bx', ['fft']) +@pytest.mark.parametrize('num_procs', [2, 1]) +@pytest.mark.parametrize('axes', ["-1", "-2", "-1,-2"]) +def test_transform_MPI(nx, nz, bx, bz, num_procs, axes): + run_MPI_test(num_procs=num_procs, test='transform', nx=nx, nz=nz, bx=bx, bz=bz, axes=axes) + + +@pytest.mark.mpi4py +@pytest.mark.parametrize('nx', [8]) +@pytest.mark.parametrize('nz', [16]) +@pytest.mark.parametrize('bx', ['fft']) +@pytest.mark.parametrize('num_procs', [2, 1]) +@pytest.mark.parametrize('axes', ["-1", "-1,-2"]) +def test_differentiation_MPI(nx, nz, bx, num_procs, axes): + run_MPI_test(num_procs=num_procs, test='diff', nx=nx, nz=nz, bx=bx, bz='cheby', axes=axes) + + +@pytest.mark.mpi4py +@pytest.mark.parametrize('nx', [8]) +@pytest.mark.parametrize('nz', [16]) +@pytest.mark.parametrize('num_procs', [2, 1]) +@pytest.mark.parametrize('axes', ["-1", "-1,-2"]) +def test_integration_MPI(nx, nz, num_procs, axes): + run_MPI_test(num_procs=num_procs, test='int', nx=nx, nz=nz, axes=axes) + + +@pytest.mark.base +@pytest.mark.parametrize('N', [8, 32]) +@pytest.mark.parametrize('bc_val', [0, 3.1415]) +def test_tau_method_integral(N, bc_val): + test_tau_method(bc=0, N=N, bc_val=bc_val, kind='integral') + + +@pytest.mark.base +@pytest.mark.parametrize('bc', [-1, 0, 1]) +@pytest.mark.parametrize('N', [8, 32]) +@pytest.mark.parametrize('bc_val', [-99, 3.1415]) +def test_tau_method(bc, N, bc_val, kind='Dirichlet'): + ''' + solve u_x - u + tau P = 0, u(bc) = bc_val + + We choose P = T_N or U_N. We replace the last row in the matrix with the boundary condition to get a + unique solution for the given resolution. + + The test verifies that the solution satisfies the perturbed equation and the boundary condition. + ''' + from pySDC.helpers.spectral_helper import SpectralHelper + import numpy as np + import scipy.sparse as sp + + helper = SpectralHelper(debug=True) + helper.add_component('u') + helper.add_axis(base='cheby', N=N) + helper.setup_fft() + + if kind == 'integral': + helper.add_BC('u', 'u', 0, v=bc_val, kind='integral') + else: + helper.add_BC('u', 'u', 0, x=bc, v=bc_val, kind='Dirichlet') + helper.setup_BCs() + + C = helper.get_basis_change_matrix() + D = helper.get_differentiation_matrix(axes=(-1,)) + Id = helper.get_Id() + + _A = helper.get_empty_operator_matrix() + helper.add_equation_lhs(_A, 'u', {'u': D - Id}) + A = helper.convert_operator_matrix_to_operator(_A) + A = helper.put_BCs_in_matrix(A) + + rhs = helper.put_BCs_in_rhs(np.zeros((1, N))) + rhs_hat = helper.transform(rhs, axes=(-1,)) + + sol_hat = sp.linalg.spsolve(A, rhs_hat.flatten()) + + sol_poly = np.polynomial.Chebyshev(sol_hat) + d_sol_poly = sol_poly.deriv(1) + x = np.linspace(-1, 1, 100) + + if kind == 'integral': + integral = sol_poly.integ(1, lbnd=-1) + assert np.isclose(integral(1), bc_val), 'Solution does not satisfy boundary condition' + else: + assert np.isclose(sol_poly(bc), bc_val), 'Solution does not satisfy boundary condition' + + coef = np.append(np.zeros(N - 1), [1]) + P = np.polynomial.Chebyshev(C @ coef) + tau = (d_sol_poly(x) - sol_poly(x)) / P(x) + assert np.allclose(tau, tau[0]), 'Solution does not satisfy perturbed equation' + + +@pytest.mark.base +@pytest.mark.parametrize('variant', ['T2T', 'T2U']) +@pytest.mark.parametrize('nx', [4, 8]) +@pytest.mark.parametrize('nz', [4, 8]) +@pytest.mark.parametrize('bc_val', [-2, 1.0]) +def test_tau_method2D(variant, nz, nx, bc_val, bc=-1, useMPI=False, plotting=False, **kwargs): + ''' + solve u_z - 0.1u_xx -u_x + tau P = 0, u(bc) = sin(bc_val*x) -> space-time discretization of advection-diffusion + problem. We do FFT in x-direction and Chebychov in z-direction. + ''' + from pySDC.helpers.spectral_helper import SpectralHelper + import numpy as np + + if useMPI: + from mpi4py import MPI + + comm = MPI.COMM_WORLD + else: + comm = None + + helper = SpectralHelper(comm=comm, debug=True) + helper.add_axis('fft', N=nx) + helper.add_axis('cheby', N=nz) + helper.add_component(['u']) + helper.setup_fft() + + Z, X = helper.get_grid() + x = X[:, 0] + z = Z[0, :] + shape = helper.init[0][1:] + + bcs = np.sin(bc_val * x) + helper.add_BC('u', 'u', 1, kind='dirichlet', x=bc, v=bcs) + helper.setup_BCs() + + # generate matrices + Dz = helper.get_differentiation_matrix(axes=(1,)) + Dx = helper.get_differentiation_matrix(axes=(0,)) + Dxx = helper.get_differentiation_matrix(axes=(0,), p=2) + + # generate operator + _A = helper.get_empty_operator_matrix() + helper.add_equation_lhs(_A, 'u', {'u': Dz - Dxx * 1e-1 - Dx}) + A = helper.convert_operator_matrix_to_operator(_A) + + # prepare system to solve + A = helper.put_BCs_in_matrix(A) + rhs_hat = helper.put_BCs_in_rhs_hat(helper.u_init_forward) + + # solve the system + sol_hat = helper.u_init_forward + sol_hat[0] = (helper.sparse_lib.linalg.spsolve(A, rhs_hat.flatten())).reshape(X.shape) + sol = helper.itransform(sol_hat, axes=(-2, -1)).real + + # construct polynomials for testing + sol_cheby = helper.transform(sol, axes=(-1,)) + polys = [np.polynomial.Chebyshev(sol_cheby[0, i, :]) for i in range(shape[0])] + + Pz = np.polynomial.Chebyshev(np.append(np.zeros(nz - 1), [1])) + tau_term, _ = np.meshgrid(Pz(z), np.ones(shape[0])) + error = ((A @ sol_hat.flatten()).reshape(X.shape) / tau_term).real + + if plotting: + import matplotlib.pyplot as plt + + im = plt.pcolormesh(X, Z, sol[0]) + plt.colorbar(im) + plt.xlabel('x') + plt.ylabel('t') + plt.show() + + for i in range(shape[0]): + + assert np.isclose(polys[i](bc), bcs[i]), f'Solution does not satisfy boundary condition x={x[i]}' + + assert np.allclose( + polys[i](z), sol[0, i, :] + ), f'Solution is incorrectly transformed back to real space at x={x[i]}' + + assert np.allclose(error, error[0, 0]), 'Solution does not satisfy perturbed equation' + + +@pytest.mark.mpi4py +@pytest.mark.parametrize('variant', ['T2T', 'T2U']) +@pytest.mark.parametrize('nx', [4, 8]) +@pytest.mark.parametrize('nz', [4, 8]) +@pytest.mark.parametrize('bc_val', [-2]) +@pytest.mark.parametrize('num_procs', [2, 1]) +def test_tau_method2D_MPI(variant, nz, nx, bc_val, num_procs, **kwargs): + run_MPI_test(variant=variant, nz=nz, nx=nx, bc_val=bc_val, num_procs=num_procs, test='tau') + + +@pytest.mark.mpi4py +@pytest.mark.parametrize('num_procs', [1]) +@pytest.mark.parametrize('axis', [-1, -2]) +@pytest.mark.parametrize('bx', ['fft']) +@pytest.mark.parametrize('bz', ['cheby']) +def test_dealias_MPI(num_procs, axis, bx, bz, nx=32, nz=64, **kwargs): + run_MPI_test(num_procs=num_procs, axis=axis, nx=nx, nz=nz, bx=bx, bz=bz, test='dealias') + + +if __name__ == '__main__': + str_to_bool = lambda me: False if me == 'False' else True + str_to_tuple = lambda arg: tuple(int(me) for me in arg.split(',')) + + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('--nx', type=int, help='Dof in x direction') + parser.add_argument('--nz', type=int, help='Dof in z direction') + parser.add_argument('--axes', type=str_to_tuple, help='Axes over which to transform') + parser.add_argument('--axis', type=int, help='Direction of the action') + parser.add_argument('--bz', type=str, help='Base in z direction') + parser.add_argument('--bx', type=str, help='Base in x direction') + parser.add_argument('--bc_val', type=int, help='Value of boundary condition') + parser.add_argument('--test', type=str, help='type of test', choices=['transform', 'diff', 'int', 'tau', 'dealias']) + parser.add_argument('--variant', type=str, help='Chebychov mode', choices=['T2T', 'T2U'], default='T2U') + parser.add_argument('--useMPI', type=str_to_bool, help='use MPI or not', choices=[True, False], default=True) + args = parser.parse_args() + + if args.test == 'transform': + test_transform(**vars(args)) + elif args.test == 'diff': + test_differentiation_matrix2D(**vars(args)) + elif args.test == 'int': + test_integration_matrix2D(**vars(args)) + elif args.test == 'tau': + test_tau_method2D(**vars(args)) + elif args.test == 'dealias': + _test_transform_dealias(**vars(args)) + elif args.test is None: + # test_transform(8, 3, 'fft', 'cheby', (-1,)) + # test_differentiation_matrix2D(2**5, 2**5, 'T2U', bx='fft', bz='fft', axes=(-2, -1)) + # test_matrix1D(4, 'cheby', 'int') + # test_tau_method(-1, 8, 99, kind='Dirichlet') + test_tau_method2D('T2U', 2**8, 2**8, -2, plotting=True) + # test_filter(6, 6, (0,)) + # _test_transform_dealias('fft', 'cheby', (-1, -2)) + else: + raise NotImplementedError + print('done') diff --git a/pySDC/tests/test_helpers/test_spectral_helper_1d_chebychev.py b/pySDC/tests/test_helpers/test_spectral_helper_1d_chebychev.py new file mode 100644 index 0000000000..37a1ee7175 --- /dev/null +++ b/pySDC/tests/test_helpers/test_spectral_helper_1d_chebychev.py @@ -0,0 +1,384 @@ +import pytest + + +@pytest.mark.base +@pytest.mark.parametrize('N', [4, 32]) +def test_D2T_conversion_matrices(N): + import numpy as np + from pySDC.helpers.spectral_helper import ChebychevHelper + + cheby = ChebychevHelper(N) + + x = np.linspace(-1, 1, N) + D2T = cheby.get_conv('D2T') + + for i in range(N): + coeffs = np.zeros(N) + coeffs[i] = 1.0 + T_coeffs = D2T @ coeffs + + Dn = np.polynomial.Chebyshev(T_coeffs)(x) + + expect_left = (-1) ** i if i < 2 else 0 + expect_right = 1 if i < 2 else 0 + + assert Dn[0] == expect_left + assert Dn[-1] == expect_right + + +@pytest.mark.base +@pytest.mark.parametrize('N', [4, 32]) +def test_T_U_conversion(N): + import numpy as np + from scipy.special import chebyt, chebyu + from pySDC.helpers.spectral_helper import ChebychevHelper + + cheby = ChebychevHelper(N) + + T2U = cheby.get_conv('T2U') + U2T = cheby.get_conv('U2T') + + coeffs = np.random.random(N) + x = cheby.get_1dgrid() + + def eval_poly(poly, coeffs, x): + return np.array([coeffs[i] * poly(i)(x) for i in range(len(coeffs))]).sum(axis=0) + + u = eval_poly(chebyu, coeffs, x) + t_from_u = eval_poly(chebyt, U2T @ coeffs, x) + t_from_u_r = eval_poly(chebyt, coeffs @ U2T.T, x) + + t = eval_poly(chebyt, coeffs, x) + u_from_t = eval_poly(chebyu, T2U @ coeffs, x) + u_from_t_r = eval_poly(chebyu, coeffs @ T2U.T, x) + + assert np.allclose(u, t_from_u) + assert np.allclose(u, t_from_u_r) + assert np.allclose(t, u_from_t) + assert np.allclose(t, u_from_t_r) + + +@pytest.mark.base +@pytest.mark.parametrize('name', ['T2U', 'T2D', 'T2T']) +def test_conversion_inverses(name): + from pySDC.helpers.spectral_helper import ChebychevHelper + import numpy as np + + N = 8 + cheby = ChebychevHelper(N) + P = cheby.get_conv(name) + Pinv = cheby.get_conv(name[::-1]) + assert np.allclose((P @ Pinv).toarray(), np.diag(np.ones(N))) + + +@pytest.mark.base +@pytest.mark.parametrize('N', [4, 32]) +def test_differentiation_matrix(N): + import numpy as np + import scipy + from pySDC.helpers.spectral_helper import ChebychevHelper + + cheby = ChebychevHelper(N) + x = np.cos(np.pi / N * (np.arange(N) + 0.5)) + coeffs = np.random.random(N) + norm = cheby.get_norm() + + D = cheby.get_differentiation_matrix(1) + + du = scipy.fft.idct(D @ coeffs / norm) + exact = np.polynomial.Chebyshev(coeffs).deriv(1)(x) + + assert np.allclose(exact, du) + + +@pytest.mark.base +@pytest.mark.parametrize('N', [4, 32]) +def test_integration_matrix(N): + import numpy as np + from pySDC.helpers.spectral_helper import ChebychevHelper + + cheby = ChebychevHelper(N) + coeffs = np.random.random(N) + coeffs[-1] = 0 + + D = cheby.get_integration_matrix() + + du = D @ coeffs + exact = np.polynomial.Chebyshev(coeffs).integ(1, lbnd=0) + + assert np.allclose(exact.coef[:-1], du) + + +@pytest.mark.base +@pytest.mark.parametrize('N', [4]) +@pytest.mark.parametrize('d', [1, 2, 3]) +@pytest.mark.parametrize('transform_type', ['dct', 'fft']) +def test_transform(N, d, transform_type): + import scipy + import numpy as np + from pySDC.helpers.spectral_helper import ChebychevHelper + + cheby = ChebychevHelper(N, transform_type=transform_type) + u = np.random.random((d, N)) + norm = cheby.get_norm() + x = cheby.get_1dgrid() + + itransform = cheby.itransform(u, axis=-1).real + + for i in range(d): + assert np.allclose(np.polynomial.Chebyshev(u[i])(x), itransform[i]) + assert np.allclose(u.shape, itransform.shape) + assert np.allclose(scipy.fft.dct(u, axis=-1) * norm, cheby.transform(u, axis=-1).real) + assert np.allclose(scipy.fft.idct(u / norm, axis=-1), itransform) + assert np.allclose(cheby.transform(cheby.itransform(u)), u) + assert np.allclose(cheby.itransform(cheby.transform(u)), u) + + +@pytest.mark.base +@pytest.mark.parametrize('N', [8, 32]) +def test_integration_BC(N): + from pySDC.helpers.spectral_helper import ChebychevHelper + import numpy as np + + cheby = ChebychevHelper(N) + coef = np.random.random(N) + + BC = cheby.get_integ_BC_row() + + polynomial = np.polynomial.Chebyshev(coef) + reference_integral = polynomial.integ(lbnd=-1, k=0) + + integral = sum(coef * BC) + assert np.isclose(integral, reference_integral(1)) + + +@pytest.mark.base +@pytest.mark.parametrize('N', [4, 32]) +def test_norm(N): + from pySDC.helpers.spectral_helper import ChebychevHelper + import numpy as np + import scipy + + cheby = ChebychevHelper(N) + coeffs = np.random.random(N) + x = cheby.get_1dgrid() + norm = cheby.get_norm() + + u = np.polynomial.Chebyshev(coeffs)(x) + u_dct = scipy.fft.idct(coeffs / norm) + coeffs_dct = scipy.fft.dct(u) * norm + + assert np.allclose(u, u_dct) + assert np.allclose(coeffs, coeffs_dct) + + +@pytest.mark.base +@pytest.mark.parametrize('bc', [-1, 0, 1]) +@pytest.mark.parametrize('N', [3, 32]) +@pytest.mark.parametrize('bc_val', [-99, 3.1415]) +def test_tau_method(bc, N, bc_val): + ''' + solve u_x - u + tau P = 0, u(bc) = bc_val + + We choose P = T_N or U_N. We replace the last row in the matrix with the boundary condition to get a + unique solution for the given resolution. + + The test verifies that the solution satisfies the perturbed equation and the boundary condition. + ''' + from pySDC.helpers.spectral_helper import ChebychevHelper + import numpy as np + + cheby = ChebychevHelper(N) + x = cheby.get_1dgrid() + + coef = np.append(np.zeros(N - 1), [1]) + rhs = np.append(np.zeros(N - 1), [bc_val]) + + P = np.polynomial.Chebyshev(coef) + D = cheby.get_differentiation_matrix() + Id = np.diag(np.ones(N)) + + A = D - Id + A[-1, :] = cheby.get_Dirichlet_BC_row(bc) + + sol_hat = np.linalg.solve(A, rhs) + + sol_poly = np.polynomial.Chebyshev(sol_hat) + d_sol_poly = sol_poly.deriv(1) + x = np.linspace(-1, 1, 100) + + assert np.isclose(sol_poly(bc), bc_val), 'Solution does not satisfy boundary condition' + + tau = (d_sol_poly(x) - sol_poly(x)) / P(x) + assert np.allclose(tau, tau[0]), 'Solution does not satisfy perturbed equation' + + +@pytest.mark.base +@pytest.mark.parametrize('bc', [-1, 1]) +@pytest.mark.parametrize('nx', [4, 8]) +@pytest.mark.parametrize('nz', [4, 8]) +@pytest.mark.parametrize('bc_val', [-2, 1.0]) +def test_tau_method2D(bc, nz, nx, bc_val, plotting=False): + ''' + solve u_z - 0.1u_xx -u_x + tau P = 0, u(bc) = sin(bc_val*x) -> space-time discretization of advection-diffusion + problem. We do FFT in x-direction and Chebychev in z-direction. + ''' + from pySDC.helpers.spectral_helper import ChebychevHelper, FFTHelper + import numpy as np + import scipy.sparse as sp + + cheby = ChebychevHelper(nz) + fft = FFTHelper(nx) + + # generate grid + x = fft.get_1dgrid() + z = cheby.get_1dgrid() + Z, X = np.meshgrid(z, x) + + # put BCs in right hand side + bcs = np.sin(bc_val * x) + rhs = np.zeros_like(X) + rhs[:, -1] = bcs + rhs_hat = fft.transform(rhs, axis=-2) # the rhs is already in Chebychev spectral space + + # generate matrices + Dx = fft.get_differentiation_matrix(p=2) * 1e-1 + fft.get_differentiation_matrix() + Ix = fft.get_Id() + Dz = cheby.get_differentiation_matrix() + Iz = cheby.get_Id() + A = sp.kron(Ix, Dz) - sp.kron(Dx, Iz) + + # put BCs in the system matrix + BCz = sp.eye(nz, format='lil') * 0 + BCz[-1, :] = cheby.get_Dirichlet_BC_row(bc) + BC = sp.kron(Ix, BCz, format='lil') + A[BC != 0] = BC[BC != 0] + + # solve the system + sol_hat = (sp.linalg.spsolve(A, rhs_hat.flatten())).reshape(rhs.shape) + + # transform back to real space + _sol = fft.itransform(sol_hat, axis=-2).real + sol = cheby.itransform(_sol, axis=-1) + + # construct polynomials for testing + polys = [np.polynomial.Chebyshev(_sol[i, :]) for i in range(nx)] + # d_polys = [me.deriv(1) for me in polys] + # _z = np.linspace(-1, 1, 100) + + if plotting: + import matplotlib.pyplot as plt + + im = plt.pcolormesh(X, Z, sol) + plt.colorbar(im) + plt.xlabel('x') + plt.ylabel('t') + plt.show() + + for i in range(nx): + + assert np.isclose(polys[i](bc), bcs[i]), f'Solution does not satisfy boundary condition x={x[i]}' + + assert np.allclose( + polys[i](z), sol[i, :] + ), f'Solution is incorrectly transformed back to real space at x={x[i]}' + + # coef = np.append(np.zeros(nz - 1), [1]) + # Pz = np.polynomial.Chebyshev(coef) + # tau = (d_polys[i](_z) - polys[i](_z)) / Pz(_z) + # plt.plot(_z, tau) + # plt.show() + # assert np.allclose(tau, tau[0]), f'Solution does not satisfy perturbed equation at x={x[i]}' + + +@pytest.mark.base +@pytest.mark.parametrize('nx', [4, 8]) +@pytest.mark.parametrize('nz', [16]) +@pytest.mark.parametrize('bc_val', [4.0]) +def test_tau_method2D_diffusion(nz, nx, bc_val, plotting=False): + ''' + Solve a Poisson problem with funny Dirichlet BCs in z-direction and periodic in x-direction. + ''' + from pySDC.helpers.spectral_helper import ChebychevHelper, FFTHelper + import numpy as np + import scipy.sparse as sp + + cheby = ChebychevHelper(nz) + fft = FFTHelper(nx) + + # generate grid + x = fft.get_1dgrid() + z = cheby.get_1dgrid() + Z, X = np.meshgrid(z, x) + + # put BCs in right hand side + rhs = np.zeros((2, nx, nz)) # components u and u_x + rhs[0, :, -1] = np.sin(bc_val * x) + 1 + rhs[1, :, -1] = 3 * np.exp(-((x - 3.6) ** 2)) + np.cos(x) + rhs_hat = fft.transform(rhs, axis=-2) # the rhs is already in Chebychev spectral space + + # generate 1D matrices + Dx = fft.get_differentiation_matrix() + Ix = fft.get_Id() + Dz = cheby.get_differentiation_matrix() + Iz = cheby.get_Id() + + # generate 2D matrices + D = sp.kron(Ix, Dz) + sp.kron(Dx, Iz) + I = sp.kron(Ix, Iz) + O = I * 0 + + # generate system matrix + A = sp.bmat([[O, D], [D, -I]], format='lil') + + # generate BC matrices + BCa = sp.eye(nz, format='lil') * 0 + BCa[-1, :] = cheby.get_Dirichlet_BC_row(-1) + BCa = sp.kron(Ix, BCa, format='lil') + + BCb = sp.eye(nz, format='lil') * 0 + BCb[-1, :] = cheby.get_Dirichlet_BC_row(1) + BCb = sp.kron(Ix, BCb, format='lil') + BC = sp.bmat([[BCa, O], [BCb, O]], format='lil') + + # put BCs in the system matrix + A[BC != 0] = BC[BC != 0] + + # solve the system + sol_hat = (sp.linalg.spsolve(A, rhs_hat.flatten())).reshape(rhs.shape) + + # transform back to real space + _sol = fft.itransform(sol_hat, axis=-2).real + sol = cheby.itransform(_sol, axis=-1) + + polys = [np.polynomial.Chebyshev(_sol[0, i, :]) for i in range(nx)] + + if plotting: + import matplotlib.pyplot as plt + + im = plt.pcolormesh(X, Z, sol[0]) + plt.colorbar(im) + plt.xlabel('x') + plt.ylabel('z') + plt.show() + + for i in range(nx): + + assert np.isclose(polys[i](-1), rhs[0, i, -1]), f'Solution does not satisfy lower boundary condition x={x[i]}' + assert np.isclose(polys[i](1), rhs[1, i, -1]), f'Solution does not satisfy upper boundary condition x={x[i]}' + + assert np.allclose( + polys[i](z), sol[0, i, :] + ), f'Solution is incorrectly transformed back to real space at x={x[i]}' + + +if __name__ == '__main__': + test_differentiation_matrix(4, 'T2U') + # test_transform(6, 1, 'fft') + # test_tau_method('T2U', -1.0, N=4, bc_val=3.0) + # test_tau_method2D('T2T', -1, nx=2**7, nz=2**6, bc_val=4.0, plotting=True) + # test_integration_matrix(5, 'T2U') + # test_integration_matrix2D(2**0, 2**2, 'T2U', 'z') + # test_differentiation_matrix2D(2**7, 2**7, 'T2U', 'mixed') + # test_integration_BC(6) + # test_filter(12, 2, 5, 'T2U') diff --git a/pySDC/tests/test_helpers/test_spectral_helper_1d_fft.py b/pySDC/tests/test_helpers/test_spectral_helper_1d_fft.py new file mode 100644 index 0000000000..15fd93464e --- /dev/null +++ b/pySDC/tests/test_helpers/test_spectral_helper_1d_fft.py @@ -0,0 +1,119 @@ +import pytest + + +@pytest.mark.base +@pytest.mark.parametrize('N', [9, 64]) +@pytest.mark.parametrize('x0', [-4, 0, 1]) +@pytest.mark.parametrize('x1', [None, 4, 8]) +def test_differentiation_matrix(N, x0, x1, plot=False): + import numpy as np + from pySDC.helpers.spectral_helper import FFTHelper + + x1 = 2 * np.pi if x1 is None else x1 + helper = FFTHelper(N=N, x0=x0, x1=x1) + + x = helper.get_1dgrid() + D = helper.get_differentiation_matrix() + + u = np.zeros_like(x) + expect = np.zeros_like(u) + + num_coef = N // 2 + f = 2 * np.pi / helper.L + coeffs = np.random.random((2, N)) + for i in range(num_coef): + u += coeffs[0, i] * np.sin(i * x * f) + u += coeffs[1, i] * np.cos(i * x * f) + expect += coeffs[0, i] * i * np.cos(i * x * f) * f + expect -= coeffs[1, i] * i * np.sin(i * x * f) * f + + u_hat = np.fft.fft(u) + Du_hat = D @ u_hat + Du = np.fft.ifft(Du_hat) + + if plot: + import matplotlib.pyplot as plt + + plt.plot(x, u) + plt.plot(x, Du) + plt.plot(x, expect, '--') + plt.show() + + assert np.allclose(expect, Du) + + +@pytest.mark.base +def test_transform(N=8): + import numpy as np + from pySDC.helpers.spectral_helper import FFTHelper + + u = np.random.random(N) + helper = FFTHelper(N=N) + u_hat = helper.transform(u) + assert np.allclose(u, helper.itransform(u_hat)) + + +@pytest.mark.base +@pytest.mark.parametrize('N', [8, 64]) +def test_integration_matrix(N, plot=False): + import numpy as np + from pySDC.helpers.spectral_helper import FFTHelper + + helper = FFTHelper(N=N) + + x = helper.get_1dgrid() + D = helper.get_integration_matrix() + + u = np.zeros_like(x) + expect = np.zeros_like(u) + + num_coef = N // 2 - 1 + coeffs = np.random.random((2, N)) + for i in range(1, num_coef + 1): + u += coeffs[0, i] * np.sin(i * x) + u += coeffs[1, i] * np.cos(i * x) + expect -= coeffs[0, i] / i * np.cos(i * x) + expect += coeffs[1, i] / i * np.sin(i * x) + + u_hat = np.fft.fft(u) + Du_hat = D @ u_hat + Du = np.fft.ifft(Du_hat) + + if plot: + import matplotlib.pyplot as plt + + plt.plot(x, u) + plt.plot(x, Du) + plt.plot(x, expect, '--') + plt.show() + + assert np.allclose(expect, Du) + + +@pytest.mark.base +@pytest.mark.parametrize('N', [4, 32]) +@pytest.mark.parametrize('v', [0, 4.78]) +def test_tau_method(N, v): + import numpy as np + from pySDC.helpers.spectral_helper import FFTHelper + + helper = FFTHelper(N=N) + + D = helper.get_differentiation_matrix() + bc_line = 0 + BC = (helper.get_Id() * 0).tolil() + BC[bc_line, :] = helper.get_integ_BC_row() + + A = D + BC + rhs = np.zeros(N) + rhs[bc_line] = v + u_hat = helper.sparse_lib.linalg.spsolve(A, rhs) + + u = helper.itransform(u_hat) + assert np.allclose(u * helper.L, v) + + +if __name__ == '__main__': + # test_differentiation_matrix(64, 4, True) + # test_integration_matrix(8, True) + test_tau_method(6, 1) diff --git a/pySDC/tests/test_helpers/test_spectral_helper_1d_ultraspherical.py b/pySDC/tests/test_helpers/test_spectral_helper_1d_ultraspherical.py new file mode 100644 index 0000000000..d5ba10afbf --- /dev/null +++ b/pySDC/tests/test_helpers/test_spectral_helper_1d_ultraspherical.py @@ -0,0 +1,102 @@ +import pytest + + +@pytest.mark.base +@pytest.mark.parametrize('N', [4, 7, 32]) +@pytest.mark.parametrize('p', [1, 2, 3, 4]) +def test_differentiation_matrix(N, p): + import numpy as np + from pySDC.helpers.spectral_helper import UltrasphericalHelper + + helper = UltrasphericalHelper(N) + x = helper.get_1dgrid() + coeffs = np.random.random(N) + + D = helper.get_differentiation_matrix(p=p) + Q = helper.get_basis_change_matrix(p_in=p, p_out=0) + + du = helper.itransform(Q @ D @ coeffs) + exact = np.polynomial.Chebyshev(coeffs).deriv(p)(x) + + assert np.allclose(exact, du) + + +@pytest.mark.base +@pytest.mark.parametrize('N', [4, 7, 32]) +def test_integration(N): + import numpy as np + from pySDC.helpers.spectral_helper import UltrasphericalHelper + + helper = UltrasphericalHelper(N) + coeffs = np.random.random(N) + coeffs[-1] = 0 + + poly = np.polynomial.Chebyshev(coeffs) + + S = helper.get_integration_matrix() + U_hat = S @ coeffs + U_hat[0] = helper.get_integration_constant(U_hat, axis=-1) + + assert np.allclose(poly.integ(lbnd=-1).coef[:-1], U_hat) + + +@pytest.mark.base +@pytest.mark.parametrize('N', [6, 33]) +@pytest.mark.parametrize('deg', [1, 3]) +@pytest.mark.parametrize('Dirichlet_recombination', [False, True]) +def test_poisson_problem(N, deg, Dirichlet_recombination): + import numpy as np + import scipy.sparse as sp + from pySDC.helpers.spectral_helper import UltrasphericalHelper + + a = 0 + b = 4 + + helper = UltrasphericalHelper(N) + x = helper.get_1dgrid() + + f = x**deg * (deg + 1) * (deg + 2) * (a - b) / 2 + + Dxx = helper.get_differentiation_matrix(p=2) + BC_l = helper.get_Dirichlet_BC_row(x=-1) + BC_r = helper.get_Dirichlet_BC_row(x=1) + P = helper.get_basis_change_matrix(p_in=0, p_out=2) + + A = Dxx.tolil() + A[-1, :] = BC_l + A[-2, :] = BC_r + A = A.tocsr() + + if Dirichlet_recombination: + Pr = helper.get_Dirichlet_recombination_matrix() + + BC_D_r = np.zeros(N) + BC_D_r[0] = 1 + BC_D_r[1] = 1 + + BC_D_l = np.zeros(N) + BC_D_l[0] = 1 + BC_D_l[1] = -1 + assert np.allclose((A @ Pr).toarray()[-1], BC_D_l) + assert np.allclose((A @ Pr).toarray()[-2], BC_D_r) + else: + Pr = helper.get_Id() + + rhs = P @ helper.transform(f) + rhs[-2] = a + rhs[-1] = b + + u_hat = Pr @ sp.linalg.spsolve(A @ Pr, rhs) + + u = helper.itransform(u_hat) + + u_exact = (a - b) / 2 * x ** (deg + 2) + (b + a) / 2 + + assert np.allclose(u_hat[deg + 3 :], 0) + assert np.allclose(u_exact, u) + + +if __name__ == '__main__': + # test_differentiation_matrix(6, 2) + test_poisson_problem(6, 1, True) + # test_integration() diff --git a/pySDC/tests/test_problems/test_Burgers.py b/pySDC/tests/test_problems/test_Burgers.py new file mode 100644 index 0000000000..ca540e0ab5 --- /dev/null +++ b/pySDC/tests/test_problems/test_Burgers.py @@ -0,0 +1,217 @@ +import pytest + + +@pytest.mark.base +@pytest.mark.parametrize('mode', ['T2T', 'T2U']) +def test_Burgers_f(mode): + import numpy as np + from pySDC.implementations.problem_classes.Burgers import Burgers1D + + P = Burgers1D(N=2**4, epsilon=8e-3, mode=mode) + + u = P.u_init + u[0] = np.sin(P.x * np.pi) + u[1] = np.cos(P.x * np.pi) * np.pi + f = P.eval_f(u) + + assert np.allclose(f.impl[0], -np.sin(P.x * np.pi) * P.epsilon * np.pi**2) + assert np.allclose(f.expl[0], -u[0] * u[1]) + assert np.allclose(f.impl[1], 0) + assert np.allclose(f.expl[1], 0) + + +@pytest.mark.base +@pytest.mark.parametrize('mode', ['T2U', 'T2T']) +def test_Burgers_solver(mode, N=2**8, plotting=False): + import numpy as np + import matplotlib.pyplot as plt + from pySDC.implementations.problem_classes.Burgers import Burgers1D + + P = Burgers1D(N=N, epsilon=1e-2, f=0, mode=mode) + + u = P.u_exact() + + def imex_euler(u, dt): + f = P.eval_f(u) + u = P.solve_system(u + dt * f.expl, dt) + return u + + small_step_size = 1e-9 + small_step = imex_euler(u.copy(), small_step_size) + + assert np.allclose(u[0], small_step[0], atol=small_step_size * 1e1) + + dt = 1e-2 + forward = imex_euler(u, dt) + f_u = P.eval_f(u) + f = P.eval_f(forward) + backward = forward - dt * (f_u.expl + f.impl) + assert np.allclose(backward[0], u[0]) + + u_exact_steady = P.u_exact(np.inf) + dt = 1e-2 + tol = 1e-4 + fig = P.get_fig() + ax = fig.get_axes()[0] + for i in range(900): + u_old = u.copy() + u = imex_euler(u, dt) + + if plotting: + P.plot(u, t=i * dt, fig=fig) + P.plot(u_exact_steady, fig=fig) + plt.pause(1e-8) + ax.cla() + + if abs(u_old[0] - u[0]) < tol: + print(f'stopping after {i} steps') + break + + assert np.allclose(u[0], u_exact_steady[0], atol=tol * 1e1) + + +@pytest.mark.base +@pytest.mark.parametrize('mode', ['T2T', 'T2U']) +@pytest.mark.parametrize('direction', ['x', 'z', 'mixed']) +def test_Burgers2D_f(mode, direction, plotting=False): + import numpy as np + from pySDC.implementations.problem_classes.Burgers import Burgers2D + + nx = 2**5 + nz = 2**5 + P = Burgers2D(nx=nx, nz=nz, epsilon=8e-3, mode=mode) + + u = P.u_init + iu, iv, iux, iuz, ivx, ivz = P.index(P.components) + + f_expect = P.f_init + + if direction == 'x': + u[iu] = np.sin(P.X * 2) + u[iux] = np.cos(P.X * 2) * 2 + f_expect.impl[iu] = -P.epsilon * u[iu] * 2**2 + elif direction == 'z': + u[iv] = P.Z**3 + 2 + u[ivz] = P.Z**2 * 3 + f_expect.impl[iv] = P.epsilon * P.Z * 6 + elif direction == 'mixed': + u[iu] = np.sin(P.X * 2) * (P.Z**3 + 2) + u[iv] = np.sin(P.X * 2) * (P.Z**3 + 2) + u[iux] = np.cos(P.X * 2) * 2 * (P.Z**3 + 2) + u[ivz] = np.sin(P.X * 2) * P.Z**2 * 3 + f_expect.impl[iu] = -P.epsilon * np.sin(P.X * 2) * 2**2 * (P.Z**3 + 2) + f_expect.impl[iv] = P.epsilon * np.sin(P.X * 2) * P.Z * 6 + + f = P.eval_f(u) + f_expect.expl[iu] = -(u[iu] * u[iux] + u[iv] * u[iuz]) + f_expect.expl[iv] = -(u[iu] * u[ivx] + u[iv] * u[ivz]) + + if plotting: + import matplotlib.pyplot as plt + + fig, axs = plt.subplots(1, 2) + i = iv + axs[0].pcolormesh(P.X, P.Z, f_expect.impl[i].real) + axs[1].pcolormesh(P.X, P.Z, (f.impl[i]).real) + plt.show() + + for comp in P.components: + i = P.index(comp) + assert np.allclose(f.expl[i], f_expect.expl[i]), f'Error in component {comp}!' + assert np.allclose(f.impl[i], f_expect.impl[i]), f'Error in component {comp}!' + + +@pytest.mark.base +@pytest.mark.parametrize('mode', ['T2U']) +def test_Burgers2D_solver(mode, nx=2**6, nz=2**6, plotting=False): + import numpy as np + import matplotlib.pyplot as plt + from pySDC.implementations.problem_classes.Burgers import Burgers2D + + try: + from mpi4py import MPI + + comm = MPI.COMM_WORLD + except ModuleNotFoundError: + comm = None + + P = Burgers2D( + nx=nx, + nz=nz, + epsilon=1e-2, + fuz=1, + fux=1, + mode=mode, + comm=comm, + left_preconditioner=True, + ) + + iu, iv, iux, iuz, ivx, ivz = P.index(P.components) + + u = P.u_exact() + + def imex_euler(u, dt): + f = P.eval_f(u) + u = P.solve_system(u + dt * f.expl, dt) + return u + + def implicit_euler(u, dt): + u = P.solve_system(u, dt) + return u + + small_step_size = 1e-8 + small_step = imex_euler(u.copy(), small_step_size) + f = P.eval_f(small_step) + + for comp in ['u', 'v', 'ux', 'uz', 'vz', 'vx']: + i = P.index(comp) + assert np.allclose(u[i], small_step[i], atol=1e-4), f'Error in component {comp}!: {abs(u[i]-small_step[i]):.2e}' + + dt = 1e0 + forward = implicit_euler(u, dt) + f = P.eval_f(forward) + backward = forward - dt * (f.impl) + + for comp in ['u', 'v']: + i = P.index(comp) + assert np.allclose( + u[i], backward[i], atol=1e-8 + ), f'Error without convection in component {comp}: {abs(u[i]-backward[i]):.2e}!' + + dt = 1e0 + f_before = P.eval_f(u) + forward = imex_euler(u, dt) + f = P.eval_f(forward) + backward = forward - dt * (f.impl + f_before.expl) + + for comp in ['u', 'v']: + i = P.index(comp) + assert np.allclose(u[i], backward[i], atol=1e-8), f'Error in component {comp}: {abs(u[i]-backward[i]):.2e}!' + + if not plotting: + return None + + dt = 1e-2 + tol = 1e-3 + + fig = P.get_fig() + u = P.u_exact(noise_level=0e-2) + for i in range(1000): + u_old = u.copy() + u = imex_euler(u, dt) + + if plotting: + P.plot(u, fig=fig, t=i * dt) + plt.pause(1e-7) + + if abs(u_old - u) < tol: + print(f'stopping after {i} steps') + break + + plt.show() + + +if __name__ == '__main__': + # test_Burgers_solver('T2U', plotting=True) + # test_Burgers2D_f('T2U', 'z', plotting=True) + test_Burgers2D_solver('T2U', plotting=True) diff --git a/pySDC/tests/test_problems/test_RayleighBenard.py b/pySDC/tests/test_problems/test_RayleighBenard.py new file mode 100644 index 0000000000..b3b4c260d5 --- /dev/null +++ b/pySDC/tests/test_problems/test_RayleighBenard.py @@ -0,0 +1,297 @@ +import pytest + + +@pytest.mark.mpi4py +@pytest.mark.parametrize('direction', ['x', 'z', 'mixed']) +@pytest.mark.parametrize('nx', [16]) +@pytest.mark.parametrize('nz', [8]) +@pytest.mark.parametrize('spectral_space', [True, False]) +def test_eval_f(nx, nz, direction, spectral_space): + import numpy as np + from pySDC.implementations.problem_classes.RayleighBenard import RayleighBenard + + P = RayleighBenard(nx=nx, nz=nz, Rayleigh=1, spectral_space=spectral_space) + iu, iv, ip, iT = P.index(['u', 'v', 'p', 'T']) + X, Z = P.X, P.Z + cos, sin = np.cos, np.sin + + kappa = (P.Rayleigh * P.Prandl) ** (-1 / 2) + nu = (P.Rayleigh / P.Prandl) ** (-1 / 2) + + if direction == 'x': + y = sin(X * np.pi) + y_x = cos(X * np.pi) * np.pi + y_xx = -sin(X * np.pi) * np.pi**2 + y_z = 0 + y_zz = 0 + elif direction == 'z': + y = Z**2 + y_x = 0 + y_xx = 0 + y_z = 2 * Z + y_zz = 2.0 + elif direction == 'mixed': + y = sin(X * np.pi) * Z**2 + y_x = cos(X * np.pi) * np.pi * Z**2 + y_xx = -sin(X * np.pi) * np.pi**2 * Z**2 + y_z = sin(X * np.pi) * 2 * Z + y_zz = sin(X * np.pi) * 2 + else: + raise NotImplementedError + + assert np.allclose(P.eval_f(P.u_init), 0), 'Non-zero time derivative in static 0 configuration' + + u = P.u_init + for i in [iu, iv, iT, ip]: + u[i][:] = y + + if spectral_space: + u = P.transform(u) + + f = P.eval_f(u) + + f_expect = P.f_init + f_expect.expl[iT] = -y * (y_x + y_z) + f_expect.impl[iT] = kappa * (y_xx + y_zz) + f_expect.expl[iu] = -y * y_z - y * y_x + f_expect.impl[iu] = -y_x + nu * (y_xx + y_zz) + f_expect.expl[iv] = -y * (y_z + y_x) + f_expect.impl[iv] = -y_z + nu * (y_xx + y_zz) + y + f_expect.impl[ip] = -(y_x + y_z) + + if spectral_space: + f.impl = P.itransform(f.impl).real + f.expl = P.itransform(f.expl).real + + for comp in ['u', 'v', 'T', 'p']: + i = P.spectral.index(comp) + assert np.allclose(f.impl[i], f_expect.impl[i]), f'Unexpected implicit function evaluation in component {comp}' + assert np.allclose(f.expl[i], f_expect.expl[i]), f'Unexpected explicit function evaluation in component {comp}' + + +@pytest.mark.mpi4py +@pytest.mark.parametrize('nx', [16]) +@pytest.mark.parametrize('nz', [4]) +@pytest.mark.parametrize('direction', ['x', 'z', 'mixed']) +def test_vorticity(nx, nz, direction): + import numpy as np + from pySDC.implementations.problem_classes.RayleighBenard import RayleighBenard + + assert nz > 3 + assert nx > 8 + + P = RayleighBenard(nx=nx, nz=nz, spectral_space=False) + iu, iv = P.index(['u', 'v']) + + u = P.u_init + + if direction == 'x': + u[iv] = np.sin(P.X * np.pi) + u[iu] = np.cos(P.X * np.pi) + expect = np.cos(P.X * np.pi) * np.pi + elif direction == 'z': + u[iv] = P.Z**2 + u[iu] = P.Z**3 + expect = 3 * P.Z**2 + elif direction == 'mixed': + u[iv] = np.sin(P.X * np.pi) * P.Z**2 + u[iu] = np.cos(P.X * np.pi) * P.Z**3 + expect = np.cos(P.X * np.pi) * np.pi * P.Z**2 + np.cos(P.X * np.pi) * 3 * P.Z**2 + else: + raise NotImplementedError + + assert np.allclose(P.compute_vorticity(u), expect) + + +@pytest.mark.mpi4py +@pytest.mark.parametrize('v', [0, 3.14]) +def test_Nusselt_numbers(v, nx=5, nz=4): + import numpy as np + from pySDC.implementations.problem_classes.RayleighBenard import RayleighBenard + + BCs = { + 'v_top': v, + 'v_bottom': v, + } + + P = RayleighBenard(nx=nx, nz=nz, BCs=BCs) + + u = P.u_exact(noise_level=0) + + nusselt = P.compute_Nusselt_numbers(u) + expect = {'V': 1 + v, 't': 1, 'b': +1 + 2 * v, 'b_no_v': 1, 't_no_v': 1} + for key in nusselt.keys(): + assert np.isclose(nusselt[key], expect[key]) + + +def test_viscous_dissipation(nx=2**5 + 1, nz=2**3 + 1): + import numpy as np + from pySDC.implementations.problem_classes.RayleighBenard import RayleighBenard + + P = RayleighBenard(nx=nx, nz=nz, spectral_space=False) + iu, iv = P.index(['u', 'v']) + X, Z = P.X, P.Z + + u = P.u_init + u[iu] = np.sin(X * np.pi) + u[iv] = Z**3 + + expect = P.u_init + expect[iu] = u[iu] * (-np.pi) ** 2 * u[iu] + expect[iv] = Z**3 * 6 * Z + + viscous_dissipation = P.compute_viscous_dissipation(u) + assert np.isclose(viscous_dissipation, abs(expect)) + + +def test_buoyancy_computation(nx=9, nz=6): + import numpy as np + from pySDC.implementations.problem_classes.RayleighBenard import RayleighBenard + + P = RayleighBenard(nx=nx, nz=nz, spectral_space=False) + iT, iv = P.index(['T', 'v']) + Z = P.Z + + u = P.u_init + u[iT] = Z - 1 + u[iv] = Z**3 + + expect = P.u_init + expect[iv] = u[iv] * P.Rayleigh * u[iT] + + buoyancy_production = P.compute_buoyancy_generation(u) + assert np.isclose(buoyancy_production, abs(expect)) + + +@pytest.mark.mpi4py +@pytest.mark.parametrize('nx', [1, 8]) +@pytest.mark.parametrize('component', ['u', 'T']) +def test_Poisson_problems(nx, component): + """ + When forgetting about convection and the time-dependent part, you get Poisson problems in u and T that are easy to solve. We check that we get the exact solution in a simple test here. + """ + import numpy as np + from pySDC.implementations.problem_classes.RayleighBenard import RayleighBenard + + BCs = { + 'u_top': 0, + 'u_bottom': 0, + 'v_top': 0, + 'v_bottom': 0, + 'T_top': 0, + 'T_bottom': 0, + } + P = RayleighBenard(nx=nx, nz=6, BCs=BCs, Rayleigh=1.0) + rhs = P.u_init + + idx = P.index(f'{component}') + + A = P.put_BCs_in_matrix(-P.L) + rhs[idx][0, 2] = 6 + rhs[idx][0, 0] = 6 + u = P.sparse_lib.linalg.spsolve(A, P.M @ rhs.flatten()).reshape(rhs.shape).real + + u_exact = P.u_init + u_exact[idx][0, 4] = 1 / 8 + u_exact[idx][0, 2] = 1 / 2 + u_exact[idx][0, 0] = -5 / 8 + + if component == 'T': + ip = P.index('p') + u_exact[ip][0, 5] = 1 / (16 * 5) + u_exact[ip][0, 3] = 5 / (16 * 5) + u_exact[ip][0, 1] = -70 / (16 * 5) + + assert np.allclose(u_exact, u) + + +@pytest.mark.mpi4py +def test_Poisson_problem_v(): + """ + Here we don't really solve a Poisson problem. v can only be constant due to the incompressibility, then we have a Possion problem in T with a linear solution and p absorbs all the rest. This is therefore mainly a test for the pressure computation. We don't test that the boundary condition is enforced because the constant pressure offset is entirely irrelevant to anything. + """ + import numpy as np + from pySDC.implementations.problem_classes.RayleighBenard import RayleighBenard + + BCs = { + 'u_top': 0, + 'u_bottom': 0, + 'v_top': 0, + 'v_bottom': 0, + 'T_top': 0, + 'T_bottom': 2, + } + P = RayleighBenard(nx=2, nz=2**3, BCs=BCs, Rayleigh=1.0) + iv = P.index('v') + + rhs_real = P.u_init + rhs_real[iv] = 32 * 6 * P.Z**5 + + rhs = P.transform(rhs_real) + rhs = (P.M @ rhs.flatten()).reshape(rhs.shape) + rhs_real = P.itransform(rhs) + + rhs_real = P.put_BCs_in_rhs(rhs_real) + rhs = P.transform(rhs_real) + + A = P.put_BCs_in_matrix(-P.L) + u = P.sparse_lib.linalg.spsolve(A, rhs.flatten()).reshape(rhs.shape).real + + u_exact_real = P.u_init + iT = P.index('T') + u_exact_real[iT] = 1 - P.Z + + ip = P.index('p') + u_exact_real[ip] = P.Z - 1 / 2 * P.Z**2 - 32 * P.Z**6 + + u_exact = P.transform(u_exact_real) + u_exact[ip, 0, 0] = u[ip, 0, 0] # nobody cares about the constant offset + assert np.allclose(u_exact, u) + + +@pytest.mark.mpi4py +def test_CFL(): + from pySDC.implementations.problem_classes.RayleighBenard import RayleighBenard, CFLLimit + import numpy as np + + P = RayleighBenard(nx=5, nz=2, spectral_space=False) + iu, iv = P.index(['u', 'v']) + + u = P.u_init + u[iu] = 2.77 + u[iv] = 1e-3 + + dt = CFLLimit.compute_max_step_size(P, u) + assert np.allclose(dt, P.X[1, 0] / u[iu]) + + u2 = P.u_init + u2[iu] = 1e-3 + u2[iv] = 3.14 + + dt2 = CFLLimit.compute_max_step_size(P, u2) + assert np.allclose(dt2, 1 / u2[iv]) + + +@pytest.mark.mpi4py +def test_Nyquist_mode_elimination(): + from pySDC.implementations.problem_classes.RayleighBenard import RayleighBenard + import numpy as np + + P = RayleighBenard(nx=32, nz=8) + u0 = P.u_exact(noise_level=1e-3) + + u = P.solve_system(u0, dt=1e-3) + + Nyquist_mode_index = P.axes[0].get_Nyquist_mode_index() + assert np.allclose(u[:, Nyquist_mode_index, :], 0) + + +if __name__ == '__main__': + # test_eval_f(2**0, 2**2, 'z', True) + # test_Poisson_problem(1, 'T') + test_Poisson_problem_v() + # test_Nusselt_numbers(1) + # test_buoyancy_computation() + # test_viscous_dissipation() + # test_CFL() + # test_Nyquist_mode_elimination() diff --git a/pySDC/tests/test_problems/test_heat_chebychev.py b/pySDC/tests/test_problems/test_heat_chebychev.py new file mode 100644 index 0000000000..8a5f150d7b --- /dev/null +++ b/pySDC/tests/test_problems/test_heat_chebychev.py @@ -0,0 +1,152 @@ +import pytest + + +@pytest.mark.base +@pytest.mark.parametrize('a', [0, 19]) +@pytest.mark.parametrize('b', [0, 66]) +@pytest.mark.parametrize('f', [0, 1]) +@pytest.mark.parametrize('noise', [0, 1e-3]) +@pytest.mark.parametrize('use_ultraspherical', [True, False]) +@pytest.mark.parametrize('spectral_space', [True, False]) +def test_heat1d_chebychev(a, b, f, noise, use_ultraspherical, spectral_space, nvars=2**4): + import numpy as np + + if use_ultraspherical: + from pySDC.implementations.problem_classes.HeatEquation_Chebychev import Heat1DUltraspherical as problem_class + else: + from pySDC.implementations.problem_classes.HeatEquation_Chebychev import Heat1DChebychev as problem_class + + P = problem_class( + nvars=nvars, + a=a, + b=b, + f=f, + nu=1e-2, + left_preconditioner=False, + debug=True, + spectral_space=spectral_space, + ) + + u0 = P.u_exact(0, noise=noise) + dt = 1e-1 + u = P.solve_system(u0, dt) + u02 = u - dt * P.eval_f(u) + + if noise == 0: + assert np.allclose(u, P.u_exact(dt), atol=1e-2), 'Error in solver' + + if noise > 0 and use_ultraspherical: + tol = 2e-2 + elif noise > 0: + tol = 1e-4 + else: + tol = 1e-8 + assert np.allclose(u0[0], u02[0], atol=tol), 'Error in eval_f' + + +@pytest.mark.base +@pytest.mark.parametrize('a', [0, 7]) +@pytest.mark.parametrize('b', [0, -2.77]) +@pytest.mark.parametrize('c', [0, 3.1415]) +@pytest.mark.parametrize('fx', [2, 1]) +@pytest.mark.parametrize('fy', [2, 1]) +@pytest.mark.parametrize('base_x', ['fft', 'chebychev', 'ultraspherical']) +@pytest.mark.parametrize('base_y', ['fft', 'chebychev', 'ultraspherical']) +def test_heat2d_chebychev(a, b, c, fx, fy, base_x, base_y, nx=2**5 + 1, ny=2**5 + 1): + import numpy as np + + if base_x == 'ultraspherical' or base_y == 'ultraspherical': + from pySDC.implementations.problem_classes.HeatEquation_Chebychev import Heat2DUltraspherical as problem_class + else: + from pySDC.implementations.problem_classes.HeatEquation_Chebychev import Heat2DChebychev as problem_class + + if base_x == 'chebychev' and base_y == 'ultraspherical' or base_y == 'chebychev' and base_x == 'ultraspherical': + return None + + if base_y == 'fft' and (b != c): + return None + if base_x == 'fft' and (b != a): + return None + + P = problem_class( + nx=nx, + ny=ny, + a=a, + b=b, + c=c, + fx=fx, + fy=fy, + base_x=base_x, + base_y=base_y, + nu=1e-3, + ) + + u0 = P.u_exact(0) + dt = 1e-2 + u = P.solve_system(u0, dt) + u02 = u - dt * P.eval_f(u) + + assert np.allclose( + u, P.u_exact(dt), atol=1e-3 + ), f'Error in solver larger than expected, got {abs((u - P.u_exact(dt))):.2e}' + assert np.allclose(u0[0], u02[0], atol=1e-4), 'Error in eval_f' + + +def test_SDC(): + import numpy as np + from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit + from pySDC.implementations.problem_classes.HeatEquation_Chebychev import Heat1DChebychev + from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI + from pySDC.implementations.problem_classes.generic_spectral import compute_residual_DAE + from pySDC.helpers.stats_helper import get_sorted + from pySDC.implementations.hooks.log_work import LogSDCIterations + + generic_implicit.compute_residual = compute_residual_DAE + + dt = 1e-1 + Tend = 2 * dt + + level_params = {} + level_params['dt'] = dt + level_params['restol'] = 1e-10 + + sweeper_params = {} + sweeper_params['quad_type'] = 'RADAU-RIGHT' + sweeper_params['num_nodes'] = 4 + sweeper_params['QI'] = 'LU' + + problem_params = {} + + step_params = {} + step_params['maxiter'] = 99 + + controller_params = {} + controller_params['logger_level'] = 15 + controller_params['hook_class'] = [LogSDCIterations] + + description = {} + description['problem_class'] = Heat1DChebychev + description['problem_params'] = problem_params + description['sweeper_class'] = generic_implicit + description['sweeper_params'] = sweeper_params + description['level_params'] = level_params + description['step_params'] = step_params + + controller = controller_nonMPI(num_procs=1, controller_params=controller_params, description=description) + + P = controller.MS[0].levels[0].prob + uinit = P.u_exact(t=0) + + uend, stats = controller.run(u0=uinit, t0=0.0, Tend=Tend) + u_exact = P.u_exact(t=Tend) + assert np.allclose(uend, u_exact, atol=1e-10) + + k = get_sorted(stats, type='k') + assert all(me[1] < step_params['maxiter'] for me in k) + assert all(me[1] > 0 for me in k) + + +if __name__ == '__main__': + test_SDC() + # test_heat1d_chebychev(1, 0, 1, 0e-3, True, True, 2**4) + # test_heat2d_chebychev(0, 0, 0, 2, 2, 'ultraspherical', 'fft', 2**6, 2**6)