Merge remote-tracking branch 'upstream/master' into merge_refactor

Author: brownbaerchen

pySDC/helpers/spectral_helper.py

@@ -1090,16 +1090,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):
         """
@@ -1369,7 +1375,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.
@@ -1404,11 +1410,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.
@@ -1420,9 +1431,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):
         """