slight tidy, everything seems to work

gusto/core/configuration.py

@@ -202,7 +202,7 @@ class SUPGOptions(WrapperOptions):
     tau = None
     default = 1/(sqrt(15))
     ibp = IntegrateByParts.TWICE
-    field_names = ["theta", "water_vapour", "cloud_water"]
+    field_names = []
 
 
 class MixedFSOptions(WrapperOptions):

gusto/time_discretisation/imex_runge_kutta.py

@@ -4,7 +4,7 @@
                        NonlinearVariationalSolver)
 from firedrake.fml import replace_subject, all_terms, drop
 from firedrake.utils import cached_property
-from gusto.core.labels import time_derivative, implicit, explicit, prognostic
+from gusto.core.labels import time_derivative, implicit, explicit
 from gusto.time_discretisation.time_discretisation import (
     TimeDiscretisation, wrapper_apply
 )
@@ -236,42 +236,19 @@ def apply(self, x_out, x_in):
         self.x1.assign(x_in)
         self.x_out.assign(x_in)
         solver_list = self.solvers
-        u_form = self.residual.label_map(lambda t: t.get(prognostic) =='u',
-                                         map_if_false=drop)
-        rho_form = self.residual.label_map(lambda t: t.get(prognostic) =='rho',
-                                         map_if_false=drop)
-        theta_form = self.residual.label_map(lambda t: t.get(prognostic) =='theta',
-                                         map_if_false=drop)
-        water_vapour_form = self.residual.label_map(lambda t: t.get(prognostic) =='water_vapour',
-                                         map_if_false=drop)
-        cloud_water_form = self.residual.label_map(lambda t: t.get(prognostic) =='cloud_water',
-                                         map_if_false=drop)
-        # print("u_form")
-        # print(u_form.form.__str__())
-        # print("rho_form")
-        # print(rho_form.form.__str__())
-        # print("theta_form")
-        # print(theta_form.form.__str__())
-        # print("water_vapour_form")
-        # print(water_vapour_form.form.__str__())
-        # print("cloud_water_form")
-        # print(cloud_water_form.form.__str__())
 
         for stage in range(self.nStages):
             self.solver = solver_list[stage]
             # Set initial solver guess
             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)
-        for evaluate in self.evaluate_source:
-                evaluate(self.xs[self.nStages-1], self.dt)
+
         self.final_solver.solve()
 
         # Apply limiter

gusto/time_discretisation/time_discretisation.py

@@ -101,7 +101,13 @@ 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.wrapper_field_names = options.field_names
+                if len(options.field_names) > 0:
+                    self.wrapper_field_names = options.field_names
+                elif len(options.field_names) == 0 and self.field_name is not None:
+                    self.wrapper_field_names = [self.field_name]
+                else:
+                    raise ValueError("No field names provided for SUPG wrapper applied to a MixedFunctionSpace")
+
                 self.wrapper = SUPGWrapper(self, options)
             else:
                 raise RuntimeError(
@@ -221,7 +227,6 @@ def setup(self, equation, apply_bcs=True, *active_labels):
             else:
                 if self.wrapper_name == "supg":
                     for field_name in self.wrapper_field_names:
-                        print("Wrapper field_name:", field_name)
                         self.wrapper.setup(field_name)
                         new_test = self.wrapper.test
                         self.residual = self.residual.label_map(
@@ -247,10 +252,10 @@ def setup(self, equation, apply_bcs=True, *active_labels):
 
         if not apply_bcs:
             self.bcs = None
-        # elif self.wrapper is not None:
-        #     # Transfer boundary conditions onto test function space
-        #     self.bcs = [DirichletBC(self.fs, bc.function_arg, bc.sub_domain)
-        #                 for bc in bcs]
+        elif self.wrapper is not None and not "supg":
+            # Transfer boundary conditions onto test function space
+            self.bcs = [DirichletBC(self.fs, bc.function_arg, bc.sub_domain)
+                        for bc in bcs]
         else:
             self.bcs = bcs
 

gusto/time_discretisation/wrappers.py

@@ -353,11 +353,40 @@ def setup(self, field_name):
         self.x_out = Function(self.function_space)
         self.field_name = field_name
 
+
         # -------------------------------------------------------------------- #
         # Work out SUPG parameter
         # -------------------------------------------------------------------- #
 
-        k = domain.k
+        # construct tau, if it is not specified
+        dim = domain.mesh.topological_dimension()
+        if self.options.tau is not None:
+            # if tau is provided, check that is has the right size
+            self.tau = self.options.tau
+            assert as_ufl(self.tau).ufl_shape == (dim, dim), "Provided tau has incorrect shape!"
+        else:
+            # create tuple of default values of size dim
+            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.test_space):
+                vals = default_vals
+            else:
+                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)]
+                else:
+                    raise ValueError("I don't know what to do with space %s" % space)
+            self.tau = Constant(tuple([
+                tuple(
+                    [vals[j] if i == j else 0. for i, v in enumerate(vals)]
+                ) for j in range(dim)])
+            )
+            self.solver_parameters = {'ksp_type': 'gmres',
+                                      'pc_type': 'bjacobi',
+                                      'sub_pc_type': 'ilu'}
+
         # -------------------------------------------------------------------- #
         # Set up test function
         # -------------------------------------------------------------------- #
@@ -370,9 +399,7 @@ def setup(self, field_name):
             uadv = Function(domain.spaces('HDiv'))
             test = TestFunction(self.test_space)
 
-        tau = Constant(self.options.default * self.time_discretisation.dt)*dot(domain.k, uadv)
-
-        self.test = test + tau*dot(domain.k, grad(test))
+        self.test = test + dot(dot(uadv, self.tau), grad(test))
         self.transporting_velocity = uadv
 
     def pre_apply(self, x_in):