Merge branch 'main' into non_inc_imp_rk

README.md

@@ -16,7 +16,7 @@ These compatible methods underpin the Met Office's next-generation model, [LFRic
 
 The best way to install Gusto is as an additional package when installing [Firedrake](http://firedrakeproject.org). Usually, for a Mac with Homebrew or an Ubuntu installation this is done by downloading the Firedrake install script and executing it:
 ```
-curl -0 https://raw.githubusercontent/com/firedrakeproject/firedrake/master/scripts/firedrake-install
+curl -0 https://raw.githubusercontent.com/firedrakeproject/firedrake/master/scripts/firedrake-install
 python3 firedrake-install --install gusto
 ```
 For an up-to-date installation guide, see the [firedrake installation instructions](http://firedrakeproject.org/download.html). Once installed, Gusto must be run from within the Firedrake virtual environment, which is activated via

examples/compressible_euler/mountain_hydrostatic.py

@@ -18,7 +18,7 @@
     TrapeziumRule, SUPGOptions, ZComponent, Perturbation,
     CompressibleParameters, HydrostaticCompressibleEulerEquations,
     CompressibleSolver, compressible_hydrostatic_balance, HydrostaticImbalance,
-    SpongeLayerParameters, MinKernel, MaxKernel, remove_initial_w, logger
+    SpongeLayerParameters, MinKernel, MaxKernel, logger
 )
 
 mountain_hydrostatic_defaults = {
@@ -243,8 +243,7 @@ def mountain_hydrostatic(
 
     theta0.assign(theta_b)
     rho0.assign(rho_b)
-    u0.project(as_vector([initial_wind, 0.0]))
-    remove_initial_w(u0)
+    u0.project(as_vector([initial_wind, 0.0]), bcs=eqns.bcs['u'])
 
     stepper.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
 

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/equations/prognostic_equations.py

@@ -305,12 +305,13 @@ def add_tracers_to_prognostics(self, domain, active_tracers):
                 name of the active tracer.
         """
 
-        # Check if there are any conservatively transported tracers.
-        # If so, ensure that the reference density is indexed before this tracer.
+        # If there are any conservatively transported tracers, ensure
+        # that the reference density, if it is also an active tracer,
+        # is indexed earlier.
         for i in range(len(active_tracers) - 1):
             tracer = active_tracers[i]
             if tracer.transport_eqn == TransportEquationType.tracer_conservative:
-                ref_density = next(x for x in active_tracers if x.name == tracer.density_name)
+                ref_density = next((x for x in active_tracers if x.name == tracer.density_name), tracer)
                 j = active_tracers.index(ref_density)
                 if j > i:
                     # Swap the indices of the tracer and the reference density

gusto/initialisation/hydrostatic_initialisation.py

@@ -12,9 +12,10 @@
 from gusto.recovery import Recoverer, BoundaryMethod
 
 
-__all__ = ["boussinesq_hydrostatic_balance",
-           "compressible_hydrostatic_balance", "remove_initial_w",
-           "saturated_hydrostatic_balance", "unsaturated_hydrostatic_balance"]
+__all__ = [
+    "boussinesq_hydrostatic_balance", "compressible_hydrostatic_balance",
+    "saturated_hydrostatic_balance", "unsaturated_hydrostatic_balance"
+]
 
 
 def boussinesq_hydrostatic_balance(equation, b0, p0, top=False, params=None):
@@ -219,23 +220,6 @@ def compressible_hydrostatic_balance(equation, theta0, rho0, exner0=None,
         rho0.interpolate(thermodynamics.rho(parameters, theta0, exner))
 
 
-def remove_initial_w(u):
-    """
-    Removes the vertical component of a velocity field.
-
-    Args:
-        u (:class:`Function`): the velocity field to be altered.
-    """
-    Vu = u.function_space()
-    Vv = FunctionSpace(Vu._ufl_domain, Vu.ufl_element()._elements[-1])
-    bc = DirichletBC(Vu[0], 0.0, "bottom")
-    bc.apply(u)
-    uv = Function(Vv).project(u)
-    ustar = Function(u.function_space()).project(uv)
-    uin = Function(u.function_space()).assign(u - ustar)
-    u.assign(uin)
-
-
 def saturated_hydrostatic_balance(equation, state_fields, theta_e, mr_t,
                                   exner0=None, top=False,
                                   exner_boundary=Constant(1.0),

gusto/solvers/preconditioners.py

@@ -8,7 +8,6 @@
 from firedrake.petsc import PETSc
 from gusto.recovery.recovery_kernels import AverageKernel, AverageWeightings
 from pyop2.profiling import timed_region, timed_function
-from pyop2.utils import as_tuple
 from functools import partial
 
 
@@ -144,7 +143,7 @@ def initialize(self, pc):
                 if isinstance(subdom, str):
                     neumann_subdomains |= set([subdom])
                 else:
-                    neumann_subdomains |= set(as_tuple(subdom, int))
+                    neumann_subdomains |= set(subdom)
 
             # separate out the top and bottom bcs
             extruded_neumann_subdomains = neumann_subdomains & {"top", "bottom"}

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(
@@ -132,26 +133,58 @@ def setup(self, equation, apply_bcs=True, *active_labels):
         self.residual = equation.residual
 
         if self.field_name is not None and hasattr(equation, "field_names"):
-            self.idx = equation.field_names.index(self.field_name)
-            self.fs = equation.spaces[self.idx]
-            self.residual = self.residual.label_map(
-                lambda t: t.get(prognostic) == self.field_name,
-                lambda t: Term(
-                    split_form(t.form)[self.idx].form,
-                    t.labels),
-                drop)
+            if isinstance(self.field_name, list):
+                # Multiple fields are being solved for simultaneously.
+                # This enables conservative transport to be implemented with SIQN.
+                # Use the full mixed space for self.fs, with the
+                # field_name, residual, and BCs being set up later.
+                self.fs = equation.function_space
+                self.idx = None
+            else:
+                self.idx = equation.field_names.index(self.field_name)
+                self.fs = equation.spaces[self.idx]
+                self.residual = self.residual.label_map(
+                    lambda t: t.get(prognostic) == self.field_name,
+                    lambda t: Term(
+                        split_form(t.form)[self.idx].form,
+                        t.labels),
+                    drop)
 
         else:
             self.field_name = equation.field_name
             self.fs = equation.function_space
             self.idx = None
 
-        bcs = equation.bcs[self.field_name]
-
         if len(active_labels) > 0:
-            self.residual = self.residual.label_map(
-                lambda t: any(t.has_label(time_derivative, *active_labels)),
-                map_if_false=drop)
+            if isinstance(self.field_name, list):
+                # Multiple fields are being solved for simultaneously.
+                # Keep all time derivative terms:
+                residual = self.residual.label_map(
+                    lambda t: t.has_label(time_derivative),
+                    map_if_false=drop)
+
+                # Only keep active labels for prognostics in the list
+                # of simultaneously transported variables:
+                for subname in self.field_name:
+                    field_residual = self.residual.label_map(
+                        lambda t: t.get(prognostic) == subname,
+                        map_if_false=drop)
+
+                    residual += field_residual.label_map(
+                        lambda t: t.has_label(*active_labels),
+                        map_if_false=drop)
+
+                self.residual = residual
+            else:
+                self.residual = self.residual.label_map(
+                    lambda t: any(t.has_label(time_derivative, *active_labels)),
+                    map_if_false=drop)
+
+        # Set the field name if using simultaneous transport.
+        if isinstance(self.field_name, list):
+            self.field_name = equation.field_name
+
+        bcs = equation.bcs[self.field_name]
 
         self.evaluate_source = []
         self.physics_names = []
@@ -175,7 +208,10 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                     # timestepper should be used instead.
                     if len(field_terms.label_map(lambda t: t.has_label(mass_weighted), map_if_false=drop)) > 0:
                         if len(field_terms.label_map(lambda t: not t.has_label(mass_weighted), map_if_false=drop)) > 0:
-                            raise ValueError(f"Mass-weighted and non-mass-weighted terms are present in a timestepping equation for {field}. As these terms cannot be solved for simultaneously, a split timestepping method should be used instead.")
+                            raise ValueError('Mass-weighted and non-mass-weighted terms are present in a '
+                                             + f'timestepping equation for {field}. As these terms cannot '
+                                             + 'be solved for simultaneously, a split timestepping method '
+                                             + 'should be used instead.')
                         else:
                             # Replace the terms with a mass_weighted label with the
                             # mass_weighted form. It is important that the labels from
@@ -199,10 +235,11 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                 for field, subwrapper in self.wrapper.subwrappers.items():
 
                     if field not in equation.field_names:
-                        raise ValueError(f"The option defined for {field} is for a field that does not exist in the equation set")
+                        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[field_idx], wrapper_bcs)
+                    subwrapper.setup(equation.spaces[field_idx], equation.bcs[field])
 
                     # Update the function space to that needed by the wrapper
                     self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space
@@ -219,34 +256,60 @@ 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:
-            # 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 self.wrapper_name != "supg":
+            if self.wrapper_name == 'mixed_options':
+                # Define new Dirichlet BCs on the wrapper-modified
+                # mixed function space.
+                self.bcs = []
+                for idx, field_name in enumerate(self.equation.field_names):
+                    for bc in equation.bcs[field_name]:
+                        self.bcs.append(DirichletBC(self.fs.sub(idx),
+                                                    bc.function_arg,
+                                                    bc.sub_domain))
+            else:
+                # 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

@@ -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