Get SUPG working for MixedFunctionSpace (#580)

Co-authored-by: Thomas Bendall <[email protected]>

gusto/core/configuration.py

@@ -203,14 +203,25 @@ class SUPGOptions(WrapperOptions):
     default = 1/sqrt(15)
     ibp = IntegrateByParts.TWICE
 
+    # Dictionary containing keys field_name and values term_labels
+    # field_name (str): name of the field for SUPG to be applied to
+    # term_label (list): labels of terms for test function to be altered
+    #                    by SUPG
+    suboptions = None
+
 
 class MixedFSOptions(WrapperOptions):
     """Specifies options for a mixed finite element formulation
     where different suboptions are applied to different
     prognostic variables."""
 
     name = "mixed_options"
-    suboptions = {}
+
+    # Dictionary containing keys field_name and values suboption
+    # field_name (str): name of the field for suboption to be applied to
+    # suboption (:class:`WrapperOptions`): Wrapper options to be applied
+    #                                      to the provided field
+    suboptions = None
 
 
 class SpongeLayerParameters(Configuration):

gusto/time_discretisation/imex_runge_kutta.py

@@ -248,7 +248,6 @@ def apply(self, x_out, x_in):
             if self.limiter is not None:
                 self.limiter.apply(self.x_out)
             self.xs[stage].assign(self.x_out)
-
         self.final_solver.solve()
 
         # Apply limiter

gusto/time_discretisation/sdc.py

@@ -163,9 +163,9 @@ def __init__(self, base_scheme, domain, M, maxk, quad_type, node_type, qdelta_im
 
         # Get Q_delta matrices
         self.Qdelta_imp = genQDeltaCoeffs(qdelta_imp, form=formulation,
-                                          nodes=self.nodes, Q=self.Q)
+                                          nodes=self.nodes, Q=self.Q, nNodes=M, nodeType=node_type, quadType=quad_type)
         self.Qdelta_exp = genQDeltaCoeffs(qdelta_exp, form=formulation,
-                                          nodes=self.nodes, Q=self.Q)
+                                          nodes=self.nodes, Q=self.Q, nNodes=M, nodeType=node_type, quadType=quad_type)
 
         # Set default linear and nonlinear solver options if none passed in
         if linear_solver_parameters is None:

gusto/time_discretisation/time_discretisation.py

@@ -91,7 +91,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
                         self.wrapper.subwrappers.update({field: RecoveryWrapper(self, suboption)})
                     elif suboption.name == "supg":
                         raise RuntimeError(
-                            'Time discretisation: suboption SUPG is currently not implemented within MixedOptions')
+                            'Time discretisation: suboption SUPG is not implemented within MixedOptions')
                     else:
                         raise RuntimeError(
                             f'Time discretisation: suboption wrapper {suboption.name} not implemented')
@@ -101,6 +101,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None,
             elif self.wrapper_name == "recovered":
                 self.wrapper = RecoveryWrapper(self, options)
             elif self.wrapper_name == "supg":
+                self.suboptions = options.suboptions
                 self.wrapper = SUPGWrapper(self, options)
             else:
                 raise RuntimeError(
@@ -255,31 +256,47 @@ def setup(self, equation, apply_bcs=True, *active_labels):
 
             else:
                 if self.wrapper_name == "supg":
-                    self.wrapper.setup()
+                    if self.suboptions is not None:
+                        for field_name, term_labels in self.suboptions.items():
+                            self.wrapper.setup(field_name)
+                            new_test = self.wrapper.test
+                            if term_labels is not None:
+                                for term_label in term_labels:
+                                    self.residual = self.residual.label_map(
+                                        lambda t: t.get(prognostic) == field_name and t.has_label(term_label),
+                                        map_if_true=replace_test_function(new_test, old_idx=self.wrapper.idx))
+                            else:
+                                self.residual = self.residual.label_map(
+                                    lambda t: t.get(prognostic) == field_name,
+                                    map_if_true=replace_test_function(new_test, old_idx=self.wrapper.idx))
+                            self.residual = self.wrapper.label_terms(self.residual)
+                    else:
+                        self.wrapper.setup(self.field_name)
+                        new_test = self.wrapper.test
+                        self.residual = self.residual.label_map(
+                            all_terms,
+                            map_if_true=replace_test_function(new_test))
+                        self.residual = self.wrapper.label_terms(self.residual)
                 else:
                     self.wrapper.setup(self.fs, wrapper_bcs)
-                self.fs = self.wrapper.function_space
+                    self.fs = self.wrapper.function_space
+                    new_test = TestFunction(self.wrapper.test_space)
+                    # Replace the original test function with the one from the wrapper
+                    self.residual = self.residual.label_map(
+                        all_terms,
+                        map_if_true=replace_test_function(new_test))
+
+                    self.residual = self.wrapper.label_terms(self.residual)
                 if self.solver_parameters is None:
                     self.solver_parameters = self.wrapper.solver_parameters
-                new_test = TestFunction(self.wrapper.test_space)
-                # SUPG has a special wrapper
-                if self.wrapper_name == "supg":
-                    new_test = self.wrapper.test
-
-                # Replace the original test function with the one from the wrapper
-                self.residual = self.residual.label_map(
-                    all_terms,
-                    map_if_true=replace_test_function(new_test))
-
-                self.residual = self.wrapper.label_terms(self.residual)
 
         # -------------------------------------------------------------------- #
         # Make boundary conditions
         # -------------------------------------------------------------------- #
 
         if not apply_bcs:
             self.bcs = None
-        elif self.wrapper is not None:
+        elif self.wrapper is not None and self.wrapper_name != "supg":
             if self.wrapper_name == 'mixed_options':
                 # Define new Dirichlet BCs on the wrapper-modified
                 # mixed function space.

gusto/time_discretisation/wrappers.py

@@ -334,16 +334,22 @@ class SUPGWrapper(Wrapper):
     test function space that is used to solve the problem.
     """
 
-    def setup(self):
+    def setup(self, field_name):
         """Sets up function spaces and fields needed for this wrapper."""
 
         assert isinstance(self.options, SUPGOptions), \
             'SUPG wrapper can only be used with SUPG Options'
 
         domain = self.time_discretisation.domain
+        if self.options.suboptions is not None:
+            self.idx = self.time_discretisation.equation.field_names.index(field_name)
+            self.test_space = self.time_discretisation.equation.spaces[self.idx]
+        else:
+            self.idx = None
+            self.test_space = self.time_discretisation.fs
         self.function_space = self.time_discretisation.fs
-        self.test_space = self.function_space
         self.x_out = Function(self.function_space)
+        self.field_name = field_name
 
         # -------------------------------------------------------------------- #
         # Work out SUPG parameter
@@ -360,10 +366,10 @@ def setup(self):
             default_vals = [self.options.default*self.time_discretisation.dt]*dim
             # check for directions is which the space is discontinuous
             # so that we don't apply supg in that direction
-            if is_cg(self.function_space):
+            if is_cg(self.test_space):
                 vals = default_vals
             else:
-                space = self.function_space.ufl_element().sobolev_space
+                space = self.test_space.ufl_element().sobolev_space
                 if space.name in ["HDiv", "DirectionalH"]:
                     vals = [default_vals[i] if space[i].name == "H1"
                             else 0. for i in range(dim)]
@@ -381,8 +387,11 @@ def setup(self):
         # -------------------------------------------------------------------- #
         # Set up test function
         # -------------------------------------------------------------------- #
+        if self.options.suboptions is not None:
+            test = self.time_discretisation.equation.tests[self.idx]
+        else:
+            test = TestFunction(self.test_space)
 
-        test = TestFunction(self.test_space)
         uadv = Function(domain.spaces('HDiv'))
         self.test = test + dot(dot(uadv, self.tau), grad(test))
         self.transporting_velocity = uadv

integration-tests/transport/test_supg_transport.py

@@ -13,6 +13,59 @@ def run(timestepper, tmax, f_end):
     return norm(timestepper.fields("f") - f_end) / norm(f_end)
 
 
+def run_coupled(timestepper, tmax, f_end):
+    timestepper.run(0, tmax)
+    norm1 = norm(timestepper.fields("f1") - f_end) / norm(f_end)
+    norm2 = norm(timestepper.fields("f2") - f_end) / norm(f_end)
+    return norm1, norm2
+
+
+@pytest.mark.parametrize("scheme", ["ssprk", "implicit_midpoint"])
+def test_supg_transport_mixed_scalar(tmpdir, scheme, tracer_setup):
+    setup = tracer_setup(tmpdir, geometry="slice")
+    domain = setup.domain
+
+    ibp = IntegrateByParts.TWICE
+
+    opts = SUPGOptions(ibp=ibp)
+
+    tracer1 = ActiveTracer(name='f1', space="theta",
+                           variable_type=TracerVariableType.mixing_ratio,
+                           transport_eqn=TransportEquationType.advective)
+    tracer2 = ActiveTracer(name='f2', space="theta",
+                           variable_type=TracerVariableType.mixing_ratio,
+                           transport_eqn=TransportEquationType.conservative)
+    tracers = [tracer1, tracer2]
+    Vu = domain.spaces("HDiv")
+    eqn = CoupledTransportEquation(domain, active_tracers=tracers, Vu=Vu)
+    suboptions = {}
+    suboptions.update({'f1': [time_derivative, transport]})
+    suboptions.update({'f2': None})
+    opts = SUPGOptions(suboptions=suboptions)
+    transport_method = [DGUpwind(eqn, "f1", ibp=ibp), DGUpwind(eqn, "f2", ibp=ibp)]
+
+    if scheme == "ssprk":
+        transport_scheme = SSPRK3(domain, options=opts)
+    elif scheme == "implicit_midpoint":
+        transport_scheme = TrapeziumRule(domain, options=opts)
+
+    time_varying_velocity = False
+    timestepper = PrescribedTransport(
+        eqn, transport_scheme, setup.io, time_varying_velocity, transport_method
+    )
+
+    # Initial conditions
+    timestepper.fields("f1").interpolate(setup.f_init)
+    timestepper.fields("f2").interpolate(setup.f_init)
+    timestepper.fields("u").project(setup.uexpr)
+
+    error1, error2 = run_coupled(timestepper, setup.tmax, setup.f_end)
+    assert error1 < setup.tol, \
+        'The transport error for f1 is greater than the permitted tolerance'
+    assert error2 < setup.tol, \
+        'The transport error for f2 is greater than the permitted tolerance'
+
+
 @pytest.mark.parametrize("equation_form", ["advective", "continuity"])
 @pytest.mark.parametrize("scheme", ["ssprk", "implicit_midpoint"])
 @pytest.mark.parametrize("space", ["CG", "theta"])
@@ -29,28 +82,26 @@ def test_supg_transport_scalar(tmpdir, equation_form, scheme, space,
         V = domain.spaces("theta")
         ibp = IntegrateByParts.TWICE
 
-    opts = SUPGOptions(ibp=ibp)
-
     if equation_form == "advective":
         eqn = AdvectionEquation(domain, V, "f")
     else:
         eqn = ContinuityEquation(domain, V, "f")
 
+    opts = SUPGOptions(ibp=ibp)
+    transport_method = DGUpwind(eqn, "f", ibp=ibp)
+
     if scheme == "ssprk":
         transport_scheme = SSPRK3(domain, options=opts)
     elif scheme == "implicit_midpoint":
         transport_scheme = TrapeziumRule(domain, options=opts)
 
-    transport_method = DGUpwind(eqn, "f", ibp=ibp)
     time_varying_velocity = False
     timestepper = PrescribedTransport(
         eqn, transport_scheme, setup.io, time_varying_velocity, transport_method
     )
 
-    # Initial conditions
     timestepper.fields("f").interpolate(setup.f_init)
     timestepper.fields("u").project(setup.uexpr)
-
     error = run(timestepper, setup.tmax, setup.f_end)
     assert error < setup.tol, \
         'The transport error is greater than the permitted tolerance'
@@ -81,12 +132,14 @@ def test_supg_transport_vector(tmpdir, equation_form, scheme, space,
     else:
         eqn = ContinuityEquation(domain, V, "f")
 
+    opts = SUPGOptions(ibp=ibp)
+    transport_method = DGUpwind(eqn, "f", ibp=ibp)
+
     if scheme == "ssprk":
         transport_scheme = SSPRK3(domain, options=opts)
     elif scheme == "implicit_midpoint":
         transport_scheme = TrapeziumRule(domain, options=opts)
 
-    transport_method = DGUpwind(eqn, "f", ibp=ibp)
     time_varying_velocity = False
     timestepper = PrescribedTransport(
         eqn, transport_scheme, setup.io, time_varying_velocity, transport_method