Put BCs in rhs with fewer transforms

pySDC/helpers/spectral_helper.py

@@ -877,6 +877,10 @@ def setup_BCs(self):
         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
@@ -906,9 +910,37 @@ def put_BCs_in_matrix(self, A):
         """
         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
+                    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
@@ -918,7 +950,7 @@ def put_BCs_in_rhs(self, rhs):
         """
         assert rhs.ndim > 1, 'rhs must not be flattened here!'
 
-        ndim = len(self.axes)
+        ndim = self.ndim
 
         for axis in range(ndim):
             _rhs_hat = self.transform(rhs, axes=(axis - ndim,))

pySDC/implementations/problem_classes/generic_spectral.py

@@ -176,8 +176,6 @@ def solve_system(self, rhs, dt, u0=None, *args, skip_itransform=False, **kwargs)
         Solve (M + dt*L)u=rhs. This requires that you setup the operators before using the functions ``GenericSpectralLinear.setup_L`` and ``GenericSpectralLinear.setup_M``. Note that the mass matrix need not be invertible, as long as (M + dt*L) is. This allows to solve some differential algebraic equations.
 
         Note that in implicit Euler, the right hand side will be composed of the initial conditions. We don't want that in lines that don't depend on time. Therefore, we multiply the right hand side by the mass matrix. This means you can only do algebraic conditions that add up to zero. But you can easily overload this function with something more generic if needed.
-
-        We use a tau method to enforce boundary conditions in Chebychov methods. This means we replace a line in the system matrix by the polynomials evaluated at a boundary and put the value we want there in the rhs at the respective position. Since we have to do that in spectral space along only the axis we want to apply the boundary condition to, we transform back to real space after applying the mass matrix, and then transform only along one axis, apply the boundary conditions and transform back. Then we transform along all dimensions again. If you desire speed, you may wish to overload this function with something less generic that avoids a few transformations.
         """
         if dt == 0:
             return rhs
@@ -188,10 +186,8 @@ def solve_system(self, rhs, dt, u0=None, *args, skip_itransform=False, **kwargs)
 
         rhs_hat = self.spectral.transform(rhs)
         rhs_hat = (self.M @ rhs_hat.flatten()).reshape(rhs_hat.shape)
-        rhs = self.spectral.itransform(rhs_hat).real
-
-        rhs = self.spectral.put_BCs_in_rhs(rhs)
-        rhs_hat = self.Pl @ self.spectral.transform(rhs).flatten()
+        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

pySDC/tests/test_helpers/test_spectral_helper.py

@@ -563,8 +563,9 @@ def test_tau_method2D(variant, nz, nx, bc_val, bc=-1, useMPI=False, plotting=Fal
 
     # prepare system to solve
     A = helper.put_BCs_in_matrix(A)
-    rhs = helper.put_BCs_in_rhs(helper.u_init)
-    rhs_hat = helper.transform(rhs, axes=(-1, -2))
+    # rhs = helper.put_BCs_in_rhs(helper.u_init)
+    # rhs_hat = helper.transform(rhs, axes=(-1, -2))
+    rhs_hat = helper.put_BCs_in_rhs_hat(helper.u_init_forward)
 
     # solve the system
     sol_hat = helper.u_init_forward
@@ -653,9 +654,9 @@ def test_dealias_MPI(num_procs, axis, bx, bz, nx=32, nz=64, **kwargs):
         # 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**2, 2**2, -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_transform_dealias('fft', 'cheby', (-1, -2))
     else:
         raise NotImplementedError
     print('done')

pySDC/tests/test_problems/test_heat_chebychev.py

@@ -138,6 +138,6 @@ def test_SDC():
 
 
 if __name__ == '__main__':
-    # test_SDC()
+    test_SDC()
     # test_heat1d_chebychev(1, 0, 1, 1e-3, True, 2**6)
-    test_heat2d_chebychev(0, 0, 0, 2, 2, 'ultraspherical', 'fft', 2**6, 2**6)
+    # test_heat2d_chebychev(0, 0, 0, 2, 2, 'ultraspherical', 'fft', 2**6, 2**6)