diff --git a/pySDC/implementations/problem_classes/TestEquation_0D.py b/pySDC/implementations/problem_classes/TestEquation_0D.py index 5262b165fd..155f78f8b5 100644 --- a/pySDC/implementations/problem_classes/TestEquation_0D.py +++ b/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) diff --git a/pySDC/implementations/sweeper_classes/imex_1st_order.py b/pySDC/implementations/sweeper_classes/imex_1st_order.py index 12600c64d7..181776b635 100644 --- a/pySDC/implementations/sweeper_classes/imex_1st_order.py +++ b/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 diff --git a/pySDC/implementations/sweeper_classes/imex_1st_order_MPI.py b/pySDC/implementations/sweeper_classes/imex_1st_order_MPI.py index b32d9477a0..cdadd721f5 100644 --- a/pySDC/implementations/sweeper_classes/imex_1st_order_MPI.py +++ b/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() diff --git a/pySDC/tests/test_sweepers/test_preconditioners.py b/pySDC/tests/test_sweepers/test_preconditioners.py index e19fe54d8d..dab7575fee 100644 --- a/pySDC/tests/test_sweepers/test_preconditioners.py +++ b/pySDC/tests/test_sweepers/test_preconditioners.py @@ -19,7 +19,7 @@ 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) @@ -27,7 +27,7 @@ def test_MIN_SR(node_type, quad_type, M): # 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)