Part E: pySDC and Firedrake#

Firedrake is a finite element library with similar features as FEniCS. This example runs the same forced heat equation as the FEniCS example, but implemented in Firedrake. The setup proceeds very much as in earlier tutorials; the interesting part is rather the problem class. See pySDC/implementations/problem_classes/HeatFiredrake.py as a blueprint for how to implement problems with Firedrake in a way that pySDC understands.

Single level, serial or parallel across the nodes#

SDC with the diagonal preconditioner MIN-SR-S, whose nodes can be solved in parallel. The space-time parallelism comes from a Firedrake ensemble: with --useMPIsweeper and a multiple of 3 ranks, the ranks are split across the 3 collocation nodes, and the ranks of each node share the spatial problem. See the Firedrake documentation on parallelism.

from pathlib import Path

import numpy as np
from mpi4py import MPI


def setup(useMPIsweeper):
    """
    Helper routine to set up parameters

    Returns:
        description and controller_params parameter dictionaries
    """
    from pySDC.implementations.problem_classes.HeatFiredrake import Heat1DForcedFiredrake
    from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order
    from pySDC.implementations.sweeper_classes.imex_1st_order_MPI import imex_1st_order_MPI
    from pySDC.implementations.hooks.log_errors import LogGlobalErrorPostRun
    from pySDC.implementations.hooks.log_work import LogWork
    from pySDC.helpers.firedrake_ensemble_communicator import FiredrakeEnsembleCommunicator

    # setup space-time parallelism via ensemble for Firedrake, see https://www.firedrakeproject.org/firedrake/parallelism.html
    num_nodes = 3
    ensemble = FiredrakeEnsembleCommunicator(MPI.COMM_WORLD, max([MPI.COMM_WORLD.size // num_nodes, 1]))

    level_params = dict()
    level_params['restol'] = 5e-10
    level_params['dt'] = 0.2

    step_params = dict()
    step_params['maxiter'] = 20

    sweeper_params = dict()
    sweeper_params['quad_type'] = 'RADAU-RIGHT'
    sweeper_params['num_nodes'] = num_nodes
    sweeper_params['QI'] = 'MIN-SR-S'
    sweeper_params['QE'] = 'PIC'
    sweeper_params['comm'] = ensemble

    problem_params = dict()
    problem_params['nu'] = 0.1
    problem_params['n'] = 128
    problem_params['c'] = 1.0
    problem_params['comm'] = ensemble.space_comm

    controller_params = dict()
    controller_params['logger_level'] = 15 if MPI.COMM_WORLD.rank == 0 else 30
    controller_params['hook_class'] = [LogGlobalErrorPostRun, LogWork]

    description = dict()
    description['problem_class'] = Heat1DForcedFiredrake
    description['problem_params'] = problem_params
    description['sweeper_class'] = imex_1st_order_MPI if useMPIsweeper else imex_1st_order
    description['sweeper_params'] = sweeper_params
    description['level_params'] = level_params
    description['step_params'] = step_params

    return description, controller_params

Three levels#

Multilevel SDC, serial, coarsened in space from 128 to 32 to 4 elements, with interpolation of the right-hand side.

def setup_ML():
    """
    Helper routine to set up parameters

    Returns:
        description and controller_params parameter dictionaries
    """
    from pySDC.implementations.problem_classes.HeatFiredrake import Heat1DForcedFiredrake
    from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order
    from pySDC.implementations.sweeper_classes.imex_1st_order_MPI import imex_1st_order_MPI
    from pySDC.implementations.transfer_classes.TransferFiredrakeMesh import MeshToMeshFiredrake
    from pySDC.implementations.hooks.log_errors import LogGlobalErrorPostRun
    from pySDC.implementations.hooks.log_work import LogWork
    from pySDC.helpers.firedrake_ensemble_communicator import FiredrakeEnsembleCommunicator

    level_params = dict()
    level_params['restol'] = 5e-10
    level_params['dt'] = 0.2

    step_params = dict()
    step_params['maxiter'] = 20

    sweeper_params = dict()
    sweeper_params['quad_type'] = 'RADAU-RIGHT'
    sweeper_params['num_nodes'] = 3
    sweeper_params['QI'] = 'MIN-SR-S'
    sweeper_params['QE'] = 'PIC'

    problem_params = dict()
    problem_params['nu'] = 0.1
    problem_params['n'] = [128, 32, 4]
    problem_params['c'] = 1.0

    base_transfer_params = dict()
    base_transfer_params['finter'] = True

    controller_params = dict()
    controller_params['logger_level'] = 15 if MPI.COMM_WORLD.rank == 0 else 30
    controller_params['hook_class'] = [LogGlobalErrorPostRun, LogWork]

    description = dict()
    description['problem_class'] = Heat1DForcedFiredrake
    description['problem_params'] = problem_params
    description['sweeper_class'] = imex_1st_order
    description['sweeper_params'] = sweeper_params
    description['level_params'] = level_params
    description['step_params'] = step_params
    description['space_transfer_class'] = MeshToMeshFiredrake
    description['base_transfer_params'] = base_transfer_params

    return description, controller_params

Running it#

The run prints the error and the work: SDC iterations, solver setups, solves and right-hand side evaluations on the finest level, and checks them against what we got last time.

def runHeatFiredrake(useMPIsweeper=False, ML=False):
    """
    Run the example defined by the above parameters
    """
    from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
    from pySDC.helpers.stats_helper import get_sorted

    Tend = 1.0
    t0 = 0.0

    if ML:
        assert not useMPIsweeper, 'MPI parallel diagonal SDC and ML SDC are not compatible at the moment'
        description, controller_params = setup_ML()
    else:
        description, controller_params = setup(useMPIsweeper)

    controller = controller_nonMPI(num_procs=1, controller_params=controller_params, description=description)

    # get initial values
    P = controller.MS[0].levels[0].prob
    uinit = P.u_exact(0.0)

    # call main function to get things done...
    uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend)

    # see what we get
    error = get_sorted(stats, type='e_global_post_run')
    work_solver_setup = get_sorted(stats, type='work_solver_setup')
    work_solves = get_sorted(stats, type='work_solves')
    work_rhs = get_sorted(stats, type='work_rhs')
    niter = get_sorted(stats, type='niter')

    tot_iter = np.sum([me[1] for me in niter])
    tot_solver_setup = np.sum([me[1] for me in work_solver_setup])
    tot_solves = np.sum([me[1] for me in work_solves])
    tot_rhs = np.sum([me[1] for me in work_rhs])

    time_rank = description["sweeper_params"]["comm"].rank if useMPIsweeper else 0
    print(
        f'Finished with error {error[0][1]:.2e}. Used {tot_iter} SDC iterations, with {tot_solver_setup} solver setups, {tot_solves} solves and {tot_rhs} right hand side evaluations on the finest level of time task {time_rank}.'
    )

    # the results the website shows, from the first rank
    if MPI.COMM_WORLD.rank == 0:
        timing = get_sorted(stats, type='timing_run')[0][1]
        variant = 'three-level SDC' if ML else ('SDC parallel across the nodes' if useMPIsweeper else 'serial SDC')
        Path('data').mkdir(parents=True, exist_ok=True)
        with open('data/step_7_E_out.txt', 'a') as file:
            file.write(
                f'{variant:30s} error {error[0][1]:.2e}, {tot_iter:2d} iterations, {tot_solves:3d} solves on the '
                f'finest level, time to solution {timing:.2f} s\n'
            )

    # do tests that we got the same as last time
    n_nodes = 1 if useMPIsweeper else description['sweeper_params']['num_nodes']
    assert error[0][1] < 2e-7
    assert tot_iter == (10 if ML else 29)
    assert tot_solver_setup == n_nodes
    assert tot_solves == n_nodes * tot_iter
    assert tot_rhs == n_nodes * tot_iter + (n_nodes + 1) * len(niter)

The script runs in three different ways:

  • python E_pySDC_with_Firedrake.py for single-level serial SDC,

  • mpiexec -np 3 python E_pySDC_with_Firedrake.py --useMPIsweeper for single-level SDC, parallel across the nodes,

  • python E_pySDC_with_Firedrake.py --ML for three-level serial SDC.

Multilevel SDC reduces the number of SDC iterations quite a bit, but not necessarily the time to solution, partly because more solvers are constructed on the coarse levels, and the parallel variant does the work of one node per rank. The results below show both; mind that they come from a CI machine, not from a dedicated one. We do not claim to have found the best parameters: this is just an example of how to use it.

if __name__ == "__main__":
    from argparse import ArgumentParser

    parser = ArgumentParser()
    parser.add_argument(
        '--ML',
        help='Whether you want to run multi-level',
        default=False,
        required=False,
        action='store_const',
        const=True,
    )
    parser.add_argument(
        '--useMPIsweeper',
        help='Whether you want to use MPI parallel diagonal SDC',
        default=False,
        required=False,
        action='store_const',
        const=True,
    )

    args = parser.parse_args()

    runHeatFiredrake(**vars(args))

Firedrake does not run in the browser, nor in the environment this website is built in. Our CI runs all three variants in a Firedrake container, and checks their errors and work counts. These are its results, in the run that built this page:

SDC parallel across the nodes  error 1.47e-08, 29 iterations,  29 solves on the finest level, time to solution 6.10 s
three-level SDC                error 1.51e-08, 10 iterations,  30 solves on the finest level, time to solution 10.50 s
serial SDC                     error 1.47e-08, 29 iterations,  87 solves on the finest level, time to solution 2.35 s