Non-split physics for all explicit or IMEX time discretisations

gusto/time_discretisation/explicit_runge_kutta.py

@@ -379,11 +379,8 @@ def solve_stage(self, x0, stage):
 
             # Update field_i for physics / limiters
             for evaluate in self.evaluate_source:
-                # TODO: not implemented! Here we need to evaluate the m-th term
                 # in the i-th RHS with field_m
-                raise NotImplementedError(
-                    'Physics not implemented with RK schemes that use the '
-                    + 'predictor form')
+                evaluate(self.field_i[stage], self.dt)
             if self.limiter is not None:
                 self.limiter.apply(self.field_i[stage])
 

gusto/time_discretisation/imex_runge_kutta.py

@@ -239,16 +239,23 @@ def apply(self, x_out, x_in):
 
         for stage in range(self.nStages):
             self.solver = solver_list[stage]
-            # Set initial solver guess
+            # Set initial solver guess and evalute physics terms
             if (stage > 0):
                 self.x_out.assign(self.xs[stage-1])
+                for evaluate in self.evaluate_source:
+                    evaluate(self.xs[stage-1], self.dt)
             self.solver.solve()
 
             # Apply limiter
             if self.limiter is not None:
                 self.limiter.apply(self.x_out)
             self.xs[stage].assign(self.x_out)
 
+        # Evaluate physics terms
+        for evaluate in self.evaluate_source:
+            evaluate(self.xs[self.nStages-1], self.dt)
+
+        # Solve final stage
         self.final_solver.solve()
 
         # Apply limiter

gusto/time_discretisation/multi_level_schemes.py

@@ -223,6 +223,11 @@ def apply(self, x_out, *x_in):
 
         self.xnm1.assign(x_in[0])
         self.x1.assign(x_in[1])
+
+        # Evaluate physics terms
+        for evaluate in self.evaluate_source:
+            evaluate(self.x1, self.dt)
+
         # Set initial solver guess
         self.x_out.assign(x_in[1])
         solver.solve()
@@ -356,6 +361,9 @@ def apply(self, x_out, *x_in):
 
         for n in range(self.nlevels):
             self.x[n].assign(x_in[n])
+            # Evaluate physics terms
+            for evaluate in self.evaluate_source:
+                evaluate(self.x[n], self.dt)
         # Set initial solver guess
         self.x_out.assign(x_in[-1])
         solver.solve()

gusto/time_discretisation/sdc.py

@@ -206,17 +206,14 @@ def setup(self, equation, apply_bcs=True, *active_labels):
         self.base.setup(equation, apply_bcs, *active_labels)
         self.equation = self.base.equation
         self.residual = self.base.residual
+        self.evaluate_source = self.base.evaluate_source
 
         for t in self.residual:
             # Check all terms are labeled implicit or explicit
             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")
 
-            # Check we are not using wrappers for implicit schemes
-            if (t.has_label(implicit) and self.wrapper is not None):
-                raise NotImplementedError("Implicit terms not supported with wrappers")
-
             # Check we are not using limiters for implicit schemes
             if (t.has_label(implicit) and self.limiter is not None):
                 raise NotImplementedError("Implicit terms not supported with limiters")
@@ -234,7 +231,7 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                         raise ValueError(f"The option defined for {field} is for a field that does not exist in the equation set")
 
                     field_idx = equation.field_names.index(field)
-                    subwrapper.setup(equation.spaces[Wrappersfield_idx])
+                    subwrapper.setup(equation.spaces[field_idx])
 
                     # Update the function space to that needed by the wrapper
                     self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space
@@ -484,6 +481,9 @@ def apply(self, x_out, x_in):
             # for Z2N: sum(j=1,M) (q_mj*F(y_m^k) +  q_mj*S(y_m^k))
             for m in range(1, self.M+1):
                 self.Uin.assign(self.Unodes[m])
+                # Include physics terms
+                for evaluate in self.evaluate_source:
+                    evaluate(self.Uin, self.base.dt)
                 self.solver_rhs.solve()
                 self.fUnodes[m-1].assign(self.Urhs)
             self.compute_quad()
@@ -500,6 +500,10 @@ def apply(self, x_out, x_in):
                 self.solver = solver_list[m-1]
                 self.U_SDC.assign(self.Unodes[m])
 
+                # Evaluate physics terms
+                for evaluate in self.evaluate_source:
+                    evaluate(self.Unodes[m], self.base.dt)
+
                 # Compute
                 # for N2N:
                 # y_m^(k+1) = y_(m-1)^(k+1) + dtau_m*(F(y_(m)^(k+1)) - F(y_(m)^k)
@@ -522,10 +526,17 @@ def apply(self, x_out, x_in):
             if self.final_update:
                 for m in range(1, self.M+1):
                     self.Uin.assign(self.Unodes1[m])
+                    # Include physics terms
+                    for evaluate in self.evaluate_source:
+                        evaluate(self.Uin, self.base.dt)
                     self.solver_rhs.solve()
                     self.fUnodes[m-1].assign(self.Urhs)
                 self.compute_quad_final()
                 # Compute y_(n+1) = y_n + sum(j=1,M) q_j*F(y_j)
+
+                # Evaluate physics terms
+                for evaluate in self.evaluate_source:
+                    evaluate(self.Unodes[-1], self.base.dt)
                 self.U_fin.assign(self.Unodes[-1])
                 self.solver_fin.solve()
                 # Apply limiter if required

integration-tests/model/test_nonsplit_physics.py

@@ -0,0 +1,147 @@
+"""
+This script tests the non-split timestepper against the split timestepper
+using an advection equation with a physics parametrisation.
+One split method is tested, whilst different nonsplit IMEX and explicit time
+discretisations are used for the dynamics and physics.
+"""
+
+from firedrake import (SpatialCoordinate, PeriodicIntervalMesh, exp, as_vector,
+                       norm, Constant, conditional, sqrt, VectorFunctionSpace)
+from gusto import *
+import pytest
+
+
+def run_nonsplit_adv_physics(tmpdir, timestepper):
+    """
+    Runs the advection equation with a physics parametrisation using different timesteppers.
+    """
+
+    # ------------------------------------------------------------------------ #
+    # Set up model objects
+    # ------------------------------------------------------------------------ #
+
+    # Domain
+    dt = 0.01
+    tmax = 0.75
+    L = 10
+    mesh = PeriodicIntervalMesh(20, L)
+    domain = Domain(mesh, dt, "CG", 1)
+
+    # Equation
+    V = domain.spaces("DG")
+    Vu = VectorFunctionSpace(mesh, "CG", 1)
+    equation = ContinuityEquation(domain, V, "f", Vu=Vu)
+    spatial_methods = [DGUpwind(equation, "f")]
+
+    x = SpatialCoordinate(mesh)
+
+    # Add a source term to inject mass into the domain.
+    source_expression = -Constant(0.5)
+    physics_schemes = [(SourceSink(equation, "f", source_expression), SSPRK3(domain))]
+
+    # I/O
+    output = OutputParameters(dirname=str(tmpdir), dumpfreq=25)
+    io = IO(domain, output)
+
+    time_varying_velocity = False
+
+    # Time stepper
+    if timestepper == 'split':
+        # Split with no defined weights
+        dynamics_schemes = {'transport': ImplicitMidpoint(domain)}
+        term_splitting = ['transport', 'physics']
+        stepper = SplitTimestepper(equation, term_splitting, dynamics_schemes,
+                                   io, spatial_methods=spatial_methods,
+                                   physics_schemes=physics_schemes)
+    elif timestepper == 'nonsplit_imex_rk':
+        # Split continuity term
+        equation = split_continuity_form(equation)
+        # Label terms as implicit and explicit
+        equation.label_terms(lambda t: not any(t.has_label(time_derivative, transport)), implicit)
+        equation.label_terms(lambda t: t.has_label(transport), explicit)
+        dynamics_schemes = IMEX_SSP3(domain)
+        stepper = PrescribedTransport(equation, dynamics_schemes,
+                                      io, time_varying_velocity,
+                                      transport_method=spatial_methods)
+    elif timestepper == 'nonsplit_exp_rk':
+        dynamics_schemes = SSPRK3(domain)
+        stepper = PrescribedTransport(equation, dynamics_schemes,
+                                      io, time_varying_velocity,
+                                      transport_method=spatial_methods)
+    elif timestepper == 'nonsplit_imex_sdc':
+        # Split continuity term
+        equation = split_continuity_form(equation)
+        # Label terms as implicit and explicit
+        equation.label_terms(lambda t: not any(t.has_label(time_derivative, transport)), implicit)
+        equation.label_terms(lambda t: t.has_label(transport), explicit)
+
+        node_type = "LEGENDRE"
+        qdelta_imp = "LU"
+        qdelta_exp = "FE"
+        quad_type = "RADAU-RIGHT"
+        M = 2
+        k = 2
+        base_scheme = IMEX_Euler(domain)
+        dynamics_schemes = SDC(base_scheme, domain, M, k, quad_type, node_type, qdelta_imp,
+                               qdelta_exp, formulation="Z2N", final_update=True, initial_guess="base")
+        stepper = PrescribedTransport(equation, dynamics_schemes,
+                                      io, time_varying_velocity,
+                                      transport_method=spatial_methods)
+    elif timestepper == 'nonsplit_exp_multistep':
+        dynamics_schemes = AdamsBashforth(domain, order=2)
+        stepper = PrescribedTransport(equation, dynamics_schemes,
+                                      io, time_varying_velocity,
+                                      transport_method=spatial_methods)
+
+    # ------------------------------------------------------------------------ #
+    # Initial conditions
+    # ------------------------------------------------------------------------ #
+
+    xc_init = 0.25 * L
+    xc_end = 0.75 * L
+    umax = 0.5 * L / tmax
+
+    # Get minimum distance on periodic interval to xc
+    x_init = conditional(sqrt((x[0] - xc_init) ** 2) < 0.5 * L,
+                         x[0] - xc_init, L + x[0] - xc_init)
+
+    x_end = conditional(sqrt((x[0] - xc_end) ** 2) < 0.5 * L,
+                        x[0] - xc_end, L + x[0] - xc_end)
+
+    f_init = 5.0
+    f_end = f_init
+    f_width_init = L / 10.0
+    f_width_end = f_width_init
+    f_init_expr = f_init * exp(-(x_init / f_width_init) ** 2)
+
+    # The end Gaussian should be advected by half the domain
+    # length and include more mass due to the source term.
+    f_end_expr = 0.5 + f_end * exp(-(x_end / f_width_end) ** 2)
+
+    stepper.fields('f').interpolate(f_init_expr)
+    stepper.fields('u').interpolate(as_vector([Constant(umax)]))
+    f_end = stepper.fields('f_end', space=V)
+    f_end.interpolate(f_end_expr)
+
+    # ------------------------------------------------------------------------ #
+    # Run
+    # ------------------------------------------------------------------------ #
+
+    stepper.run(0, tmax=tmax)
+
+    error = norm(stepper.fields('f') - f_end) / norm(f_end)
+
+    return error
+
+
+@pytest.mark.parametrize("timestepper", ["split", "nonsplit_imex_rk", "nonsplit_imex_sdc",
+                                         "nonsplit_exp_rk", "nonsplit_exp_multistep"])
+def test_nonsplit_adv_physics(tmpdir, timestepper):
+    """
+    Test the nonsplit timestepper in the advection equation with source physics.
+    """
+    tol = 0.12
+    error = run_nonsplit_adv_physics(tmpdir, timestepper)
+    assert error < tol, 'The nonsplit timestepper in the advection' + \
+                        'equation with source physics has an error greater than ' + \
+                        'the permitted tolerance'