Fix for spatial preconditioners in generic spectral problem class (#568)

* Bugfix for spatial preconditioners in generic spectral problem class

* Skip test on python 3.9

Author: brownbaerchen

pySDC/implementations/problem_classes/generic_spectral.py

@@ -84,9 +84,11 @@ def __init__(
             debug (bool): Make additional tests at extra computational cost
         """
         solver_args = {} if solver_args is None else solver_args
+
         preconditioner_args = {} if preconditioner_args is None else preconditioner_args
         preconditioner_args['drop_tol'] = preconditioner_args.get('drop_tol', 1e-3)
         preconditioner_args['fill_factor'] = preconditioner_args.get('fill_factor', 100)
+
         self._makeAttributeAndRegister(
             'max_cached_factorizations',
             'useGPU',
@@ -209,7 +211,7 @@ def setup_preconditioner(self, Dirichlet_recombination=True, left_preconditioner
 
             self.Pl = self.spectral.sparse_lib.csc_matrix(R)
 
-        if Dirichlet_recombination and type(self.axes[-1]).__name__ in ['ChebychevHelper, Ultraspherical']:
+        if Dirichlet_recombination and type(self.axes[-1]).__name__ in ['ChebychevHelper', 'UltrasphericalHelper']:
             _Pr = self.spectral.get_Dirichlet_recombination_matrix(axis=-1)
         else:
             _Pr = Id
@@ -234,18 +236,21 @@ def solve_system(self, rhs, dt, u0=None, *args, skip_itransform=False, **kwargs)
         if self.spectral_space:
             rhs_hat = rhs.copy()
             if u0 is not None:
-                u0_hat = self.Pr.T @ u0.copy().flatten()
+                u0_hat = u0.copy().flatten()
             else:
                 u0_hat = None
         else:
             rhs_hat = self.spectral.transform(rhs)
             if u0 is not None:
-                u0_hat = self.Pr.T @ self.spectral.transform(u0).flatten()
+                u0_hat = self.spectral.transform(u0).flatten()
             else:
                 u0_hat = None
 
-        if self.useGPU:
-            self.xp.cuda.Device().synchronize()
+        # apply inverse right preconditioner to initial guess
+        if u0_hat is not None and 'direct' not in self.solver_type:
+            if not hasattr(self, '_Pr_inv'):
+                self._PR_inv = self.linalg.splu(self.Pr.astype(complex)).solve
+            u0_hat[...] = self._PR_inv(u0_hat)
 
         rhs_hat = (self.M @ rhs_hat.flatten()).reshape(rhs_hat.shape)
         rhs_hat = self.spectral.put_BCs_in_rhs_hat(rhs_hat)
@@ -255,17 +260,13 @@ def solve_system(self, rhs, dt, u0=None, *args, skip_itransform=False, **kwargs)
             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
+            # if A.shape[0] < 200e20:
+            #     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()
+            #     # M = self.spectral.put_BCs_in_matrix(self.L.copy())
+            #     M = A  # self.L
+            #     im = plt.spy(M)
+            #     plt.show()
 
         if 'ilu' in self.solver_type.lower():
             if dt not in self.cached_factorizations.keys():
@@ -330,9 +331,6 @@ def solve_system(self, rhs, dt, u0=None, *args, skip_itransform=False, **kwargs)
         sol_hat = self.spectral.u_init_forward
         sol_hat[...] = (self.Pr @ _sol_hat).reshape(sol_hat.shape)
 
-        if self.useGPU:
-            self.xp.cuda.Device().synchronize()
-
         if self.spectral_space:
             return sol_hat
         else:

pySDC/tests/test_problems/test_RayleighBenard3D.py

@@ -1,4 +1,5 @@
 import pytest
+import sys
 
 
 @pytest.mark.mpi4py
@@ -202,8 +203,37 @@ def test_libraries():
     assert P.axes[2].fft_lib == fft
 
 
+@pytest.mark.mpi4py
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="requires python3.10 or higher")
+@pytest.mark.parametrize('preconditioning', [True, False])
+def test_banded_matrix(preconditioning):
+    from pySDC.implementations.problem_classes.RayleighBenard3D import RayleighBenard3D
+
+    P = RayleighBenard3D(
+        nx=16, ny=16, nz=16, left_preconditioner=preconditioning, Dirichlet_recombination=preconditioning
+    )
+    sp = P.spectral.linalg
+
+    A = P.M + P.L
+    A = P.Pl @ P.spectral.put_BCs_in_matrix(A) @ P.Pr
+
+    bandwidth = sp.spbandwidth(A)
+    if preconditioning:
+        assert all(250 * me < A.shape[0] for me in bandwidth), 'Preconditioned matrix is not banded!'
+    else:
+        assert all(2 * me > A.shape[0] for me in bandwidth), 'Unpreconditioned matrix is unexpectedly banded!'
+
+    LU = sp.splu(A)
+    for me in [LU.L, LU.U]:
+
+        assert max(sp.spbandwidth(me)) <= max(
+            bandwidth
+        ), 'One-sided bandwidth of LU decomposition is larger than that of the full matrix!'
+
+
 if __name__ == '__main__':
     # test_eval_f(2**2, 2**1, 'x', False)
     # test_libraries()
     # test_Poisson_problems(4, 'u')
-    test_Poisson_problem_w()
+    # test_Poisson_problem_w()
+    test_banded_matrix(False)

pySDC/tests/test_problems/test_generic_spectral.py

@@ -0,0 +1,57 @@
+import pytest
+
+
+@pytest.mark.base
+@pytest.mark.parametrize('use_ultraspherical', [True, False])
+@pytest.mark.parametrize('spectral_space', [True, False])
+@pytest.mark.parametrize('solver_type', ['bicgstab'])
+@pytest.mark.parametrize('left_preconditioner', [True, False])
+@pytest.mark.parametrize('Dirichlet_recombination', [True, False])
+def test_initial_guess(
+    use_ultraspherical, spectral_space, solver_type, left_preconditioner, Dirichlet_recombination, 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
+
+    params = {
+        'nvars': nvars,
+        'a': 0,
+        'b': 1,
+        'f': 0,
+        'nu': 1e-2,
+        'debug': True,
+        'spectral_space': spectral_space,
+        'solver_type': solver_type,
+        'solver_args': {'rtol': 1e-12},
+        'left_preconditioner': left_preconditioner,
+        'Dirichlet_recombination': Dirichlet_recombination,
+    }
+
+    P = problem_class(
+        **params,
+    )
+
+    u0 = P.u_exact(0)
+    dt = 1e-1
+    u = P.solve_system(u0, dt)
+
+    if spectral_space:
+        error = max(abs((P.M + dt * P.L) @ u.flatten() - P.M @ u0.flatten()))
+        assert error < 1e-12, error
+
+    iter_1 = P.work_counters[P.solver_type].niter
+    assert iter_1 > 0
+
+    u2 = P.solve_system(u0, u0=u, dt=dt)
+    iter_2 = P.work_counters[P.solver_type].niter - iter_1
+
+    assert iter_2 == 0, f'Did {iter_2} extra iterations after doing {iter_1} iterations first time'
+    assert np.allclose(u, u2)
+
+
+if __name__ == '__main__':
+    test_initial_guess(False, True, 'bicgstab', True, True, 2**4)

pySDC/tests/test_problems/test_heat_chebychev.py

@@ -8,7 +8,7 @@
 @pytest.mark.parametrize('noise', [0, 1e-3])
 @pytest.mark.parametrize('use_ultraspherical', [True, False])
 @pytest.mark.parametrize('spectral_space', [True, False])
-@pytest.mark.parametrize('solver_type', ['cached_direct', 'direct', 'gmres', 'bicgstab', 'gmres+ilu', 'bicgstab+ilu'])
+@pytest.mark.parametrize('solver_type', ['cached_direct', 'direct', 'bicgstab', 'gmres+ilu', 'bicgstab+ilu'])
 def test_heat1d_chebychev(a, b, f, noise, use_ultraspherical, spectral_space, solver_type, nvars=2**4):
     import numpy as np
 
@@ -150,6 +150,6 @@ def test_SDC():
 
 
 if __name__ == '__main__':
-    test_SDC()
-    # test_heat1d_chebychev(1, 0, 1, 0e-3, True, True, 2**4)
+    # test_SDC()
+    test_heat1d_chebychev(1, 0, 1, 0e-3, True, True, 'bicgstab', 2**4)
     # test_heat2d_chebychev(0, 0, 0, 2, 2, 'ultraspherical', 'fft', 2**6, 2**6)