Flex preconditioners in IMEX sweepers (#572)

* Added test for flex preconditioner

* Implemented FLEX preconditioners in IMEX sweepers

Author: brownbaerchen

pySDC/implementations/problem_classes/TestEquation_0D.py

@@ -173,11 +173,8 @@ def __init__(self, lambdas_implicit=None, lambdas_explicit=None, u0=0.0):
                 [[complex(re[i], im[j]) for i in range(len(re))] for j in range(len(im))]
             ).reshape((len(re) * len(im)))
         if lambdas_explicit is None:
-            re = self.xp.linspace(-30, 19, 50)
-            im = self.xp.linspace(-50, 49, 50)
-            lambdas_implicit = self.xp.array(
-                [[complex(re[i], im[j]) for i in range(len(re))] for j in range(len(im))]
-            ).reshape((len(re) * len(im)))
+            lambdas_explicit = lambdas_implicit.copy()
+
         lambdas_implicit = self.xp.asarray(lambdas_implicit)
         lambdas_explicit = self.xp.asarray(lambdas_explicit)
 

pySDC/implementations/sweeper_classes/imex_1st_order.py

@@ -72,6 +72,9 @@ def update_nodes(self):
         # get number of collocation nodes for easier access
         M = self.coll.num_nodes
 
+        # update the MIN-SR-FLEX preconditioner
+        self.updateVariableCoeffs(L.status.sweep)
+
         # gather all terms which are known already (e.g. from the previous iteration)
         # this corresponds to u0 + QF(u^k) - QIFI(u^k) - QEFE(u^k) + tau
 

pySDC/implementations/sweeper_classes/imex_1st_order_MPI.py

@@ -55,6 +55,9 @@ def update_nodes(self):
         # gather all terms which are known already (e.g. from the previous iteration)
         # this corresponds to u0 + QF(u^k) - QdF(u^k) + tau
 
+        # update the MIN-SR-FLEX preconditioner
+        self.updateVariableCoeffs(L.status.sweep)
+
         # get QF(u^k)
         rhs = self.integrate()
 

pySDC/tests/test_sweepers/test_preconditioners.py

@@ -19,15 +19,15 @@ def test_MIN_SR(node_type, quad_type, M):
 
     # Check non-stiff limit
     QDelta = sweeper.get_Qdelta_implicit('MIN-SR-NS')[1:, 1:]
-    assert np.all(np.diag(np.diag(QDelta)) == QDelta), "no diagonal QDelta"
+    assert np.all(np.diag(np.diag(QDelta)) == QDelta), "QDelta not diagonal"
     K = Q - QDelta
     Km = np.linalg.matrix_power(K, M)
     nilpotency = np.linalg.norm(Km, ord=np.inf)
     assert nilpotency < 1e-10, "Q-QDelta not nilpotent " f"(M={M}, norm={nilpotency})"
 
     # Check stiff limit
     QDelta = sweeper.get_Qdelta_implicit('MIN-SR-S')[1:, 1:]
-    assert np.all(np.diag(np.diag(QDelta)) == QDelta), "no diagonal QDelta"
+    assert np.all(np.diag(np.diag(QDelta)) == QDelta), "QDelta not diagonal"
 
     if params['quad_type'] in ['LOBATTO', 'RADAU-LEFT']:
         QDelta = np.diag(1 / np.diag(QDelta[1:, 1:]))
@@ -41,6 +41,92 @@ def test_MIN_SR(node_type, quad_type, M):
     assert nilpotency < 1e-10, "I-QDelta^{-1}Q not nilpotent " f"(M={M}, norm={nilpotency})"
 
 
+@pytest.mark.base
+@pytest.mark.parametrize("node_type", node_types)
+@pytest.mark.parametrize("quad_type", quad_types)
+@pytest.mark.parametrize("M", num_nodes)
+def test_MIN_SR_FLEX(node_type, quad_type, M):
+    params = {'num_nodes': M, 'quad_type': quad_type, 'node_type': node_type}
+    sweeper = Sweeper(params)
+
+    start_idx = 1
+    for i in range(M):
+        if sweeper.coll.nodes[i] == 0:
+            start_idx += 1
+        else:
+            break
+
+    Q = sweeper.coll.Qmat[start_idx:, start_idx:]
+
+    QDelta = [sweeper.get_Qdelta_implicit('MIN-SR-FLEX', k=i + 1)[start_idx:, start_idx:] for i in range(M)]
+    for QD in QDelta:
+        assert np.all(np.diag(np.diag(QD)) == QD), "QDelta not diagonal"
+
+    I = np.eye(M + 1 - start_idx)
+    K = np.eye(M + 1 - start_idx)
+    for QD in QDelta:
+        K = (I - np.linalg.inv(QD) @ Q) @ K
+
+    nilpotency = np.linalg.norm(K, ord=np.inf)
+    assert (
+        nilpotency < 1e-10
+    ), f"Applying FLEX preconditioner does not give nilpotent SDC iteration matrix after {M} iterations! (M={M}, norm={nilpotency})"
+
+
+@pytest.mark.base
+@pytest.mark.parametrize('imex', [True, False])
+@pytest.mark.parametrize('num_nodes', num_nodes)
+def test_FLEX_preconditioner_in_sweepers(imex, num_nodes, MPI=False):
+    from pySDC.core.level import Level
+
+    if imex:
+        from pySDC.implementations.problem_classes.TestEquation_0D import test_equation_IMEX as problem_class
+
+        if MPI:
+            from pySDC.implementations.sweeper_classes.imex_1st_order_MPI import imex_1st_order_MPI as sweeper_class
+        else:
+            from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order as sweeper_class
+    else:
+        from pySDC.implementations.problem_classes.TestEquation_0D import testequation0d as problem_class
+
+        if MPI:
+            from pySDC.implementations.sweeper_classes.generic_implicit_MPI import generic_implicit_MPI as sweeper_class
+        else:
+            from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit as sweeper_class
+
+    sweeper_params = {'quad_type': 'RADAU-RIGHT', 'num_nodes': num_nodes, 'QI': 'MIN-SR-FLEX', 'QE': 'PIC'}
+    if MPI:
+        from mpi4py import MPI
+
+        sweeper_params['comm'] = MPI.COMM_WORLD
+    level_params = {'nsweeps': num_nodes, 'dt': 1}
+
+    lvl = Level(problem_class, {}, sweeper_class, sweeper_params, level_params, 0)
+
+    lvl.status.unlocked = True
+    lvl.u[0] = lvl.prob.u_exact(0)
+    lvl.status.time = 0
+
+    sweep = lvl.sweep
+    sweep.predict()
+
+    for k in range(1, level_params['nsweeps'] + 1):
+        lvl.status.sweep = k
+        sweep.update_nodes()
+        assert np.allclose(
+            sweep.QI, sweep.get_Qdelta_implicit(sweeper_params['QI'], k)
+        ), f'Got incorrect FLEX preconditioner in sweep {k}'
+
+
+@pytest.mark.mpi4py
+@pytest.mark.parametrize('imex', [True, False])
+@pytest.mark.mpi(ranks=[3])
+def test_FLEX_preconditioner_in_MPI_sweepers(mpi_ranks, imex):
+    from mpi4py import MPI
+
+    test_FLEX_preconditioner_in_sweepers(imex, num_nodes=MPI.COMM_WORLD.size, MPI=True)
+
+
 @pytest.mark.base
 @pytest.mark.parametrize("node_type", node_types)
 @pytest.mark.parametrize("quad_type", quad_types)
@@ -122,8 +208,12 @@ def test_PIC(node_type, quad_type, M):
 
 
 if __name__ == '__main__':
+    test_MIN_SR_FLEX('LEGENDRE', 'LOBATTO', 4)
+
     test_MIN_SR('LEGENDRE', 'RADAU-RIGHT', 4)
     test_MIN_SR('EQUID', 'LOBATTO', 5)
 
     test_LU('LEGENDRE', 'RADAU-RIGHT', 4)
     test_LU('EQUID', 'LOBATTO', 5)
+
+    test_FLEX_preconditioner_in_sweepers(True)