Reduced memory footprint of spectral discretizations via block-diagonal sparse operators

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):
         """

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)
 

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):
         """

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')