diff --git a/pySDC/helpers/spectral_helper.py b/pySDC/helpers/spectral_helper.py index 6491b99f5b..b1b147a26d 100644 --- a/pySDC/helpers/spectral_helper.py +++ b/pySDC/helpers/spectral_helper.py @@ -1024,16 +1024,22 @@ def index(self, name): else: raise NotImplementedError(f'Don\'t know how to compute index for {type(name)=}') - def get_empty_operator_matrix(self): + def get_empty_operator_matrix(self, diag=False): """ Return a matrix of operators to be filled with the connections between the solution components. + Args: + diag (bool): Whether operator is block-diagonal + Returns: list containing sparse zeros """ S = len(self.components) O = self.get_Id() * 0 - return [[O for _ in range(S)] for _ in range(S)] + if diag: + return [O for _ in range(S)] + else: + return [[O for _ in range(S)] for _ in range(S)] def get_BC(self, axis, kind, line=-1, scalar=False, **kwargs): """ @@ -1296,7 +1302,7 @@ def put_BCs_in_rhs(self, rhs): return rhs - def add_equation_lhs(self, A, equation, relations): + def add_equation_lhs(self, A, equation, relations, diag=False): """ 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. @@ -1331,11 +1337,16 @@ def add_equation_lhs(self, A, equation, relations): 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 + diag (bool): Whether operator is block-diagonal """ for k, v in relations.items(): - A[self.index(equation)][self.index(k)] = v + if diag: + assert k == equation, 'You are trying to put a non-diagonal equation into a diagonal operator' + A[self.index(equation)] = v + else: + A[self.index(equation)][self.index(k)] = v - def convert_operator_matrix_to_operator(self, M): + def convert_operator_matrix_to_operator(self, M, diag=False): """ 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. @@ -1347,9 +1358,14 @@ def convert_operator_matrix_to_operator(self, M): sparse linear operator """ if len(self.components) == 1: - return M[0][0] + if diag: + return M[0] + else: + return M[0][0] + elif diag: + return self.sparse_lib.block_diag(M, format='csc') else: - return self.sparse_lib.bmat(M, format='csc') + return self.sparse_lib.block_array(M, format='csc') def get_wavenumbers(self): """ diff --git a/pySDC/implementations/problem_classes/RayleighBenard.py b/pySDC/implementations/problem_classes/RayleighBenard.py index e69a86d496..509bf3ca14 100644 --- a/pySDC/implementations/problem_classes/RayleighBenard.py +++ b/pySDC/implementations/problem_classes/RayleighBenard.py @@ -213,8 +213,12 @@ def eval_f(self, u, *args, **kwargs): # start by computing derivatives if not hasattr(self, '_Dx_expanded') or not hasattr(self, '_Dz_expanded'): - self._Dx_expanded = self._setup_operator({'u': {'u': Dx}, 'v': {'v': Dx}, 'T': {'T': Dx}, 'p': {}}) - self._Dz_expanded = self._setup_operator({'u': {'u': Dz}, 'v': {'v': Dz}, 'T': {'T': Dz}, 'p': {}}) + self._Dx_expanded = self._setup_operator( + {'u': {'u': Dx}, 'v': {'v': Dx}, 'T': {'T': Dx}, 'p': {}}, diag=True + ) + self._Dz_expanded = self._setup_operator( + {'u': {'u': Dz}, 'v': {'v': Dz}, 'T': {'T': Dz}, 'p': {}}, diag=True + ) Dx_u_hat = (self._Dx_expanded @ u_hat.flatten()).reshape(u_hat.shape) Dz_u_hat = (self._Dz_expanded @ u_hat.flatten()).reshape(u_hat.shape) diff --git a/pySDC/implementations/problem_classes/generic_spectral.py b/pySDC/implementations/problem_classes/generic_spectral.py index 0410cb3ad9..c14e60a22a 100644 --- a/pySDC/implementations/problem_classes/generic_spectral.py +++ b/pySDC/implementations/problem_classes/generic_spectral.py @@ -136,20 +136,21 @@ def __getattr__(self, name): """ return getattr(self.spectral, name) - def _setup_operator(self, LHS): + def _setup_operator(self, LHS, diag=False): """ 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 + diag (bool): Whether operator is block-diagonal Returns: sparse linear operator """ - operator = self.spectral.get_empty_operator_matrix() + operator = self.spectral.get_empty_operator_matrix(diag=diag) for line, equation in LHS.items(): - self.spectral.add_equation_lhs(operator, line, equation) - return self.spectral.convert_operator_matrix_to_operator(operator) + self.spectral.add_equation_lhs(operator, line, equation, diag=diag) + return self.spectral.convert_operator_matrix_to_operator(operator, diag=diag) def setup_L(self, LHS): """ @@ -171,13 +172,13 @@ def setup_L(self, LHS): """ self.L = self._setup_operator(LHS) - def setup_M(self, LHS): + def setup_M(self, LHS, diag=True): ''' 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) + self.M = self._setup_operator(LHS, diag=diag) def setup_preconditioner(self, Dirichlet_recombination=True, left_preconditioner=True): """ @@ -192,7 +193,7 @@ def setup_preconditioner(self, Dirichlet_recombination=True, left_preconditioner Id = sp.eye(N) Pl_lhs = {comp: {comp: Id} for comp in self.components} - self.Pl = self._setup_operator(Pl_lhs) + self.Pl = self._setup_operator(Pl_lhs, diag=True) if left_preconditioner: # reverse Kronecker product @@ -214,7 +215,7 @@ def setup_preconditioner(self, Dirichlet_recombination=True, left_preconditioner _Pr = Id Pr_lhs = {comp: {comp: _Pr} for comp in self.components} - self.Pr = self._setup_operator(Pr_lhs) @ self.Pl.T + self.Pr = self._setup_operator(Pr_lhs, diag=True) @ self.Pl.T def solve_system(self, rhs, dt, u0=None, *args, skip_itransform=False, **kwargs): """ diff --git a/pySDC/tests/test_helpers/test_spectral_helper.py b/pySDC/tests/test_helpers/test_spectral_helper.py index 639a4e71af..7df1eff3ac 100644 --- a/pySDC/tests/test_helpers/test_spectral_helper.py +++ b/pySDC/tests/test_helpers/test_spectral_helper.py @@ -491,9 +491,10 @@ def test_tau_method2D(variant, nz, nx, bc_val, bc=-1, useMPI=False, plotting=Fal 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) + diag = True + _A = helper.get_empty_operator_matrix(diag=diag) + helper.add_equation_lhs(_A, 'u', {'u': Dz - Dxx * 1e-1 - Dx}, diag=diag) + A = helper.convert_operator_matrix_to_operator(_A, diag=diag) # prepare system to solve A = helper.put_BCs_in_matrix(A) @@ -608,6 +609,34 @@ def function(): assert track[0] == 0, "possible memory leak with the @cache" +@pytest.mark.base +def test_block_diagonal_operators(N=16): + from pySDC.helpers.spectral_helper import SpectralHelper + import numpy as np + + helper = SpectralHelper(comm=None, debug=True) + helper.add_axis('fft', N=N) + helper.add_axis('cheby', N=N) + helper.add_component(['u', 'v']) + helper.setup_fft() + + # generate matrices + Dz = helper.get_differentiation_matrix(axes=(1,)) + Dx = helper.get_differentiation_matrix(axes=(0,)) + + def get_operator(diag): + _A = helper.get_empty_operator_matrix(diag=diag) + helper.add_equation_lhs(_A, 'u', {'u': Dx}, diag=diag) + helper.add_equation_lhs(_A, 'v', {'v': Dz}, diag=diag) + return helper.convert_operator_matrix_to_operator(_A, diag=diag) + + AD = get_operator(True) + A = get_operator(False) + + assert np.allclose(A.toarray(), AD.toarray()), 'Operators don\'t match' + assert A.data.nbytes > AD.data.nbytes, 'Block diagonal operator did not conserve memory over general operator' + + 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(',')) @@ -642,9 +671,10 @@ def function(): # 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_tau_method2D('T2U', 2**8, 2**8, -2, plotting=True) # test_filter(6, 6, (0,)) # _test_transform_dealias('fft', 'cheby', (-1, -2)) + test_block_diagonal_operators() else: raise NotImplementedError print('done')