diff --git a/.github/workflows/ci_pipeline.yml b/.github/workflows/ci_pipeline.yml
index 3487393df48ca79c26e9513b707f15cf636bfccf..2bfbdfe5214b4cc4a83b2e0e0a377a319555083a 100644
--- a/.github/workflows/ci_pipeline.yml
+++ b/.github/workflows/ci_pipeline.yml
@@ -173,7 +173,7 @@ jobs:
user_firedrake_tests:
runs-on: ubuntu-latest
container:
- image: firedrakeproject/firedrake-vanilla:latest
+ image: firedrakeproject/firedrake-vanilla:2025-01
options: --user root
volumes:
- ${{ github.workspace }}:/repositories
diff --git a/docs/source/tutorial/doc_step_7_E.rst b/docs/source/tutorial/doc_step_7_E.rst
new file mode 100644
index 0000000000000000000000000000000000000000..a0d18755281c10f5f5c29a4261fe1df7530b41d5
--- /dev/null
+++ b/docs/source/tutorial/doc_step_7_E.rst
@@ -0,0 +1,3 @@
+Full code: `pySDC/tutorial/step_7/E_pySDC_with_Firedrake.py <https://github.com/Parallel-in-Time/pySDC/blob/master/pySDC/tutorial/step_7/E_pySDC_with_Firedrake.py>`_
+
+.. literalinclude:: ../../../pySDC/tutorial/step_7/E_pySDC_with_Firedrake.py
diff --git a/docs/source/tutorial/doc_step_7_F.rst b/docs/source/tutorial/doc_step_7_F.rst
new file mode 100644
index 0000000000000000000000000000000000000000..f706dd231b2d918a0fe6eaeb26d9ca84600ddd5d
--- /dev/null
+++ b/docs/source/tutorial/doc_step_7_F.rst
@@ -0,0 +1,4 @@
+Full code: `pySDC/tutorial/step_7/F_pySDC_with_Gusto.py <https://github.com/Parallel-in-Time/pySDC/blob/master/pySDC/tutorial/step_7/F_pySDC_with_Gusto.py>`_
+Find a suitable plotting script here: `pySDC/tutorial/step_7/F_2_plot_pySDC_with_Gusto_result.py <https://github.com/Parallel-in-Time/pySDC/blob/master/pySDC/tutorial/step_7/F_2_plot_pySDC_with_Gusto_result.py>`_
+
+.. literalinclude:: ../../../pySDC/tutorial/step_7/F_pySDC_with_Gusto.py
diff --git a/pySDC/helpers/pySDC_as_gusto_time_discretization.py b/pySDC/helpers/pySDC_as_gusto_time_discretization.py
new file mode 100644
index 0000000000000000000000000000000000000000..33bc7228784e519b84276e36ce1aaa43386c2b89
--- /dev/null
+++ b/pySDC/helpers/pySDC_as_gusto_time_discretization.py
@@ -0,0 +1,174 @@
+import firedrake as fd
+
+from gusto.time_discretisation.time_discretisation import TimeDiscretisation, wrapper_apply
+from gusto.core.labels import explicit
+
+from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
+from pySDC.implementations.problem_classes.GenericGusto import GenericGusto, GenericGustoImex
+from pySDC.core.hooks import Hooks
+from pySDC.helpers.stats_helper import get_sorted
+
+
+class LogTime(Hooks):
+ """
+ Utility hook for knowing how far we got when using adaptive step size selection.
+ """
+
+ def post_step(self, step, level_number):
+ L = step.levels[level_number]
+ self.add_to_stats(
+ process=step.status.slot,
+ process_sweeper=L.sweep.rank,
+ time=L.time,
+ level=-1,
+ iter=-1,
+ sweep=-1,
+ type='_time',
+ value=L.time + L.dt,
+ )
+
+
+class pySDC_integrator(TimeDiscretisation):
+ """
+ This class can be entered into Gusto as a time discretization scheme and will solve steps using pySDC.
+ It will construct a pySDC controller which can be used by itself and will be used within the time step when called
+ from Gusto. Access the controller via `pySDC_integrator.controller`. This class also has `pySDC_integrator.stats`,
+ which gathers all of the pySDC stats recorded in the hooks during every time step when used within Gusto.
+ """
+
+ def __init__(
+ self,
+ equation,
+ description,
+ controller_params,
+ domain,
+ field_name=None,
+ solver_parameters=None,
+ options=None,
+ t0=0,
+ imex=False,
+ ):
+ """
+ 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
+ mesh and the compatible function spaces.
+ field_name (str, optional): name of the field to be evolved.
+ Defaults to None.
+ solver_parameters (dict, optional): dictionary of parameters to
+ pass to the underlying solver. Defaults to None.
+ options (:class:`AdvectionOptions`, optional): an object containing
+ options to either be passed to the spatial discretisation, or
+ to control the "wrapper" methods, such as Embedded DG or a
+ recovery method. Defaults to None.
+ """
+
+ self._residual = None
+
+ super().__init__(
+ domain=domain,
+ field_name=field_name,
+ solver_parameters=solver_parameters,
+ options=options,
+ )
+
+ self.description = description
+ self.controller_params = controller_params
+ self.timestepper = None
+ self.dt_next = None
+ self.imex = imex
+
+ def setup(self, equation, apply_bcs=True, *active_labels):
+ super().setup(equation, apply_bcs, *active_labels)
+
+ # Check if any terms are explicit
+ imex = any(t.has_label(explicit) for t in equation.residual) or self.imex
+ if imex:
+ self.description['problem_class'] = GenericGustoImex
+ else:
+ self.description['problem_class'] = GenericGusto
+
+ self.description['problem_params'] = {
+ 'equation': equation,
+ 'solver_parameters': self.solver_parameters,
+ 'residual': self._residual,
+ }
+ self.description['level_params']['dt'] = float(self.domain.dt)
+
+ # add utility hook required for step size adaptivity
+ hook_class = self.controller_params.get('hook_class', [])
+ if not type(hook_class) == list:
+ hook_class = [hook_class]
+ hook_class.append(LogTime)
+ self.controller_params['hook_class'] = hook_class
+
+ # prepare controller and variables
+ self.controller = controller_nonMPI(1, description=self.description, controller_params=self.controller_params)
+ self.prob = self.level.prob
+ self.sweeper = self.level.sweep
+ self.x0_pySDC = self.prob.dtype_u(self.prob.init)
+ self.t = 0
+ self.stats = {}
+
+ @property
+ def residual(self):
+ """Make sure the pySDC problem residual and this residual are the same"""
+ if hasattr(self, 'prob'):
+ return self.prob.residual
+ else:
+ return self._residual
+
+ @residual.setter
+ def residual(self, value):
+ """Make sure the pySDC problem residual and this residual are the same"""
+ if hasattr(self, 'prob'):
+ self.prob.residual = value
+ else:
+ self._residual = value
+
+ @property
+ def level(self):
+ """Get the finest pySDC level"""
+ return self.controller.MS[0].levels[0]
+
+ @wrapper_apply
+ def apply(self, x_out, x_in):
+ """
+ Apply the time discretization to advance one whole time step.
+
+ Args:
+ x_out (:class:`Function`): the output field to be computed.
+ x_in (:class:`Function`): the input field.
+ """
+ self.x0_pySDC.functionspace.assign(x_in)
+ assert self.level.params.dt == float(self.dt), 'Step sizes have diverged between pySDC and Gusto'
+
+ if self.dt_next is not None:
+ assert (
+ self.timestepper is not None
+ ), 'You need to set self.timestepper to the timestepper in order to facilitate adaptive step size selection here!'
+ self.timestepper.dt = fd.Constant(self.dt_next)
+ self.t = self.timestepper.t
+
+ uend, _stats = self.controller.run(u0=self.x0_pySDC, t0=float(self.t), Tend=float(self.t + self.dt))
+
+ # update time variables
+ if self.level.params.dt != float(self.dt):
+ self.dt_next = self.level.params.dt
+
+ self.t = get_sorted(_stats, type='_time', recomputed=False)[-1][1]
+
+ # update time of the Gusto stepper.
+ # After this step, the Gusto stepper updates its time again to arrive at the correct time
+ if self.timestepper is not None:
+ self.timestepper.t = fd.Constant(self.t - self.dt)
+
+ self.dt = self.level.params.dt
+
+ # update stats and output
+ self.stats = {**self.stats, **_stats}
+ x_out.assign(uend.functionspace)
diff --git a/pySDC/implementations/problem_classes/GenericGusto.py b/pySDC/implementations/problem_classes/GenericGusto.py
new file mode 100644
index 0000000000000000000000000000000000000000..3d8ef50f779288df84bbdab5c33b434473c07256
--- /dev/null
+++ b/pySDC/implementations/problem_classes/GenericGusto.py
@@ -0,0 +1,264 @@
+from pySDC.core.problem import Problem, WorkCounter
+from pySDC.implementations.datatype_classes.firedrake_mesh import firedrake_mesh, IMEX_firedrake_mesh
+from gusto.core.labels import (
+ time_derivative,
+ implicit,
+ explicit,
+)
+from firedrake.fml import replace_subject, all_terms, drop
+import firedrake as fd
+import numpy as np
+
+
+class GenericGusto(Problem):
+ """
+ Set up solvers based on the equation. Keep in mind that you probably want to use the pySDC-Gusto coupling via
+ the `pySDC_integrator` class in the helpers in order to get spatial methods rather than interfacing with this
+ class directly.
+
+ Gusto equations work by a residual, which is minimized in nonlinear solvers to obtain the right hand side
+ evaluation or the solution to (IMEX) Euler steps. You control what you solve for by manipulating labeled parts
+ of the residual.
+ """
+
+ dtype_u = firedrake_mesh
+ dtype_f = firedrake_mesh
+ rhs_n_labels = 1
+
+ def __init__(
+ self,
+ equation,
+ apply_bcs=True,
+ solver_parameters=None,
+ stop_at_divergence=False,
+ LHS_cache_size=12,
+ residual=None,
+ *active_labels,
+ ):
+ """
+ Initialisation
+
+ Args:
+ equation (:class:`PrognosticEquation`): the model's equation.
+ apply_bcs (bool, optional): whether to apply the equation's boundary
+ conditions. Defaults to True.
+ solver_params (dict, optional): Solver parameters for the nonlinear variational problems
+ stop_at_divergence (bool, optional): Whether to raise an error when the variational problems do not converge. Defaults to False
+ LHS_cache_size (int, optional): Size of the cache for solvers. Defaults to 12.
+ residual (Firedrake.form, optional): Overwrite the residual of the equation, e.g. after adding spatial methods. Defaults to None.
+ *active_labels (:class:`Label`): labels indicating which terms of
+ the equation to include.
+ """
+
+ self.equation = equation
+ self.residual = equation.residual if residual is None else residual
+ self.field_name = equation.field_name
+ self.fs = equation.function_space
+ self.idx = None
+ if solver_parameters is None:
+ # default solver parameters
+ solver_parameters = {'ksp_type': 'gmres', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}
+ self.solver_parameters = solver_parameters
+ self.stop_at_divergence = stop_at_divergence
+
+ # -------------------------------------------------------------------- #
+ # Setup caches
+ # -------------------------------------------------------------------- #
+
+ self.x_out = fd.Function(self.fs)
+ self.solvers = {}
+ self._u = fd.Function(self.fs)
+
+ super().__init__(self.fs)
+ self._makeAttributeAndRegister('LHS_cache_size', 'apply_bcs', localVars=locals(), readOnly=True)
+ self.work_counters['rhs'] = WorkCounter()
+ self.work_counters['ksp'] = WorkCounter()
+ self.work_counters['solver_setup'] = WorkCounter()
+ self.work_counters['solver'] = WorkCounter()
+
+ @property
+ def bcs(self):
+ if not self.apply_bcs:
+ return None
+ else:
+ return self.equation.bcs[self.equation.field_name]
+
+ def invert_mass_matrix(self, rhs):
+ self._u.assign(rhs.functionspace)
+
+ if 'mass_matrix' not in self.solvers.keys():
+ mass_form = self.residual.label_map(
+ lambda t: t.has_label(time_derivative),
+ map_if_true=replace_subject(self.x_out, old_idx=self.idx),
+ map_if_false=drop,
+ )
+ rhs_form = self.residual.label_map(
+ lambda t: t.has_label(time_derivative),
+ map_if_true=replace_subject(self._u, old_idx=self.idx),
+ map_if_false=drop,
+ )
+
+ problem = fd.NonlinearVariationalProblem((mass_form - rhs_form).form, self.x_out, bcs=self.bcs)
+ solver_name = self.field_name + self.__class__.__name__
+ self.solvers['mass_matrix'] = fd.NonlinearVariationalSolver(
+ problem, solver_parameters=self.solver_parameters, options_prefix=solver_name
+ )
+ self.work_counters['solver_setup']()
+
+ self.solvers['mass_matrix'].solve()
+
+ return self.dtype_u(self.x_out)
+
+ def eval_f(self, u, *args):
+ self._u.assign(u.functionspace)
+
+ if 'eval_rhs' not in self.solvers.keys():
+ residual = self.residual.label_map(
+ lambda t: t.has_label(time_derivative),
+ map_if_false=replace_subject(self._u, old_idx=self.idx),
+ map_if_true=drop,
+ )
+ mass_form = self.residual.label_map(
+ lambda t: t.has_label(time_derivative),
+ map_if_true=replace_subject(self.x_out, old_idx=self.idx),
+ map_if_false=drop,
+ )
+
+ problem = fd.NonlinearVariationalProblem((mass_form + residual).form, self.x_out, bcs=self.bcs)
+ solver_name = self.field_name + self.__class__.__name__
+ self.solvers['eval_rhs'] = fd.NonlinearVariationalSolver(
+ problem, solver_parameters=self.solver_parameters, options_prefix=solver_name
+ )
+ self.work_counters['solver_setup']()
+
+ self.solvers['eval_rhs'].solve()
+ self.work_counters['rhs']()
+
+ return self.dtype_f(self.x_out)
+
+ def solve_system(self, rhs, factor, u0, *args):
+ self.x_out.assign(u0.functionspace) # set initial guess
+ self._u.assign(rhs.functionspace)
+
+ if factor not in self.solvers.keys():
+ if len(self.solvers) >= self.LHS_cache_size + self.rhs_n_labels:
+ self.solvers.pop(
+ [me for me in self.solvers.keys() if type(me) in [float, int, np.float64, np.float32]][0]
+ )
+
+ # setup left hand side (M - factor*f)(u)
+ # put in output variable
+ residual = self.residual.label_map(all_terms, map_if_true=replace_subject(self.x_out, old_idx=self.idx))
+ # multiply f by factor
+ residual = residual.label_map(
+ lambda t: t.has_label(time_derivative), map_if_false=lambda t: fd.Constant(factor) * t
+ )
+
+ # subtract right hand side
+ mass_form = self.residual.label_map(lambda t: t.has_label(time_derivative), map_if_false=drop)
+ residual -= mass_form.label_map(all_terms, map_if_true=replace_subject(self._u, old_idx=self.idx))
+
+ # construct solver
+ problem = fd.NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs)
+ solver_name = f'{self.field_name}-{self.__class__.__name__}-{factor}'
+ self.solvers[factor] = fd.NonlinearVariationalSolver(
+ problem, solver_parameters=self.solver_parameters, options_prefix=solver_name
+ )
+ self.work_counters['solver_setup']()
+
+ try:
+ self.solvers[factor].solve()
+ except fd.exceptions.ConvergenceError as error:
+ if self.stop_at_divergence:
+ raise error
+ else:
+ self.logger.debug(error)
+
+ self.work_counters['ksp'].niter += self.solvers[factor].snes.getLinearSolveIterations()
+ self.work_counters['solver']()
+ return self.dtype_u(self.x_out)
+
+
+class GenericGustoImex(GenericGusto):
+ dtype_f = IMEX_firedrake_mesh
+ rhs_n_labels = 2
+
+ def evaluate_labeled_term(self, u, label):
+ self._u.assign(u.functionspace)
+
+ if label not in self.solvers.keys():
+ residual = self.residual.label_map(
+ lambda t: t.has_label(label) and not t.has_label(time_derivative),
+ map_if_true=replace_subject(self._u, old_idx=self.idx),
+ map_if_false=drop,
+ )
+ mass_form = self.residual.label_map(
+ lambda t: t.has_label(time_derivative),
+ map_if_true=replace_subject(self.x_out, old_idx=self.idx),
+ map_if_false=drop,
+ )
+
+ problem = fd.NonlinearVariationalProblem((mass_form + residual).form, self.x_out, bcs=self.bcs)
+ solver_name = self.field_name + self.__class__.__name__
+ self.solvers[label] = fd.NonlinearVariationalSolver(
+ problem, solver_parameters=self.solver_parameters, options_prefix=solver_name
+ )
+ self.work_counters['solver_setup'] = WorkCounter()
+
+ self.solvers[label].solve()
+ return self.x_out
+
+ def eval_f(self, u, *args):
+ me = self.dtype_f(self.init)
+ me.impl.assign(self.evaluate_labeled_term(u, implicit))
+ me.expl.assign(self.evaluate_labeled_term(u, explicit))
+ self.work_counters['rhs']()
+ return me
+
+ def solve_system(self, rhs, factor, u0, *args):
+ self.x_out.assign(u0.functionspace) # set initial guess
+ self._u.assign(rhs.functionspace)
+
+ if factor not in self.solvers.keys():
+ if len(self.solvers) >= self.LHS_cache_size + self.rhs_n_labels:
+ self.solvers.pop(
+ [me for me in self.solvers.keys() if type(me) in [float, int, np.float64, np.float32]][0]
+ )
+
+ # setup left hand side (M - factor*f_I)(u)
+ # put in output variable
+ residual = self.residual.label_map(
+ lambda t: t.has_label(time_derivative) or t.has_label(implicit),
+ map_if_true=replace_subject(self.x_out, old_idx=self.idx),
+ map_if_false=drop,
+ )
+ # multiply f_I by factor
+ residual = residual.label_map(
+ lambda t: t.has_label(implicit) and not t.has_label(time_derivative),
+ map_if_true=lambda t: fd.Constant(factor) * t,
+ )
+
+ # subtract right hand side
+ mass_form = self.residual.label_map(lambda t: t.has_label(time_derivative), map_if_false=drop)
+ residual -= mass_form.label_map(all_terms, map_if_true=replace_subject(self._u, old_idx=self.idx))
+
+ # construct solver
+ problem = fd.NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs)
+ solver_name = f'{self.field_name}-{self.__class__.__name__}-{factor}'
+ self.solvers[factor] = fd.NonlinearVariationalSolver(
+ problem, solver_parameters=self.solver_parameters, options_prefix=solver_name
+ )
+ self.work_counters['solver_setup'] = WorkCounter()
+
+ self.solvers[factor].solve()
+ try:
+ self.solvers[factor].solve()
+ except fd.exceptions.ConvergenceError as error:
+ if self.stop_at_divergence:
+ raise error
+ else:
+ self.logger.debug(error)
+
+ self.work_counters['ksp'].niter += self.solvers[factor].snes.getLinearSolveIterations()
+ self.work_counters['solver']()
+ return self.dtype_u(self.x_out)
diff --git a/pySDC/tests/test_helpers/test_gusto_coupling.py b/pySDC/tests/test_helpers/test_gusto_coupling.py
new file mode 100644
index 0000000000000000000000000000000000000000..69f8611819189b94e1e223b65612ff7a8f5f69ef
--- /dev/null
+++ b/pySDC/tests/test_helpers/test_gusto_coupling.py
@@ -0,0 +1,626 @@
+import pytest
+
+
+def get_gusto_stepper(eqns, method, spatial_methods):
+ from gusto import IO, OutputParameters, PrescribedTransport
+ import sys
+
+ if '--running-tests' not in sys.argv:
+ sys.argv.append('--running-tests')
+
+ output = OutputParameters(dirname='./tmp', dumpfreq=15)
+ io = IO(method.domain, output)
+ return PrescribedTransport(eqns, method, io, False, transport_method=spatial_methods)
+
+
+def tracer_setup(tmpdir='./tmp', degree=1, small_dt=False):
+ from firedrake import (
+ IcosahedralSphereMesh,
+ PeriodicIntervalMesh,
+ ExtrudedMesh,
+ SpatialCoordinate,
+ as_vector,
+ sqrt,
+ exp,
+ pi,
+ )
+ from gusto import OutputParameters, Domain, IO
+ from collections import namedtuple
+
+ opts = ('domain', 'tmax', 'io', 'f_init', 'f_end', 'degree', 'uexpr', 'umax', 'radius', 'tol')
+ TracerSetup = namedtuple('TracerSetup', opts)
+ TracerSetup.__new__.__defaults__ = (None,) * len(opts)
+
+ radius = 1
+ mesh = IcosahedralSphereMesh(radius=radius, refinement_level=3, degree=1)
+ x = SpatialCoordinate(mesh)
+
+ # Parameters chosen so that dt != 1
+ # Gaussian is translated from (lon=pi/2, lat=0) to (lon=0, lat=0)
+ # to demonstrate that transport is working correctly
+ if small_dt:
+ dt = pi / 3.0 * 0.005
+ else:
+ dt = pi / 3.0 * 0.02
+
+ output = OutputParameters(dirname=str(tmpdir), dumpfreq=15)
+ domain = Domain(mesh, dt, family="BDM", degree=degree)
+ io = IO(domain, output)
+
+ umax = 1.0
+ uexpr = as_vector([-umax * x[1] / radius, umax * x[0] / radius, 0.0])
+
+ tmax = pi / 2
+ f_init = exp(-x[2] ** 2 - x[0] ** 2)
+ f_end = exp(-x[2] ** 2 - x[1] ** 2)
+
+ tol = 0.05
+
+ return TracerSetup(domain, tmax, io, f_init, f_end, degree, uexpr, umax, radius, tol)
+
+
+@pytest.fixture
+def setup():
+ return tracer_setup()
+
+
+def get_gusto_advection_setup(use_transport_scheme, imex, setup):
+ from gusto import ContinuityEquation, AdvectionEquation, split_continuity_form, DGUpwind
+ from gusto.core.labels import time_derivative, transport, implicit, explicit
+
+ domain = setup.domain
+ V = domain.spaces("DG")
+
+ eqn = ContinuityEquation(domain, V, "f")
+ eqn = split_continuity_form(eqn)
+
+ transport_methods = [DGUpwind(eqn, 'f')]
+ spatial_methods = None
+
+ if use_transport_scheme:
+ spatial_methods = transport_methods
+
+ if imex:
+ eqn.label_terms(lambda t: not any(t.has_label(time_derivative, transport)), implicit)
+ eqn.label_terms(lambda t: t.has_label(transport), explicit)
+ else:
+ eqn.label_terms(lambda t: not t.has_label(time_derivative), implicit)
+
+ return eqn, domain, spatial_methods, setup
+
+
+def get_initial_conditions(timestepper, setup):
+ timestepper.fields("f").interpolate(setup.f_init)
+ timestepper.fields("u").project(setup.uexpr)
+ return timestepper
+
+
+@pytest.mark.firedrake
+def test_generic_gusto_problem(setup):
+ from pySDC.implementations.problem_classes.GenericGusto import GenericGusto
+ from firedrake import norm, Constant
+ import numpy as np
+ from gusto import ThetaMethod
+
+ eqns, domain, spatial_methods, setup = get_gusto_advection_setup(False, False, setup)
+
+ dt = 1e-1
+ domain.dt = Constant(dt)
+
+ solver_parameters = {
+ 'snes_type': 'newtonls',
+ 'ksp_type': 'gmres',
+ 'pc_type': 'bjacobi',
+ 'sub_pc_type': 'ilu',
+ 'ksp_rtol': 1e-12,
+ 'snes_rtol': 1e-12,
+ 'ksp_atol': 1e-30,
+ 'snes_atol': 1e-30,
+ 'ksp_divtol': 1e30,
+ 'snes_divtol': 1e30,
+ 'snes_max_it': 99,
+ }
+
+ # ------------------------------------------------------------------------ #
+ # Prepare different methods
+ # ------------------------------------------------------------------------ #
+
+ problem = GenericGusto(eqns, solver_parameters=solver_parameters)
+ stepper_backward = get_gusto_stepper(
+ eqns, ThetaMethod(domain, theta=1.0, solver_parameters=solver_parameters), spatial_methods
+ )
+ stepper_forward = get_gusto_stepper(
+ eqns, ThetaMethod(domain, theta=0.0, solver_parameters=solver_parameters), spatial_methods
+ )
+
+ # ------------------------------------------------------------------------ #
+ # Run tests
+ # ------------------------------------------------------------------------ #
+
+ for stepper in [stepper_backward, stepper_forward]:
+ get_initial_conditions(stepper, setup)
+
+ u_start = problem.u_init
+ u_start.assign(stepper_backward.fields('f'))
+
+ un = problem.solve_system(u_start, dt, u_start)
+ fn = problem.eval_f(un)
+
+ u02 = un - dt * fn
+
+ error = abs(u_start - u02) / abs(u_start)
+
+ assert error < np.finfo(float).eps * 1e2
+
+ # test forward Euler step
+ stepper_forward.run(t=0, tmax=dt)
+ un_ref = problem.dtype_u(problem.init)
+ un_ref.assign(stepper_forward.fields('f'))
+ un_forward = u_start + dt * problem.eval_f(u_start)
+ error = abs(un_forward - un_ref) / abs(un_ref)
+
+ assert (
+ error < np.finfo(float).eps * 1e2
+ ), f'Forward Euler does not match reference implementation! Got relative difference of {error}'
+
+ # test backward Euler step
+ stepper_backward.run(t=0, tmax=dt)
+ un_ref = problem.dtype_u(problem.init)
+ un_ref.assign(stepper_backward.fields('f'))
+ error = abs(un - un_ref) / abs(un_ref)
+
+ assert (
+ error < np.finfo(float).eps * 1e2
+ ), f'Backward Euler does not match reference implementation! Got relative difference of {error}'
+
+
+class Method(object):
+ imex = False
+
+ @staticmethod
+ def get_pySDC_method():
+ raise NotImplementedError
+
+ @staticmethod
+ def get_Gusto_method():
+ raise NotImplementedError
+
+
+class RK4(Method):
+
+ @staticmethod
+ def get_pySDC_method():
+ from pySDC.implementations.sweeper_classes.Runge_Kutta import RK4
+
+ return RK4
+
+ @staticmethod
+ def get_Gusto_method():
+ from gusto import RK4
+
+ return RK4
+
+
+class ImplicitMidpoint(Method):
+
+ @staticmethod
+ def get_pySDC_method():
+ from pySDC.implementations.sweeper_classes.Runge_Kutta import ImplicitMidpointMethod
+
+ return ImplicitMidpointMethod
+
+ @staticmethod
+ def get_Gusto_method():
+ from gusto import ImplicitMidpoint
+
+ return ImplicitMidpoint
+
+
+class BackwardEuler(Method):
+
+ @staticmethod
+ def get_pySDC_method():
+ from pySDC.implementations.sweeper_classes.Runge_Kutta import BackwardEuler
+
+ return BackwardEuler
+
+ @staticmethod
+ def get_Gusto_method():
+ from gusto import BackwardEuler
+
+ return BackwardEuler
+
+
+@pytest.mark.firedrake
+@pytest.mark.parametrize('use_transport_scheme', [True, False])
+@pytest.mark.parametrize('method', [RK4, ImplicitMidpoint, BackwardEuler])
+def test_pySDC_integrator_RK(use_transport_scheme, method, setup):
+ from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
+ from pySDC.helpers.pySDC_as_gusto_time_discretization import pySDC_integrator
+ from firedrake import norm
+ import numpy as np
+
+ eqns, domain, spatial_methods, setup = get_gusto_advection_setup(use_transport_scheme, method.imex, setup)
+
+ solver_parameters = {
+ 'snes_type': 'newtonls',
+ 'ksp_type': 'gmres',
+ 'pc_type': 'bjacobi',
+ 'sub_pc_type': 'ilu',
+ 'ksp_rtol': 1e-12,
+ 'snes_rtol': 1e-12,
+ 'ksp_atol': 1e-30,
+ 'snes_atol': 1e-30,
+ 'ksp_divtol': 1e30,
+ 'snes_divtol': 1e30,
+ 'snes_max_it': 99,
+ }
+
+ # ------------------------------------------------------------------------ #
+ # Setup pySDC
+ # ------------------------------------------------------------------------ #
+
+ level_params = dict()
+ level_params['restol'] = -1
+
+ step_params = dict()
+ step_params['maxiter'] = 1
+
+ sweeper_params = dict()
+
+ problem_params = dict()
+
+ controller_params = dict()
+ controller_params['logger_level'] = 20
+ controller_params['mssdc_jac'] = False
+
+ description = dict()
+ description['problem_params'] = problem_params
+ description['sweeper_class'] = method.get_pySDC_method()
+ description['sweeper_params'] = sweeper_params
+ description['level_params'] = level_params
+ description['step_params'] = step_params
+
+ # ------------------------------------------------------------------------ #
+ # Setup time steppers
+ # ------------------------------------------------------------------------ #
+
+ stepper_gusto = get_gusto_stepper(
+ eqns, method.get_Gusto_method()(domain, solver_parameters=solver_parameters), spatial_methods
+ )
+ stepper_pySDC = get_gusto_stepper(
+ eqns,
+ pySDC_integrator(
+ eqns,
+ description,
+ controller_params,
+ domain,
+ solver_parameters=solver_parameters,
+ imex=method.imex,
+ ),
+ spatial_methods,
+ )
+
+ # ------------------------------------------------------------------------ #
+ # Initial conditions
+ # ------------------------------------------------------------------------ #
+
+ for stepper in [stepper_gusto, stepper_pySDC]:
+ get_initial_conditions(stepper, setup)
+
+ # ------------------------------------------------------------------------ #
+ # Run tests
+ # ------------------------------------------------------------------------ #
+
+ def run(stepper, n_steps):
+ stepper.run(t=0, tmax=n_steps * float(domain.dt))
+
+ for stepper in [stepper_gusto, stepper_pySDC]:
+ run(stepper, 5)
+
+ error = norm(stepper_gusto.fields('f') - stepper_pySDC.fields('f')) / norm(stepper_gusto.fields('f'))
+ print(error)
+
+ assert (
+ error < solver_parameters['snes_rtol'] * 1e3
+ ), f'pySDC and Gusto differ in method {method}! Got relative difference of {error}'
+
+
+@pytest.mark.firedrake
+@pytest.mark.parametrize('use_transport_scheme', [True, False])
+@pytest.mark.parametrize('imex', [True, False])
+def test_pySDC_integrator(use_transport_scheme, imex, setup):
+ from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
+ from pySDC.helpers.pySDC_as_gusto_time_discretization import pySDC_integrator
+ from gusto import BackwardEuler, SDC
+ from firedrake import norm
+ import numpy as np
+
+ eqns, domain, spatial_methods, setup = get_gusto_advection_setup(use_transport_scheme, imex, setup)
+
+ solver_parameters = {
+ 'snes_type': 'newtonls',
+ 'ksp_type': 'gmres',
+ 'pc_type': 'bjacobi',
+ 'sub_pc_type': 'ilu',
+ 'ksp_rtol': 1e-12,
+ 'snes_rtol': 1e-12,
+ 'ksp_atol': 1e-30,
+ 'snes_atol': 1e-30,
+ 'ksp_divtol': 1e30,
+ 'snes_divtol': 1e30,
+ 'snes_max_it': 99,
+ }
+
+ # ------------------------------------------------------------------------ #
+ # Setup pySDC
+ # ------------------------------------------------------------------------ #
+ if imex:
+ from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order as sweeper_cls
+ else:
+ from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit as sweeper_cls
+
+ level_params = dict()
+ level_params['restol'] = -1
+
+ step_params = dict()
+ step_params['maxiter'] = 3
+
+ sweeper_params = dict()
+ sweeper_params['quad_type'] = 'RADAU-RIGHT'
+ sweeper_params['node_type'] = 'LEGENDRE'
+ sweeper_params['num_nodes'] = 2
+ sweeper_params['QI'] = 'IE'
+ sweeper_params['QE'] = 'PIC'
+ sweeper_params['initial_guess'] = 'copy'
+
+ problem_params = dict()
+
+ controller_params = dict()
+ controller_params['logger_level'] = 20
+ controller_params['mssdc_jac'] = False
+
+ description = dict()
+ description['problem_params'] = problem_params
+ description['sweeper_class'] = sweeper_cls
+ description['sweeper_params'] = sweeper_params
+ description['level_params'] = level_params
+ description['step_params'] = step_params
+
+ # ------------------------------------------------------------------------ #
+ # Setup SDC in gusto
+ # ------------------------------------------------------------------------ #
+
+ SDC_params = {
+ 'base_scheme': BackwardEuler(domain, solver_parameters=solver_parameters),
+ 'M': sweeper_params['num_nodes'],
+ 'maxk': step_params['maxiter'],
+ 'quad_type': sweeper_params['quad_type'],
+ 'node_type': sweeper_params['node_type'],
+ 'qdelta_imp': sweeper_params['QI'],
+ 'qdelta_exp': sweeper_params['QE'],
+ 'formulation': 'Z2N',
+ 'initial_guess': 'copy',
+ 'nonlinear_solver_parameters': solver_parameters,
+ 'linear_solver_parameters': solver_parameters,
+ 'final_update': False,
+ }
+
+ # ------------------------------------------------------------------------ #
+ # Setup time steppers
+ # ------------------------------------------------------------------------ #
+
+ stepper_gusto = get_gusto_stepper(eqns, SDC(**SDC_params, domain=domain), spatial_methods)
+
+ stepper_pySDC = get_gusto_stepper(
+ eqns,
+ pySDC_integrator(
+ eqns,
+ description,
+ controller_params,
+ domain,
+ solver_parameters=solver_parameters,
+ imex=imex,
+ ),
+ spatial_methods,
+ )
+
+ # ------------------------------------------------------------------------ #
+ # Initial conditions
+ # ------------------------------------------------------------------------ #
+
+ for stepper in [stepper_gusto, stepper_pySDC]:
+ get_initial_conditions(stepper, setup)
+
+ # ------------------------------------------------------------------------ #
+ # Run tests
+ # ------------------------------------------------------------------------ #
+
+ def run(stepper, n_steps):
+ stepper.run(t=0, tmax=n_steps * float(domain.dt))
+
+ for stepper in [stepper_gusto, stepper_pySDC]:
+ run(stepper, 5)
+
+ error = norm(stepper_gusto.fields('f') - stepper_pySDC.fields('f')) / norm(stepper_gusto.fields('f'))
+ print(error)
+
+ assert (
+ error < solver_parameters['snes_rtol'] * 1e3
+ ), f'pySDC and Gusto differ in SDC! Got relative difference of {error}'
+
+
+@pytest.mark.firedrake
+@pytest.mark.parametrize('dt_initial', [1e-5, 1e-1])
+def test_pySDC_integrator_with_adaptivity(dt_initial, setup):
+ from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
+ from pySDC.implementations.convergence_controller_classes.adaptivity import Adaptivity
+ from pySDC.implementations.convergence_controller_classes.spread_step_sizes import SpreadStepSizesBlockwiseNonMPI
+ from pySDC.implementations.convergence_controller_classes.step_size_limiter import StepSizeRounding
+ from pySDC.helpers.pySDC_as_gusto_time_discretization import pySDC_integrator
+ from pySDC.helpers.stats_helper import get_sorted
+ from gusto import BackwardEuler, SDC
+ from firedrake import norm, Constant
+ import numpy as np
+
+ use_transport_scheme = True
+ imex = True
+
+ eqns, domain, spatial_methods, setup = get_gusto_advection_setup(use_transport_scheme, imex, setup)
+ domain.dt = Constant(dt_initial)
+
+ solver_parameters = {
+ 'snes_type': 'newtonls',
+ 'ksp_type': 'gmres',
+ 'pc_type': 'bjacobi',
+ 'sub_pc_type': 'ilu',
+ 'ksp_rtol': 1e-12,
+ 'snes_rtol': 1e-12,
+ 'ksp_atol': 1e-30,
+ 'snes_atol': 1e-30,
+ 'ksp_divtol': 1e30,
+ 'snes_divtol': 1e30,
+ 'snes_max_it': 99,
+ }
+
+ # ------------------------------------------------------------------------ #
+ # Setup pySDC
+ # ------------------------------------------------------------------------ #
+ if imex:
+ from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order as sweeper_cls
+ else:
+ from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit as sweeper_cls
+
+ level_params = dict()
+ level_params['restol'] = -1
+
+ step_params = dict()
+ step_params['maxiter'] = 3
+
+ sweeper_params = dict()
+ sweeper_params['quad_type'] = 'RADAU-RIGHT'
+ sweeper_params['node_type'] = 'LEGENDRE'
+ sweeper_params['num_nodes'] = 2
+ sweeper_params['QI'] = 'IE'
+ sweeper_params['QE'] = 'PIC'
+ sweeper_params['initial_guess'] = 'copy'
+
+ problem_params = dict()
+
+ convergence_controllers = {}
+ convergence_controllers[Adaptivity] = {'e_tol': 1e-5, 'rel_error': True}
+ convergence_controllers[SpreadStepSizesBlockwiseNonMPI] = {'overwrite_to_reach_Tend': False}
+ convergence_controllers[StepSizeRounding] = {}
+
+ controller_params = dict()
+ controller_params['logger_level'] = 15
+ controller_params['mssdc_jac'] = False
+
+ description = dict()
+ description['problem_params'] = problem_params
+ description['sweeper_class'] = sweeper_cls
+ description['sweeper_params'] = sweeper_params
+ description['level_params'] = level_params
+ description['step_params'] = step_params
+ description['convergence_controllers'] = convergence_controllers
+
+ # ------------------------------------------------------------------------ #
+ # Setup SDC in gusto
+ # ------------------------------------------------------------------------ #
+
+ SDC_params = {
+ 'base_scheme': BackwardEuler(domain, solver_parameters=solver_parameters),
+ 'M': sweeper_params['num_nodes'],
+ 'maxk': step_params['maxiter'],
+ 'quad_type': sweeper_params['quad_type'],
+ 'node_type': sweeper_params['node_type'],
+ 'qdelta_imp': sweeper_params['QI'],
+ 'qdelta_exp': sweeper_params['QE'],
+ 'formulation': 'Z2N',
+ 'initial_guess': 'copy',
+ 'nonlinear_solver_parameters': solver_parameters,
+ 'linear_solver_parameters': solver_parameters,
+ 'final_update': False,
+ }
+
+ # ------------------------------------------------------------------------ #
+ # Setup time steppers
+ # ------------------------------------------------------------------------ #
+
+ stepper_pySDC = get_gusto_stepper(
+ eqns,
+ pySDC_integrator(
+ eqns,
+ description,
+ controller_params,
+ domain,
+ solver_parameters=solver_parameters,
+ imex=imex,
+ ),
+ spatial_methods,
+ )
+
+ stepper_gusto = get_gusto_stepper(eqns, SDC(**SDC_params, domain=domain), spatial_methods)
+
+ stepper_pySDC.scheme.timestepper = stepper_pySDC
+
+ # ------------------------------------------------------------------------ #
+ # Initial conditions
+ # ------------------------------------------------------------------------ #
+
+ for stepper in [stepper_gusto, stepper_pySDC]:
+ get_initial_conditions(stepper, setup)
+
+ # ------------------------------------------------------------------------ #
+ # Run tests
+ # ------------------------------------------------------------------------ #
+
+ # run with pySDC first
+ get_initial_conditions(stepper_pySDC, setup)
+ stepper_pySDC.run(t=0, tmax=0.1)
+
+ # retrieve step sizes
+ stats = stepper_pySDC.scheme.stats
+ dts_pySDC = get_sorted(stats, type='dt', recomputed=False)
+
+ assert len(dts_pySDC) > 0, 'No step sizes were recorded in adaptivity test!'
+
+ # run with Gusto using same step sizes
+ get_initial_conditions(stepper_gusto, setup)
+ old_dt = float(stepper_gusto.dt)
+
+ for _dt in dts_pySDC:
+ # update step size
+ stepper_gusto.dt = Constant(_dt[1])
+
+ stepper_gusto.scheme.Q *= _dt[1] / old_dt
+ stepper_gusto.scheme.Qdelta_imp *= _dt[1] / old_dt
+ stepper_gusto.scheme.Qdelta_exp *= _dt[1] / old_dt
+ stepper_gusto.scheme.nodes *= _dt[1] / old_dt
+
+ old_dt = _dt[1] * 1.0
+
+ # run
+ stepper_gusto.run(t=_dt[0], tmax=_dt[0] + _dt[1])
+
+ # clear solver cache with old step size
+ del stepper_gusto.scheme.solvers
+
+ assert np.isclose(float(stepper_pySDC.t), float(stepper_gusto.t))
+
+ print(dts_pySDC)
+
+ error = norm(stepper_gusto.fields('f') - stepper_pySDC.fields('f')) / norm(stepper_gusto.fields('f'))
+ print(error, norm(stepper_gusto.fields('f')))
+
+ assert (
+ error < np.finfo(float).eps * 1e2
+ ), f'SDC does not match reference implementation with adaptive step size selection! Got relative difference of {error}'
+
+
+if __name__ == '__main__':
+ setup = tracer_setup()
+ # test_generic_gusto_problem(setup)
+ # test_pySDC_integrator_RK(False, RK4, setup)
+ # test_pySDC_integrator(False, False, setup)
+ test_pySDC_integrator_with_adaptivity(1e-3, setup)
diff --git a/pySDC/tests/test_tutorials/test_step_7.py b/pySDC/tests/test_tutorials/test_step_7.py
index 56aa6cb83f6c0a7a49b34d59d1f27128fb90b89c..6820a17d15a8b94eec7e42d45d8dfafff72b552c 100644
--- a/pySDC/tests/test_tutorials/test_step_7.py
+++ b/pySDC/tests/test_tutorials/test_step_7.py
@@ -151,3 +151,31 @@ def test_E_MPI():
for line in p.stderr:
print(line)
assert p.returncode == 0, 'ERROR: did not get return code 0, got %s with %2i processes' % (p.returncode, num_procs)
+
+
+@pytest.mark.firedrake
+def test_F():
+ """
+ Test that the same result is obtained using the pySDC and Gusto coupling compared to only using Gusto after a few time steps.
+ The test problem is Williamson 5, which involves huge numbers. Due to roundoff errors, we therefore cannot expect the solutions to match exactly.
+ """
+ from pySDC.tutorial.step_7.F_pySDC_with_Gusto import williamson_5
+ from firedrake import norm
+ import sys
+
+ if '--running-tests' not in sys.argv:
+ sys.argv += ['--running-tests']
+
+ params = {'dt': 900, 'tmax': 2700, 'use_adaptivity': False, 'M': 2, 'kmax': 3, 'QI': 'LU'}
+ stepper_pySDC, mesh = williamson_5(use_pySDC=True, **params)
+ stepper_gusto, mesh = williamson_5(use_pySDC=False, mesh=mesh, **params)
+
+ error = max(
+ [
+ norm(stepper_gusto.fields(comp) - stepper_pySDC.fields(comp)) / norm(stepper_gusto.fields(comp))
+ for comp in ['u', 'D']
+ ]
+ )
+ assert (
+ error < 1e-8
+ ), f'Unexpectedly large difference of {error} between pySDC and Gusto SDC implementations in Williamson 5 test case'
diff --git a/pySDC/tutorial/step_7/F_2_plot_pySDC_with_Gusto_result.py b/pySDC/tutorial/step_7/F_2_plot_pySDC_with_Gusto_result.py
new file mode 100644
index 0000000000000000000000000000000000000000..8d091e33378f1b3c2bb23fdd1f57407cc8775d74
--- /dev/null
+++ b/pySDC/tutorial/step_7/F_2_plot_pySDC_with_Gusto_result.py
@@ -0,0 +1,205 @@
+def plot_williamson_5(
+ results_file_name='./results/williamson_5/field_output.nc', plot_stem='williamson_5'
+): # pragma: no cover
+ """
+ Plot the results of the Williamson 5 test case from tutorial step 7/F.
+ This file is taken from the Gusto example scripts and stores plots in the current directory.
+
+ See https://github.com/tommbendall/tomplot for obtaining tomplot
+
+ The initial conditions are plotted:
+ (a) the velocity field, (b) the depth field.
+
+ And at the last recorded time, this plots:
+ (a) the relative vorticity field, (b) free-surface height.
+
+ Args:
+ results_file_name (str): Path to netCDF results file from Williamson 5 test case
+ plot_stem (str): Name of the resulting plots
+ """
+ import cartopy.crs as ccrs
+ import matplotlib.pyplot as plt
+ import numpy as np
+ from netCDF4 import Dataset
+ from tomplot import (
+ set_tomplot_style,
+ tomplot_cmap,
+ plot_contoured_field,
+ add_colorbar_ax,
+ plot_field_quivers,
+ tomplot_field_title,
+ extract_gusto_coords,
+ extract_gusto_field,
+ regrid_horizontal_slice,
+ )
+
+ # ---------------------------------------------------------------------------- #
+ # Initial plot details
+ # ---------------------------------------------------------------------------- #
+ init_field_names = ['u', 'D']
+ init_colour_schemes = ['Oranges', 'YlGnBu']
+ init_field_labels = [r'$|u|$ (m s$^{-1}$)', r'$D$ (m)']
+ init_contours_to_remove = [None, None, None]
+ init_contours = [np.linspace(0, 20, 9), np.linspace(3800, 6000, 12)]
+
+ # ---------------------------------------------------------------------------- #
+ # Final plot details
+ # ---------------------------------------------------------------------------- #
+ final_field_names = ['RelativeVorticity', 'D_plus_topography']
+ final_colour_schemes = ['RdBu_r', 'PiYG']
+ final_field_labels = [r'$\zeta \ / $ s$^{-1}$', r'$D+B$ (m)']
+ final_contours = [np.linspace(-1e-4, 1e-4, 21), np.linspace(5000, 6000, 11)]
+ final_contours_to_remove = [0.0, None]
+
+ # ---------------------------------------------------------------------------- #
+ # General options
+ # ---------------------------------------------------------------------------- #
+ projection = ccrs.Robinson()
+ contour_method = 'contour'
+ # xlims = [-180, 180]
+ # ylims = [-90, 90]
+
+ cbar_format = {'RelativeVorticity': '1.1e', 'u': '1.0f', 'D': '1.0f', 'D_plus_topography': '1.0f'}
+
+ # We need to regrid onto lon-lat grid -- specify that here
+ lon_1d = np.linspace(-180.0, 180.0, 120)
+ lat_1d = np.linspace(-90, 90, 120)
+ lon_2d, lat_2d = np.meshgrid(lon_1d, lat_1d, indexing='ij')
+
+ # Things that are likely the same for all plots --------------------------------
+ set_tomplot_style()
+ data_file = Dataset(results_file_name, 'r')
+
+ # ---------------------------------------------------------------------------- #
+ # INITIAL PLOTTING
+ # ---------------------------------------------------------------------------- #
+ fig = plt.figure(figsize=(15, 5))
+ time_idx = 0
+
+ for i, (field_name, colour_scheme, field_label, contour_to_remove, contours) in enumerate(
+ zip(init_field_names, init_colour_schemes, init_field_labels, init_contours_to_remove, init_contours)
+ ):
+
+ # Make axes
+ ax = fig.add_subplot(1, 2, 1 + i, projection=projection)
+
+ # Data extraction ----------------------------------------------------------
+ if field_name == 'u':
+ zonal_data = extract_gusto_field(data_file, 'u_zonal', time_idx=time_idx)
+ meridional_data = extract_gusto_field(data_file, 'u_meridional', time_idx=time_idx)
+ field_data = np.sqrt(zonal_data**2 + meridional_data**2)
+ coords_X, coords_Y = extract_gusto_coords(data_file, 'u_zonal')
+
+ else:
+ field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx)
+ coords_X, coords_Y = extract_gusto_coords(data_file, field_name)
+ time = data_file['time'][time_idx] / (24.0 * 60.0 * 60.0)
+
+ # Regrid onto lon-lat grid
+ field_data = regrid_horizontal_slice(lon_2d, lat_2d, coords_X, coords_Y, field_data, periodic_fix='sphere')
+
+ # Plot data ----------------------------------------------------------------
+ cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove)
+ cf, _ = plot_contoured_field(
+ ax,
+ lon_2d,
+ lat_2d,
+ field_data,
+ contour_method,
+ contours,
+ cmap=cmap,
+ line_contours=lines,
+ projection=projection,
+ )
+
+ add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10, cbar_format=cbar_format[field_name])
+ tomplot_field_title(ax, None, minmax=True, field_data=field_data)
+
+ # Add quivers --------------------------------------------------------------
+ if field_name == 'u':
+ # Need to re-grid to lat-lon grid to get sensible looking quivers
+ regrid_zonal_data = regrid_horizontal_slice(
+ lon_2d, lat_2d, coords_X, coords_Y, zonal_data, periodic_fix='sphere'
+ )
+ regrid_meridional_data = regrid_horizontal_slice(
+ lon_2d, lat_2d, coords_X, coords_Y, meridional_data, periodic_fix='sphere'
+ )
+ plot_field_quivers(
+ ax,
+ lon_2d,
+ lat_2d,
+ regrid_zonal_data,
+ regrid_meridional_data,
+ spatial_filter_step=6,
+ magnitude_filter=1.0,
+ projection=ccrs.PlateCarree(),
+ )
+
+ # Save figure ------------------------------------------------------------------
+ fig.subplots_adjust(wspace=0.25)
+ plt.suptitle(f't = {time:.1f} days')
+ plot_name = f'{plot_stem}_initial.png'
+ print(f'Saving figure to {plot_name}')
+ fig.savefig(plot_name, bbox_inches='tight')
+ plt.close()
+
+ # ---------------------------------------------------------------------------- #
+ # FINAL PLOTTING
+ # ---------------------------------------------------------------------------- #
+ fig = plt.figure(figsize=(15, 5))
+ time_idx = -1
+
+ for i, (field_name, colour_scheme, field_label, contours, contour_to_remove) in enumerate(
+ zip(final_field_names, final_colour_schemes, final_field_labels, final_contours, final_contours_to_remove)
+ ):
+
+ # Make axes
+ ax = fig.add_subplot(1, 2, 1 + i, projection=projection)
+
+ # Data extraction ----------------------------------------------------------
+ field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx)
+ coords_X, coords_Y = extract_gusto_coords(data_file, field_name)
+ time = data_file['time'][time_idx] / (24.0 * 60.0 * 60.0)
+
+ # Regrid onto lon-lat grid
+ field_data = regrid_horizontal_slice(lon_2d, lat_2d, coords_X, coords_Y, field_data, periodic_fix='sphere')
+
+ # Plot data ----------------------------------------------------------------
+ cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove)
+ cf, _ = plot_contoured_field(
+ ax,
+ lon_2d,
+ lat_2d,
+ field_data,
+ contour_method,
+ contours,
+ cmap=cmap,
+ line_contours=lines,
+ projection=projection,
+ )
+
+ add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10, cbar_format=cbar_format[field_name])
+ tomplot_field_title(ax, None, minmax=True, field_data=field_data)
+
+ # Save figure ------------------------------------------------------------------
+ plt.suptitle(f't = {time:.1f} days')
+ plot_name = f'{plot_stem}_final.png'
+ print(f'Saving figure to {plot_name}')
+ fig.savefig(plot_name, bbox_inches='tight')
+ plt.close()
+
+
+if __name__ == "__main__":
+ from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
+
+ parser = ArgumentParser(description=__doc__, formatter_class=ArgumentDefaultsHelpFormatter)
+ parser.add_argument(
+ '--results_file_name',
+ help='Path to netCDF result file from the Williamson 5 test case',
+ type=str,
+ default='./results/williamson_5/field_output.nc',
+ )
+ parser.add_argument('--plot_stem', help='Name of the plots', type=str, default='williamson_5')
+ args, unknown = parser.parse_known_args()
+
+ plot_williamson_5(**vars(args))
diff --git a/pySDC/tutorial/step_7/F_pySDC_with_Gusto.py b/pySDC/tutorial/step_7/F_pySDC_with_Gusto.py
new file mode 100644
index 0000000000000000000000000000000000000000..2053f79ce0ca83f51883887631a253526d66ff16
--- /dev/null
+++ b/pySDC/tutorial/step_7/F_pySDC_with_Gusto.py
@@ -0,0 +1,346 @@
+"""
+Example for running pySDC together with Gusto. This test runs a shallow water equation and may take a considerable
+amount of time. After you have run it, move on to step F_2, which includes a plotting script.
+
+This is Test Case 5 (flow over a mountain) of Williamson et al, 1992:
+``A standard test set for numerical approximations to the shallow water
+equations in spherical geometry'', JCP.
+
+This script is adapted from the Gusto example: https://github.com/firedrakeproject/gusto/blob/main/examples/shallow_water/williamson_5.py
+
+The pySDC coupling works by setting up pySDC as a time integrator within Gusto.
+To this end, you need to construct a pySDC description and controller parameters as usual and pass them when
+constructing the pySDC time discretization.
+
+After passing this to a Gusto timestepper, you have two choices:
+ - Access the `.scheme.controller` variable of the timestepper, which is the pySDC controller and use pySDC for
+ running
+ - Use the Gusto timestepper for running
+You may wonder why it is necessary to construct a Gusto timestepper if you don't want to use it. The reason is the
+setup of spatial methods, such as upwinding. These are passed to the Gusto timestepper and modify the residual of the
+equations during its instantiation. Once the residual is modified, we can choose whether to continue in Gusto or pySDC.
+
+This script supports space-time parallelism, as well as running the Gusto SDC implementation or the pySDC-Gusto coupling.
+Please run with `--help` to learn how to configure this script.
+"""
+
+import firedrake as fd
+from pySDC.helpers.pySDC_as_gusto_time_discretization import pySDC_integrator
+from pySDC.helpers.firedrake_ensemble_communicator import FiredrakeEnsembleCommunicator
+from gusto import SDC, BackwardEuler
+from gusto.core.labels import implicit, time_derivative
+
+
+from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
+from firedrake import SpatialCoordinate, as_vector, pi, sqrt, min_value, Function
+from gusto import (
+ Domain,
+ IO,
+ OutputParameters,
+ DGUpwind,
+ ShallowWaterParameters,
+ ShallowWaterEquations,
+ Sum,
+ lonlatr_from_xyz,
+ GeneralIcosahedralSphereMesh,
+ ZonalComponent,
+ MeridionalComponent,
+ RelativeVorticity,
+ Timestepper,
+)
+
+williamson_5_defaults = {
+ 'ncells_per_edge': 12, # number of cells per icosahedron edge
+ 'dt': 900.0,
+ 'tmax': 50.0 * 24.0 * 60.0 * 60.0, # 50 days
+ 'dumpfreq': 10, # output every <dumpfreq> steps
+ 'dirname': 'williamson_5', # results will go into ./results/<dirname>
+ 'time_parallelism': False, # use parallel diagonal SDC or not
+ 'QI': 'MIN-SR-S', # implicit preconditioner
+ 'M': '3', # number of collocation nodes
+ 'kmax': '5', # use fixed number of iteration up to this value
+ 'use_pySDC': True, # whether to use pySDC for time integration
+ 'use_adaptivity': True, # whether to use adaptive step size selection
+}
+
+
+def williamson_5(
+ ncells_per_edge=williamson_5_defaults['ncells_per_edge'],
+ dt=williamson_5_defaults['dt'],
+ tmax=williamson_5_defaults['tmax'],
+ dumpfreq=williamson_5_defaults['dumpfreq'],
+ dirname=williamson_5_defaults['dirname'],
+ time_parallelism=williamson_5_defaults['time_parallelism'],
+ QI=williamson_5_defaults['QI'],
+ M=williamson_5_defaults['M'],
+ kmax=williamson_5_defaults['kmax'],
+ use_pySDC=williamson_5_defaults['use_pySDC'],
+ use_adaptivity=williamson_5_defaults['use_adaptivity'],
+ mesh=None,
+):
+ """
+ Run the Williamson 5 test case.
+
+ Args:
+ ncells_per_edge (int): number of cells per icosahedron edge
+ dt (float): Initial step size
+ tmax (float): Time to integrate to
+ dumpfreq (int): Output every <dumpfreq> time steps
+ dirname (str): Output will go into ./results/<dirname>
+ time_parallelism (bool): True for parallel SDC, False for serial
+ M (int): Number of collocation nodes
+ kmax (int): Max number of SDC iterations
+ use_pySDC (bool): Use pySDC as Gusto time integrator or Gusto SDC implementation
+ """
+ if not use_pySDC and use_adaptivity:
+ raise NotImplementedError('Adaptive step size selection not yet implemented in Gusto')
+
+ # ------------------------------------------------------------------------ #
+ # Parameters for test case
+ # ------------------------------------------------------------------------ #
+
+ radius = 6371220.0 # planetary radius (m)
+ mean_depth = 5960 # reference depth (m)
+ g = 9.80616 # acceleration due to gravity (m/s^2)
+ u_max = 20.0 # max amplitude of the zonal wind (m/s)
+ mountain_height = 2000.0 # height of mountain (m)
+ R0 = pi / 9.0 # radius of mountain (rad)
+ lamda_c = -pi / 2.0 # longitudinal centre of mountain (rad)
+ phi_c = pi / 6.0 # latitudinal centre of mountain (rad)
+
+ # ------------------------------------------------------------------------ #
+ # Our settings for this set up
+ # ------------------------------------------------------------------------ #
+
+ element_order = 1
+
+ # ------------------------------------------------------------------------ #
+ # Set up model objects
+ # ------------------------------------------------------------------------ #
+
+ # parallelism
+ if time_parallelism:
+ ensemble_comm = FiredrakeEnsembleCommunicator(fd.COMM_WORLD, fd.COMM_WORLD.size // M)
+ space_comm = ensemble_comm.space_comm
+ from pySDC.implementations.sweeper_classes.generic_implicit_MPI import generic_implicit_MPI as sweeper_class
+
+ if ensemble_comm.time_comm.rank > 0:
+ dirname = f'{dirname}-{ensemble_comm.time_comm.rank}'
+ else:
+ ensemble_comm = None
+ space_comm = fd.COMM_WORLD
+ from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit as sweeper_class
+
+ # Domain
+ mesh = GeneralIcosahedralSphereMesh(radius, ncells_per_edge, degree=2, comm=space_comm) if mesh is None else mesh
+ domain = Domain(mesh, dt, 'BDM', element_order)
+ x, y, z = SpatialCoordinate(mesh)
+ lamda, phi, _ = lonlatr_from_xyz(x, y, z)
+
+ # Equation: coriolis
+ parameters = ShallowWaterParameters(H=mean_depth, g=g)
+ Omega = parameters.Omega
+ fexpr = 2 * Omega * z / radius
+
+ # Equation: topography
+ rsq = min_value(R0**2, (lamda - lamda_c) ** 2 + (phi - phi_c) ** 2)
+ r = sqrt(rsq)
+ tpexpr = mountain_height * (1 - r / R0)
+ eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, topog_expr=tpexpr)
+
+ eqns.label_terms(lambda t: not t.has_label(time_derivative), implicit)
+
+ # I/O
+ output = OutputParameters(
+ dirname=dirname,
+ dumplist_latlon=['D'],
+ dumpfreq=dumpfreq,
+ dump_vtus=True,
+ dump_nc=True,
+ dumplist=['D', 'topography'],
+ )
+ diagnostic_fields = [Sum('D', 'topography'), RelativeVorticity(), MeridionalComponent('u'), ZonalComponent('u')]
+ io = IO(domain, output, diagnostic_fields=diagnostic_fields)
+
+ # Transport schemes
+ transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")]
+
+ # ------------------------------------------------------------------------ #
+ # pySDC parameters: description and controller parameters
+ # ------------------------------------------------------------------------ #
+
+ level_params = dict()
+ level_params['restol'] = -1
+ level_params['dt'] = dt
+ level_params['residual_type'] = 'full_rel'
+
+ step_params = dict()
+ step_params['maxiter'] = kmax
+
+ sweeper_params = dict()
+ sweeper_params['quad_type'] = 'RADAU-RIGHT'
+ sweeper_params['node_type'] = 'LEGENDRE'
+ sweeper_params['num_nodes'] = M
+ sweeper_params['QI'] = QI
+ sweeper_params['QE'] = 'PIC'
+ sweeper_params['comm'] = ensemble_comm
+ sweeper_params['initial_guess'] = 'copy'
+
+ problem_params = dict()
+
+ convergence_controllers = {}
+ if use_adaptivity:
+ from pySDC.implementations.convergence_controller_classes.adaptivity import Adaptivity
+ from pySDC.implementations.convergence_controller_classes.spread_step_sizes import (
+ SpreadStepSizesBlockwiseNonMPI,
+ )
+
+ convergence_controllers[Adaptivity] = {'e_tol': 1e-6, 'rel_error': True, 'dt_max': 1e4, 'dt_rel_min_slope': 0.5}
+ # this is needed because the coupling runs on the controller level and this will almost always overwrite
+ convergence_controllers[SpreadStepSizesBlockwiseNonMPI] = {'overwrite_to_reach_Tend': False}
+
+ controller_params = dict()
+ controller_params['logger_level'] = 15 if fd.COMM_WORLD.rank == 0 else 30
+ controller_params['mssdc_jac'] = False
+
+ description = dict()
+ description['problem_params'] = problem_params
+ description['sweeper_class'] = sweeper_class
+ description['sweeper_params'] = sweeper_params
+ description['level_params'] = level_params
+ description['step_params'] = step_params
+ description['convergence_controllers'] = convergence_controllers
+
+ # ------------------------------------------------------------------------ #
+ # petsc solver parameters
+ # ------------------------------------------------------------------------ #
+
+ solver_parameters = {
+ 'snes_type': 'newtonls',
+ 'ksp_type': 'gmres',
+ 'pc_type': 'bjacobi',
+ 'sub_pc_type': 'ilu',
+ 'ksp_rtol': 1e-12,
+ 'snes_rtol': 1e-12,
+ 'ksp_atol': 1e-30,
+ 'snes_atol': 1e-30,
+ 'ksp_divtol': 1e30,
+ 'snes_divtol': 1e30,
+ 'snes_max_it': 999,
+ 'ksp_max_it': 999,
+ }
+
+ # ------------------------------------------------------------------------ #
+ # Set Gusto SDC parameters to match the pySDC ones
+ # ------------------------------------------------------------------------ #
+
+ SDC_params = {
+ 'base_scheme': BackwardEuler(domain, solver_parameters=solver_parameters),
+ 'M': sweeper_params['num_nodes'],
+ 'maxk': step_params['maxiter'],
+ 'quad_type': sweeper_params['quad_type'],
+ 'node_type': sweeper_params['node_type'],
+ 'qdelta_imp': sweeper_params['QI'],
+ 'qdelta_exp': sweeper_params['QE'],
+ 'formulation': 'Z2N',
+ 'initial_guess': 'copy',
+ 'nonlinear_solver_parameters': solver_parameters,
+ 'linear_solver_parameters': solver_parameters,
+ 'final_update': False,
+ }
+
+ # ------------------------------------------------------------------------ #
+ # Setup time stepper
+ # ------------------------------------------------------------------------ #
+
+ if use_pySDC:
+ method = pySDC_integrator(
+ eqns, description, controller_params, domain=domain, solver_parameters=solver_parameters
+ )
+ else:
+ method = SDC(**SDC_params, domain=domain)
+
+ stepper = Timestepper(eqns, method, io, spatial_methods=transport_methods)
+
+ if use_pySDC and use_adaptivity:
+ # we have to do this for adaptive time stepping, because it is a bit of a mess
+ method.timestepper = stepper
+
+ # ------------------------------------------------------------------------ #
+ # Initial conditions
+ # ------------------------------------------------------------------------ #
+
+ u0 = stepper.fields('u')
+ D0 = stepper.fields('D')
+ uexpr = as_vector([-u_max * y / radius, u_max * x / radius, 0.0])
+ Dexpr = mean_depth - tpexpr - (radius * Omega * u_max + 0.5 * u_max**2) * (z / radius) ** 2 / g
+
+ u0.project(uexpr)
+ D0.interpolate(Dexpr)
+
+ Dbar = Function(D0.function_space()).assign(mean_depth)
+ stepper.set_reference_profiles([('D', Dbar)])
+
+ # ------------------------------------------------------------------------ #
+ # Run
+ # ------------------------------------------------------------------------ #
+
+ stepper.run(t=0, tmax=tmax)
+ return stepper, mesh
+
+
+# ---------------------------------------------------------------------------- #
+# MAIN
+# ---------------------------------------------------------------------------- #
+
+
+if __name__ == "__main__":
+
+ parser = ArgumentParser(description=__doc__, formatter_class=ArgumentDefaultsHelpFormatter)
+ parser.add_argument(
+ '--ncells_per_edge',
+ help="The number of cells per edge of icosahedron",
+ type=int,
+ default=williamson_5_defaults['ncells_per_edge'],
+ )
+ parser.add_argument('--dt', help="The time step in seconds.", type=float, default=williamson_5_defaults['dt'])
+ parser.add_argument(
+ "--tmax", help="The end time for the simulation in seconds.", type=float, default=williamson_5_defaults['tmax']
+ )
+ parser.add_argument(
+ '--dumpfreq',
+ help="The frequency at which to dump field output.",
+ type=int,
+ default=williamson_5_defaults['dumpfreq'],
+ )
+ parser.add_argument(
+ '--dirname', help="The name of the directory to write to.", type=str, default=williamson_5_defaults['dirname']
+ )
+ parser.add_argument(
+ '--time_parallelism',
+ help="Whether to use parallel diagonal SDC or not.",
+ type=str,
+ default=williamson_5_defaults['time_parallelism'],
+ )
+ parser.add_argument('--kmax', help='SDC iteration count', type=int, default=williamson_5_defaults['kmax'])
+ parser.add_argument('-M', help='SDC node count', type=int, default=williamson_5_defaults['M'])
+ parser.add_argument(
+ '--use_pySDC',
+ help='whether to use pySDC or Gusto SDC implementation',
+ type=str,
+ default=williamson_5_defaults['use_pySDC'],
+ )
+ parser.add_argument(
+ '--use_adaptivity',
+ help='whether to use adaptive step size selection',
+ type=str,
+ default=williamson_5_defaults['use_adaptivity'],
+ )
+ parser.add_argument('--QI', help='Implicit preconditioner', type=str, default=williamson_5_defaults['QI'])
+ args, unknown = parser.parse_known_args()
+
+ options = vars(args)
+ for key in ['use_pySDC', 'use_adaptivity', 'time_parallelism']:
+ options[key] = options[key] not in ['False', 0, False, 'false']
+
+ williamson_5(**options)
diff --git a/pySDC/tutorial/step_7/README.rst b/pySDC/tutorial/step_7/README.rst
index d21f6cd7e066b01c1a0ffb8a2bb903bfb98964dc..7aa37a1d721329f13c19d5dfd9becf5b80ba186a 100644
--- a/pySDC/tutorial/step_7/README.rst
+++ b/pySDC/tutorial/step_7/README.rst
@@ -64,3 +64,28 @@ This is work in progress in very early stages! The tensor datatype is the simple
If you want to work on this, your input is appreciated!
.. include:: doc_step_7_D.rst
+
+
+Part E: pySDC and Firedrake
+---------------------------
+
+`Firedrake <https://github.com/firedrakeproject/firedrake>`_ is a finite element library with similar features as FEniCS.
+The below example runs the same heat equation as in the FEniCS example, but implemented in Firedrake.
+See the problem class implementation as a blueprint for how to implement problems with Firedrake in a way that pySDC can understand: `pySDC/implementations/problem_classes/HeatFiredrake.py <https://github.com/Parallel-in-Time/pySDC/blob/master/pySDC/implementations/problem_classes/HeatFiredrake.py>`_
+
+.. include:: doc_step_7_E.rst
+
+
+Part F: pySDC and Gusto
+---------------------------
+
+`Gusto <https://github.com/firedrakeproject/gusto>`_ is a toolkit for geophysical simulations that uses `Firedrake <https://github.com/firedrakeproject/firedrake>`_ for spatial discretization.
+The below example is an adaptation of the Williamson 5 test case as implemented in Gusto.
+This coupling works slightly different to the other examples within this tutorial, as timestepping is part of Gusto.
+The aim of the coupling is not a spatial discretization, but to use the equations that are implemented in Gusto.
+A Gusto equation includes the basic form of the equation set, but a crucial part is to modify terms in the discretized equations with spatial methods, such as upwinding schemes.
+We get the finished equation set into pySDC by setting up pySDC as a time discretization for Gusto and instantiating a Gusto timestepper.
+During this instantiation the equation, and the residual that is used for solving systems, is modified with all the spatial methods.
+Afterwards, you have a Gusto timestepping scheme that you can run in Gusto and a pySDC controller that you can run by itself.
+
+.. include:: doc_step_7_F.rst