diff --git a/.github/workflows/ci_pipeline.yml b/.github/workflows/ci_pipeline.yml index 3487393df4..2bfbdfe521 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 0000000000..a0d1875528 --- /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 `_ + +.. 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 0000000000..f706dd231b --- /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 `_ +Find a suitable plotting script here: `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 0000000000..33bc722878 --- /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 0000000000..3d8ef50f77 --- /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 0000000000..69f8611819 --- /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 56aa6cb83f..6820a17d15 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 0000000000..8d091e3337 --- /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 0000000000..2053f79ce0 --- /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 steps + 'dirname': 'williamson_5', # results will go into ./results/ + '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 time steps + dirname (str): Output will go into ./results/ + 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 d21f6cd7e0..7aa37a1d72 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 `_ 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 `_ + +.. include:: doc_step_7_E.rst + + +Part F: pySDC and Gusto +--------------------------- + +`Gusto `_ is a toolkit for geophysical simulations that uses `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