PR #491: from firedrakeproject/linear_boussinesq

Implement the Linear Boussinesq equations

examples/boussinesq/skamarock_klemp_linear.py

@@ -0,0 +1,98 @@
+"""
+The gravity wave test case of Skamarock and Klemp (1994), solved using the
+incompressible Boussinesq equations.
+
+Buoyancy is transported using SUPG.
+"""
+
+from gusto import *
+from firedrake import (PeriodicIntervalMesh, ExtrudedMesh,
+                       sin, SpatialCoordinate, Function, pi)
+import sys
+
+# ---------------------------------------------------------------------------- #
+# Test case parameters
+# ---------------------------------------------------------------------------- #
+
+dt = 0.5
+L = 3.0e5  # Domain length
+H = 1.0e4  # Height position of the model top
+
+if '--running-tests' in sys.argv:
+    tmax = dt
+    dumpfreq = 1
+    columns = 30  # number of columns
+    nlayers = 5  # horizontal layers
+
+else:
+    tmax = 3600.
+    dumpfreq = int(tmax / (2*dt))
+    columns = 300  # number of columns
+    nlayers = 10  # horizontal layers
+
+# ---------------------------------------------------------------------------- #
+# Set up model objects
+# ---------------------------------------------------------------------------- #
+
+# Domain
+m = PeriodicIntervalMesh(columns, L)
+mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers)
+domain = Domain(mesh, dt, 'CG', 1)
+
+# Equation
+parameters = BoussinesqParameters(cs=300)
+eqns = LinearBoussinesqEquations(domain, parameters)
+
+# I/O
+output = OutputParameters(dirname='skamarock_klemp_linear')
+# list of diagnostic fields, each defined in a class in diagnostics.py
+diagnostic_fields = [CourantNumber(), Divergence(), Perturbation('b')]
+io = IO(domain, output, diagnostic_fields=diagnostic_fields)
+
+# Transport schemes
+b_opts = SUPGOptions()
+transport_methods = [DGUpwind(eqns, "p"),
+                     DGUpwind(eqns, "b", ibp=b_opts.ibp)]
+
+
+# Time stepper
+stepper = Timestepper(eqns, RK4(domain), io, spatial_methods=transport_methods)
+
+# ---------------------------------------------------------------------------- #
+# Initial conditions
+# ---------------------------------------------------------------------------- #
+
+b0 = stepper.fields("b")
+p0 = stepper.fields("p")
+
+# spaces
+Vb = b0.function_space()
+Vp = p0.function_space()
+
+x, z = SpatialCoordinate(mesh)
+
+# first setup the background buoyancy profile
+# z.grad(bref) = N**2
+N = parameters.N
+bref = z*(N**2)
+# interpolate the expression to the function
+b_b = Function(Vb).interpolate(bref)
+
+# setup constants
+a = 5.0e3
+deltab = 1.0e-2
+b_pert = deltab*sin(pi*z/H)/(1 + (x - L/2)**2/a**2)
+# interpolate the expression to the function
+b0.interpolate(b_b + b_pert)
+
+p_b = Function(Vp)
+boussinesq_hydrostatic_balance(eqns, b_b, p_b)
+p0.assign(p_b)
+
+# set the background buoyancy
+stepper.set_reference_profiles([('p', p_b), ('b', b_b)])
+
+# ---------------------------------------------------------------------------- #
+# Run
+# ---------------------------------------------------------------------------- #
+stepper.run(t=0, tmax=tmax)

gusto/equations.py

@@ -1326,7 +1326,7 @@ def __init__(self, domain, parameters,
 
         w, phi, gamma = self.tests[0:3]
         u, p, b = split(self.X)
-        u_trial = split(self.trials)[0]
+        u_trial, p_trial, _ = split(self.trials)
         _, p_bar, b_bar = split(self.X_ref)
 
         # -------------------------------------------------------------------- #
@@ -1385,16 +1385,21 @@ def __init__(self, domain, parameters,
         # -------------------------------------------------------------------- #
         if compressible:
             cs = parameters.cs
-            divergence_form = divergence(subject(
-                prognostic(cs**2 * phi * div(u) * dx, 'p'), self.X))
+            linear_div_form = divergence(subject(
+                prognostic(cs**2 * phi * div(u_trial) * dx, 'p'), self.X))
+            divergence_form = divergence(linearisation(
+                subject(prognostic(cs**2 * phi * div(u) * dx, 'p'), self.X),
+                linear_div_form))
         else:
             # This enforces that div(u) = 0
             # The p features here so that the div(u) evaluated in the
             # "forcing" step replaces the whole pressure field, rather than
             # merely providing an increment to it.
-            divergence_form = incompressible(
+            linear_div_form = incompressible(
+                subject(prognostic(phi*(p_trial-div(u_trial))*dx, 'p'), self.X))
+            divergence_form = incompressible(linearisation(
                 subject(prognostic(phi*(p-div(u))*dx, 'p'), self.X),
-            )
+                linear_div_form))
 
         residual = (mass_form + adv_form + divergence_form
                     + pressure_gradient_form + gravity_form)
@@ -1412,3 +1417,90 @@ def __init__(self, domain, parameters,
         # -------------------------------------------------------------------- #
         # Add linearisations to equations
         self.residual = self.generate_linear_terms(residual, self.linearisation_map)
+
+
+class LinearBoussinesqEquations(BoussinesqEquations):
+    """
+    Class for the Boussinesq equations, which evolve the velocity
+    'u', the pressure 'p' and the buoyancy 'b'. Can be compressible or
+    incompressible, depending on the value of the input flag, which defaults
+    to compressible.
+
+    The compressible form of the equations is
+    ∂u/∂t + (u.∇)u + 2Ω×u + ∇p + b*k = 0,                                     \n
+    ∂p/∂t + cs**2 ∇.u = p,                                                    \n
+    ∂b/∂t + (u.∇)b = 0,                                                       \n
+    where k is the vertical unit vector, Ω is the planet's rotation vector
+    and cs is the sound speed.
+
+    For the incompressible form of the equations, the pressure features as
+    a Lagrange multiplier to enforce incompressibility. The equations are     \n
+    ∂u/∂t + (u.∇)u + 2Ω×u + ∇p + b*k = 0,                                     \n
+    ∇.u = p,                                                                  \n
+    ∂b/∂t + (u.∇)b = 0,                                                       \n
+    where k is the vertical unit vector and Ω is the planet's rotation vector.
+    """
+
+    def __init__(self, domain, parameters,
+                 compressible=True,
+                 Omega=None,
+                 space_names=None,
+                 linearisation_map='default',
+                 u_transport_option="vector_invariant_form",
+                 no_normal_flow_bc_ids=None,
+                 active_tracers=None):
+        """
+        Args:
+            domain (:class:`Domain`): the model's domain object, containing the
+                mesh and the compatible function spaces.
+            parameters (:class:`Configuration`, optional): an object containing
+                the model's physical parameters.
+            compressible (bool, optional): flag to indicate whether the
+                equations are compressible. Defaults to True
+            Omega (:class:`ufl.Expr`, optional): an expression for the planet's
+                rotation vector. Defaults to None.
+            space_names (dict, optional): a dictionary of strings for names of
+                the function spaces to use for the spatial discretisation. The
+                keys are the names of the prognostic variables. Defaults to None
+                in which case the spaces are taken from the de Rham complex.
+            linearisation_map (func, optional): a function specifying which
+                terms in the equation set to linearise. If None is specified
+                then no terms are linearised. Defaults to the string 'default',
+                in which case the linearisation includes time derivatives and
+                scalar transport terms.
+            u_transport_option (str, optional): specifies the transport term
+                used for the velocity equation. Supported options are:
+                'vector_invariant_form', 'vector_advection_form' and
+                'circulation_form'.
+                Defaults to 'vector_invariant_form'.
+            no_normal_flow_bc_ids (list, optional): a list of IDs of domain
+                boundaries at which no normal flow will be enforced. Defaults to
+                None.
+            active_tracers (list, optional): a list of `ActiveTracer` objects
+                that encode the metadata for any active tracers to be included
+                in the equations.. Defaults to None.
+
+        Raises:
+            NotImplementedError: active tracers are not implemented.
+        """
+
+        if linearisation_map == 'default':
+            # Default linearisation is time derivatives, pressure gradient,
+            # Coriolis and transport term from depth equation
+            linearisation_map = lambda t: \
+                (any(t.has_label(time_derivative, pressure_gradient, coriolis,
+                                 gravity, divergence, incompressible))
+                 or (t.get(prognostic) in ['p', 'b'] and t.has_label(transport)))
+        super().__init__(domain=domain,
+                         parameters=parameters,
+                         compressible=compressible,
+                         Omega=Omega,
+                         space_names=space_names,
+                         linearisation_map=linearisation_map,
+                         u_transport_option=u_transport_option,
+                         no_normal_flow_bc_ids=no_normal_flow_bc_ids,
+                         active_tracers=active_tracers)
+
+        # Use the underlying routine to do a first linearisation of
+        # the equations
+        self.linearise_equation_set()