IMEX Multistage

gusto/common_forms.py

@@ -2,13 +2,17 @@
 Provides some basic forms for discretising various common terms in equations for
 geophysical fluid dynamics."""
 
-from firedrake import dx, dot, grad, div, inner, outer, cross, curl
+from firedrake import (dx, dot, grad, div, inner, outer, cross, curl, split,
+                       TestFunctions, TrialFunction)
+from firedrake.fml import subject, drop
 from gusto.configuration import TransportEquationType
-from gusto.labels import transport, transporting_velocity, diffusion
+from gusto.labels import (transport, transporting_velocity, diffusion,
+                          prognostic, linearisation)
 
 __all__ = ["advection_form", "continuity_form", "vector_invariant_form",
            "kinetic_energy_form", "advection_equation_circulation_form",
-           "diffusion_form", "linear_advection_form", "linear_continuity_form"]
+           "diffusion_form", "linear_advection_form", "linear_continuity_form",
+           "split_continuity_form"]
 
 
 def advection_form(test, q, ubar):
@@ -194,3 +198,65 @@ def diffusion_form(test, q, kappa):
     form = -inner(test, div(kappa*grad(q)))*dx
 
     return diffusion(form)
+
+
+def split_continuity_form(equation):
+    u"""
+    Loops through terms in a given equation, and splits continuity terms into
+    advective and divergence terms.
+
+    This describes splitting ∇.(u*q) into u.∇q and q(∇.u),
+    for transporting velocity u and transported q.
+
+    Args:
+        equation (:class:`PrognosticEquation`): the model's equation.
+
+    Returns:
+        equation (:class:`PrognosticEquation`): the model's equation.
+    """
+
+    for t in equation.residual:
+        if (t.get(transport) == TransportEquationType.conservative):
+            # Get fields and test functions
+            subj = t.get(subject)
+            prognostic_field_name = t.get(prognostic)
+            if hasattr(equation, "field_names"):
+                idx = equation.field_names.index(prognostic_field_name)
+            else:
+                idx = equation.field_name.index(prognostic_field_name)
+            W = equation.function_space
+            test = TestFunctions(W)[idx]
+            q = split(subj)[idx]
+            # u is either a prognostic or prescribed field
+            if (hasattr(equation, "field_names")
+               and 'u' in equation.field_names):
+                u_idx = equation.field_names.index('u')
+                uadv = split(equation.X)[u_idx]
+            elif 'u' in equation.prescribed_fields._field_names:
+                uadv = equation.prescribed_fields('u')
+            else:
+                raise ValueError('Cannot get velocity field')
+
+            # Create new advective and divergence terms
+            adv_term = prognostic(advection_form(test, q, uadv), prognostic_field_name)
+            div_term = prognostic(test*q*div(uadv)*dx, prognostic_field_name)
+
+            # Add linearisations of new terms if required
+            if (t.has_label(linearisation)):
+                u_trial = TrialFunction(W)[u_idx]
+                qbar = split(equation.X_ref)[idx]
+                # Add linearisation to adv_term
+                linear_adv_term = linear_advection_form(test, qbar, u_trial)
+                adv_term = linearisation(adv_term, linear_adv_term)
+                # Add linearisation to div_term
+                linear_div_term = transporting_velocity(qbar*test*div(u_trial)*dx, u_trial)
+                div_term = linearisation(div_term, linear_div_term)
+
+            # Add new terms onto residual
+            equation.residual += subject(adv_term + div_term, subj)
+            # Drop old term
+            equation.residual = equation.residual.label_map(
+                lambda t: t.get(transport) == TransportEquationType.conservative,
+                map_if_true=drop)
+
+    return equation

gusto/time_discretisation.py

@@ -6,19 +6,22 @@
 """
 
 from abc import ABCMeta, abstractmethod, abstractproperty
-from firedrake import (Function, TestFunction, TestFunctions,
-                       NonlinearVariationalProblem, NonlinearVariationalSolver,
-                       DirichletBC, Constant, split, div, dx)
+import math
+import numpy as np
+
+from firedrake import (
+    Function, TestFunction, NonlinearVariationalProblem,
+    NonlinearVariationalSolver, DirichletBC, split, Constant
+)
+from firedrake.fml import (
+    replace_subject, replace_test_function, Term, all_terms, drop
+)
 from firedrake.formmanipulation import split_form
 from firedrake.utils import cached_property
 
-from gusto.configuration import EmbeddedDGOptions, RecoveryOptions, TransportEquationType
-from gusto.fml import (
-    replace_subject, replace_test_function, Term, all_terms, drop, subject
-)
+from gusto.configuration import EmbeddedDGOptions, RecoveryOptions
 from gusto.labels import (time_derivative, prognostic, physics_label,
-                          transport, implicit, explicit)
-from gusto.common_forms import advection_form
+                          implicit, explicit)
 from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual
 from gusto.wrappers import *
 
@@ -92,7 +95,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
             elif self.wrapper_name == "supg":
                 self.wrapper = SUPGWrapper(self, options)
             else:
-                raise NotImplementedError(
+                raise RuntimeError(
                     f'Time discretisation: wrapper {self.wrapper_name} not implemented')
         else:
             self.wrapper = None
@@ -327,38 +330,11 @@ def setup(self, equation, apply_bcs=True, *active_labels):
 
         super().setup(equation, apply_bcs, *active_labels)
 
-        # Get continuity form transport term
+        # Check all terms are labeled implicit, exlicit
         for t in self.residual:
-            if (t.get(transport) == TransportEquationType.conservative):
-                # Split continuity form term
-                subj = t.get(subject)
-                prognostic_field_name = t.get(prognostic)
-                idx = self.equation.field_names.index(prognostic_field_name)
-                W = self.fs
-                test = TestFunctions(W)[idx]
-                transported_field = split(subj)[idx]
-                u_idx = self.equation.field_names.index('u')
-                uadv = split(self.equation.X)[u_idx]
-                new_transport_term = prognostic(advection_form(test, transported_field, uadv), prognostic_field_name)
-                div_term = prognostic(test*transported_field*div(uadv)*dx, prognostic_field_name)
-                # Add onto residual
-                self.residual += subject(new_transport_term + div_term, subj)
-
-        # Drop old term
-        self.residual = self.residual.label_map(
-            lambda t: t.get(transport) == TransportEquationType.conservative,
-            map_if_true=drop)
-
-        # Label transport terms as explicit, all other terms as implicit
-        self.residual = self.residual.label_map(
-            lambda t: any(t.has_label(time_derivative, transport)),
-            map_if_false=lambda t: implicit(t))
-
-        self.residual = self.residual.label_map(
-            lambda t: t.has_label(transport),
-            map_if_true=lambda t: explicit(t))
-
-        logger.warning("Default IMEX Multistage treats transport terms explicitly, and all other terms implicitly")
+            if ((not t.has_label(implicit)) and (not t.has_label(explicit))
+               and (not t.has_label(time_derivative))):
+                raise NotImplementedError("Non time-derivative terms must be labeled as implicit or explicit")
 
         self.xs = [Function(self.fs) for i in range(self.nStages)]
 

integration-tests/model/test_IMEX.py

@@ -0,0 +1,42 @@
+from firedrake import norm
+from gusto import *
+import pytest
+
+
+def run(timestepper, tmax, f_end):
+    timestepper.run(0, tmax)
+    return norm(timestepper.fields("f") - f_end) / norm(f_end)
+
+
+@pytest.mark.parametrize("scheme", ["ssp3", "ark2", "ars3", "trap2", "euler"])
+def test_time_discretisation(tmpdir, scheme, tracer_setup):
+
+    setup = tracer_setup(tmpdir, "sphere")
+    domain = setup.domain
+    V = domain.spaces("DG")
+    eqn = ContinuityEquation(domain, V, "f")
+    # Split continuity term
+    eqn = split_continuity_form(eqn)
+    # Label terms are implicit and explicit
+    eqn.label_terms(lambda t: not any(t.has_label(time_derivative, transport)), implicit)
+    eqn.label_terms(lambda t: t.has_label(transport), explicit)
+
+    if scheme == "ssp3":
+        transport_scheme = SSP3(domain)
+    elif scheme == "ark2":
+        transport_scheme = ARK2(domain)
+    elif scheme == "ars3":
+        transport_scheme = ARS3(domain)
+    elif scheme == "trap2":
+        transport_scheme = Trap2(domain)
+    elif scheme == "euler":
+        transport_scheme = IMEX_Euler(domain)
+
+    transport_method = DGUpwind(eqn, "f")
+
+    timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method)
+
+    # Initial conditions
+    timestepper.fields("f").interpolate(setup.f_init)
+    timestepper.fields("u").project(setup.uexpr)
+    assert run(timestepper, setup.tmax, setup.f_end) < setup.tol