Added multi-level SDC for Firedrake and Gusto coupling (#520)

* Implemented MLSDC for Firedrake and Gusto coupling

* Removed unused `equation` argument from `pySDC_integrator`

Author: brownbaerchen

pySDC/helpers/pySDC_as_gusto_time_discretization.py

@@ -38,7 +38,6 @@ class pySDC_integrator(TimeDiscretisation):
 
     def __init__(
         self,
-        equation,
         description,
         controller_params,
         domain,
@@ -52,7 +51,6 @@ def __init__(
         Initialization
 
         Args:
-            equation (:class:`PrognosticEquation`): the prognostic equation.
             description (dict): pySDC description
             controller_params (dict): pySDC controller params
             domain (:class:`Domain`): the model's domain object, containing the
@@ -96,6 +94,7 @@ def setup(self, equation, apply_bcs=True, *active_labels):
             'equation': equation,
             'solver_parameters': self.solver_parameters,
             'residual': self._residual,
+            **self.description['problem_params'],
         }
         self.description['level_params']['dt'] = float(self.domain.dt)
 

pySDC/implementations/problem_classes/HeatFiredrake.py

@@ -53,26 +53,30 @@ class Heat1DForcedFiredrake(Problem):
     dtype_u = firedrake_mesh
     dtype_f = IMEX_firedrake_mesh
 
-    def __init__(self, n=30, nu=0.1, c=0.0, LHS_cache_size=12, comm=None):
+    def __init__(self, n=30, nu=0.1, c=0.0, order=4, LHS_cache_size=12, mesh=None, comm=None):
         """
         Initialization
 
         Args:
             n (int): Number of degrees of freedom
             nu (float): Diffusion parameter
             c (float): Boundary condition constant
+            order (int): Order of finite elements
             LHS_cache_size (int): Size of the cache for solvers
-            comm (mpi4pi.Intracomm): MPI communicator for spatial parallelism
+            mesh (Firedrake mesh, optional): Give a custom mesh, for instance from a hierarchy
+            comm (mpi4pi.Intracomm, optional): MPI communicator for spatial parallelism
         """
         comm = MPI.COMM_WORLD if comm is None else comm
 
         # prepare Firedrake mesh and function space
-        self.mesh = fd.UnitIntervalMesh(n, comm=comm)
-        self.V = fd.FunctionSpace(self.mesh, "CG", 4)
+        self.mesh = fd.UnitIntervalMesh(n, comm=comm) if mesh is None else mesh
+        self.V = fd.FunctionSpace(self.mesh, "CG", order)
 
         # prepare pySDC problem class infrastructure by passing the function space to super init
         super().__init__(self.V)
-        self._makeAttributeAndRegister('n', 'nu', 'c', 'LHS_cache_size', 'comm', localVars=locals(), readOnly=True)
+        self._makeAttributeAndRegister(
+            'n', 'nu', 'c', 'order', 'LHS_cache_size', 'comm', localVars=locals(), readOnly=True
+        )
 
         # prepare caches and IO variables for solvers
         self.solvers = {}
@@ -191,6 +195,10 @@ def solve_system(self, rhs, factor, *args, **kwargs):
         self.work_counters['solves']()
         return me
 
+    @fd.utils.cached_property
+    def x(self):
+        return fd.SpatialCoordinate(self.mesh)
+
     def u_exact(self, t):
         r"""
         Routine to compute the exact solution at time :math:`t`.
@@ -206,6 +214,5 @@ def u_exact(self, t):
             Exact solution.
         """
         me = self.u_init
-        x = fd.SpatialCoordinate(self.mesh)
-        me.interpolate(np.cos(t) * fd.sin(np.pi * x[0]) + self.c)
+        me.interpolate(np.cos(t) * fd.sin(np.pi * self.x[0]) + self.c)
         return me

pySDC/implementations/transfer_classes/TransferFiredrakeMesh.py

@@ -0,0 +1,124 @@
+from firedrake import assemble, prolong, inject
+from firedrake.__future__ import interpolate
+
+from pySDC.core.errors import TransferError
+from pySDC.core.space_transfer import SpaceTransfer
+from pySDC.implementations.datatype_classes.firedrake_mesh import firedrake_mesh, IMEX_firedrake_mesh
+
+
+class MeshToMeshFiredrake(SpaceTransfer):
+    """
+    This implementation can restrict and prolong between Firedrake meshes
+    """
+
+    def restrict(self, F):
+        """
+        Restrict from fine to coarse grid
+
+        Args:
+            F: the fine level data
+
+        Returns:
+            Coarse level data
+        """
+        if isinstance(F, firedrake_mesh):
+            u_coarse = self.coarse_prob.dtype_u(assemble(interpolate(F.functionspace, self.coarse_prob.init)))
+        elif isinstance(F, IMEX_firedrake_mesh):
+            u_coarse = IMEX_firedrake_mesh(self.coarse_prob.init)
+            u_coarse.impl.functionspace.assign(assemble(interpolate(F.impl.functionspace, self.coarse_prob.init)))
+            u_coarse.expl.functionspace.assign(assemble(interpolate(F.expl.functionspace, self.coarse_prob.init)))
+        else:
+            raise TransferError('Unknown type of fine data, got %s' % type(F))
+
+        return u_coarse
+
+    def prolong(self, G):
+        """
+        Prolongate from coarse to fine grid
+
+        Args:
+            G: the coarse level data
+
+        Returns:
+            fine level data
+        """
+        if isinstance(G, firedrake_mesh):
+            u_fine = self.fine_prob.dtype_u(assemble(interpolate(G.functionspace, self.fine_prob.init)))
+        elif isinstance(G, IMEX_firedrake_mesh):
+            u_fine = IMEX_firedrake_mesh(self.fine_prob.init)
+            u_fine.impl.functionspace.assign(assemble(interpolate(G.impl.functionspace, self.fine_prob.init)))
+            u_fine.expl.functionspace.assign(assemble(interpolate(G.expl.functionspace, self.fine_prob.init)))
+        else:
+            raise TransferError('Unknown type of coarse data, got %s' % type(G))
+
+        return u_fine
+
+
+class MeshToMeshFiredrakeHierarchy(SpaceTransfer):
+    """
+    This implementation can restrict and prolong between Firedrake meshes that are generated from a hierarchy.
+    Example:
+
+    ```
+    from firedrake import *
+
+    mesh = UnitSquareMesh(8, 8)
+    hierarchy = MeshHierarchy(mesh, 4)
+
+    mesh = hierarchy[-1]
+    ```
+    """
+
+    @staticmethod
+    def _restrict(u_fine, u_coarse):
+        """Perform restriction in Firedrake"""
+        inject(u_fine.functionspace, u_coarse.functionspace)
+
+    @staticmethod
+    def _prolong(u_coarse, u_fine):
+        """Perform prolongation in Firedrake"""
+        prolong(u_coarse.functionspace, u_fine.functionspace)
+
+    def restrict(self, F):
+        """
+        Restrict from fine to coarse grid
+
+        Args:
+            F: the fine level data
+
+        Returns:
+            Coarse level data
+        """
+        if isinstance(F, firedrake_mesh):
+            G = self.coarse_prob.u_init
+            self._restrict(u_fine=F, u_coarse=G)
+        elif isinstance(F, IMEX_firedrake_mesh):
+            G = IMEX_firedrake_mesh(self.coarse_prob.init)
+            self._restrict(u_fine=F.impl, u_coarse=G.impl)
+            self._restrict(u_fine=F.expl, u_coarse=G.expl)
+        else:
+            raise TransferError('Unknown type of fine data, got %s' % type(F))
+
+        return G
+
+    def prolong(self, G):
+        """
+        Prolongate from coarse to fine grid
+
+        Args:
+            G: the coarse level data
+
+        Returns:
+            fine level data
+        """
+        if isinstance(G, firedrake_mesh):
+            F = self.fine_prob.u_init
+            self._prolong(u_coarse=G, u_fine=F)
+        elif isinstance(G, IMEX_firedrake_mesh):
+            F = self.fine_prob.f_init
+            self._prolong(u_coarse=G.impl, u_fine=F.impl)
+            self._prolong(u_coarse=G.expl, u_fine=F.expl)
+        else:
+            raise TransferError('Unknown type of coarse data, got %s' % type(G))
+
+        return F

pySDC/tests/test_helpers/test_gusto_coupling.py

@@ -291,7 +291,6 @@ def test_pySDC_integrator_RK(use_transport_scheme, method, setup):
     stepper_pySDC = get_gusto_stepper(
         eqns,
         pySDC_integrator(
-            eqns,
             description,
             controller_params,
             domain,
@@ -415,7 +414,6 @@ def test_pySDC_integrator(use_transport_scheme, imex, setup):
     stepper_pySDC = get_gusto_stepper(
         eqns,
         pySDC_integrator(
-            eqns,
             description,
             controller_params,
             domain,
@@ -550,7 +548,6 @@ def test_pySDC_integrator_with_adaptivity(dt_initial, setup):
     stepper_pySDC = get_gusto_stepper(
         eqns,
         pySDC_integrator(
-            eqns,
             description,
             controller_params,
             domain,

pySDC/tests/test_transfer_classes/test_firedrake_transfer.py

@@ -0,0 +1,90 @@
+import pytest
+
+
+@pytest.mark.firedrake
+def test_Firedrake_transfer():
+    import firedrake as fd
+    from firedrake.__future__ import interpolate
+    from pySDC.implementations.problem_classes.HeatFiredrake import Heat1DForcedFiredrake
+    from pySDC.implementations.transfer_classes.TransferFiredrakeMesh import MeshToMeshFiredrake
+    import numpy as np
+
+    # prepare fine and coarse problems
+    resolutions = [2, 4]
+    probs = [Heat1DForcedFiredrake(n=n) for n in resolutions]
+
+    # prepare data that we can interpolate exactly in this discretization
+    functions = [fd.Function(me.V).interpolate(me.x**2) for me in probs]
+    data = [me.u_init for me in probs]
+    [data[i].assign(functions[i]) for i in range(len(functions))]
+    rhs = [me.f_init for me in probs]
+    [rhs[i].impl.assign(functions[i]) for i in range(len(functions))]
+    [rhs[i].expl.assign(functions[i]) for i in range(len(functions))]
+
+    # setup transfer class
+    transfer = MeshToMeshFiredrake(*probs[::-1], {})
+
+    # test that restriction and interpolation were indeed exact on the mesh
+    restriction = transfer.restrict(data[1])
+    error = abs(restriction - data[0])
+    assert np.isclose(error, 0), f'Got unexpectedly large {error=} during restriction!'
+
+    interpolation = transfer.prolong(data[0])
+    error = abs(interpolation - data[1])
+    assert np.isclose(error, 0), f'Got unexpectedly large {error=} during interpolation!'
+
+    # test that restriction and interpolation were indeed exact on the IMEX mesh
+    restriction = transfer.restrict(rhs[1])
+    error = max([abs(restriction.impl - rhs[0].impl), abs(restriction.expl - rhs[0].expl)])
+    assert np.isclose(error, 0), f'Got unexpectedly large {error=} during restriction!'
+
+    interpolation = transfer.prolong(rhs[0])
+    error = max([abs(interpolation.impl - rhs[1].impl), abs(interpolation.expl - rhs[1].expl)])
+    assert np.isclose(error, 0), f'Got unexpectedly large {error=} during interpolation!'
+
+
+@pytest.mark.firedrake
+def test_Firedrake_transfer_hierarchy():
+    import firedrake as fd
+    from firedrake.__future__ import interpolate
+    from pySDC.implementations.problem_classes.HeatFiredrake import Heat1DForcedFiredrake
+    from pySDC.implementations.transfer_classes.TransferFiredrakeMesh import MeshToMeshFiredrakeHierarchy
+    import numpy as np
+
+    # prepare fine and coarse problems with the hierarchy
+    _prob = Heat1DForcedFiredrake(n=2)
+    hierarchy = fd.MeshHierarchy(_prob.mesh, 1)
+    probs = [Heat1DForcedFiredrake(mesh=mesh) for mesh in hierarchy]
+
+    # prepare data that we can interpolate exactly in this discretization
+    functions = [fd.Function(me.V).interpolate(me.x**2) for me in probs]
+    data = [me.u_init for me in probs]
+    [data[i].assign(functions[i]) for i in range(len(functions))]
+    rhs = [me.f_init for me in probs]
+    [rhs[i].impl.assign(functions[i]) for i in range(len(functions))]
+    [rhs[i].expl.assign(functions[i]) for i in range(len(functions))]
+
+    # setup transfer class
+    transfer = MeshToMeshFiredrakeHierarchy(*probs[::-1], {})
+
+    # test that restriction and interpolation were indeed exact on the mesh
+    restriction = transfer.restrict(data[1])
+    error = abs(restriction - data[0])
+    assert np.isclose(error, 0), f'Got unexpectedly large {error=} during restriction!'
+
+    interpolation = transfer.prolong(data[0])
+    error = abs(interpolation - data[1])
+    assert np.isclose(error, 0), f'Got unexpectedly large {error=} during interpolation!'
+
+    # test that restriction and interpolation were indeed exact on the IMEX mesh
+    restriction = transfer.restrict(rhs[1])
+    error = max([abs(restriction.impl - rhs[0].impl), abs(restriction.expl - rhs[0].expl)])
+    assert np.isclose(error, 0), f'Got unexpectedly large {error=} during restriction!'
+
+    interpolation = transfer.prolong(rhs[0])
+    error = max([abs(interpolation.impl - rhs[1].impl), abs(interpolation.expl - rhs[1].expl)])
+    assert np.isclose(error, 0), f'Got unexpectedly large {error=} during interpolation!'
+
+
+if __name__ == '__main__':
+    test_Firedrake_transfer_hierarchy()

pySDC/tests/test_tutorials/test_step_7.py

@@ -130,10 +130,11 @@ def test_D():
 
 
 @pytest.mark.firedrake
-def test_E():
+@pytest.mark.parametrize('ML', [True, False])
+def test_E(ML):
     from pySDC.tutorial.step_7.E_pySDC_with_Firedrake import runHeatFiredrake
 
-    runHeatFiredrake(useMPIsweeper=False)
+    runHeatFiredrake(useMPIsweeper=False, ML=ML)
 
 
 @pytest.mark.firedrake
@@ -142,7 +143,7 @@ def test_E_MPI():
     my_env['COVERAGE_PROCESS_START'] = 'pyproject.toml'
     cwd = '.'
     num_procs = 3
-    cmd = f'mpiexec -np {num_procs} python pySDC/tutorial/step_7/E_pySDC_with_Firedrake.py'.split()
+    cmd = f'mpiexec -np {num_procs} python pySDC/tutorial/step_7/E_pySDC_with_Firedrake.py --useMPIsweeper'.split()
 
     p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=my_env, cwd=cwd)
     p.wait()
@@ -179,3 +180,32 @@ def test_F():
     assert (
         error < 1e-8
     ), f'Unexpectedly large difference of {error} between pySDC and Gusto SDC implementations in Williamson 5 test case'
+
+
+@pytest.mark.firedrake
+def test_F_ML():
+    """
+    Test that the Gusto coupling with multiple levels in space converges
+    """
+    from pySDC.tutorial.step_7.F_pySDC_with_Gusto import williamson_5
+    from pySDC.helpers.stats_helper import get_sorted, filter_stats
+    import sys
+
+    if '--running-tests' not in sys.argv:
+        sys.argv += ['--running-tests']
+
+    params = {'use_pySDC': True, 'dt': 1000, 'tmax': 1000, 'use_adaptivity': False, 'M': 2, 'kmax': 4, 'QI': 'LU'}
+    stepper_ML, _ = williamson_5(Nlevels=2, **params)
+    stepper_SL, _ = williamson_5(Nlevels=1, **params)
+
+    stats = stepper_ML.scheme.stats
+    residual_fine = get_sorted(stats, type='residual_post_sweep', level=0, sortby='iter')
+    residual_coarse = get_sorted(stats, type='residual_post_sweep', level=1, sortby='iter')
+    assert residual_fine[0][1] / residual_fine[-1][1] > 3e2, 'Fine residual did not converge as expected'
+    assert residual_coarse[0][1] / residual_coarse[-1][1] > 3e2, 'Coarse residual did not converge as expected'
+
+    stats_SL = stepper_SL.scheme.stats
+    residual_SL = get_sorted(stats_SL, type='residual_post_sweep', sortby='iter')
+    assert all(
+        res_SL > res_ML for res_SL, res_ML in zip(residual_SL, residual_fine)
+    ), 'Single level SDC converged faster than multi-level!'