Merge branch 'main' into equation_class_extruded_mesh_check

examples/compressible/unsaturated_bubble.py

@@ -67,6 +67,7 @@
 VCG1 = FunctionSpace(mesh, "CG", 1)
 Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element())
 Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1)
+Vt = domain.spaces("theta")
 
 u_opts = RecoveryOptions(embedding_space=Vu_DG1,
                          recovered_space=Vu_CG1,
@@ -92,10 +93,12 @@
 # Physics schemes
 # NB: can't yet use wrapper or limiter options with physics
 rainfall_method = DGUpwind(eqns, 'rain', outflow=True)
+zero_limiter = MixedFSLimiter(eqns, {'water_vapour': ZeroLimiter(Vt),
+                                     'cloud_water': ZeroLimiter(Vt)})
 physics_schemes = [(Fallout(eqns, 'rain', domain, rainfall_method), SSPRK3(domain)),
                    (Coalescence(eqns), ForwardEuler(domain)),
                    (EvaporationOfRain(eqns), ForwardEuler(domain)),
-                   (SaturationAdjustment(eqns), ForwardEuler(domain))]
+                   (SaturationAdjustment(eqns), ForwardEuler(domain, limiter=zero_limiter))]
 
 # Time stepper
 stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields,
@@ -116,7 +119,6 @@
 
 # spaces
 Vu = domain.spaces("HDiv")
-Vt = domain.spaces("theta")
 Vr = domain.spaces("DG")
 x, z = SpatialCoordinate(mesh)
 quadrature_degree = (4, 4)

examples/shallow_water/moist_convective_williamson2.py

@@ -0,0 +1,163 @@
+"""
+A moist convective version of the Williamson 2 shallow water test (steady state
+geostrophically-balanced flow). The saturation function depends on height,
+with a constant background buoyancy/temperature field.
+Vapour is initialised very close to saturation and small overshoots will
+generate clouds.
+"""
+from gusto import *
+from firedrake import (IcosahedralSphereMesh, SpatialCoordinate, sin, cos, exp)
+import sys
+
+# ----------------------------------------------------------------- #
+# Test case parameters
+# ----------------------------------------------------------------- #
+
+dt = 120
+
+if '--running-tests' in sys.argv:
+    tmax = dt
+    dumpfreq = 1
+else:
+    day = 24*60*60
+    tmax = 5*day
+    ndumps = 5
+    dumpfreq = int(tmax / (ndumps*dt))
+
+R = 6371220.
+u_max = 20
+phi_0 = 3e4
+epsilon = 1/300
+theta_0 = epsilon*phi_0**2
+g = 9.80616
+H = phi_0/g
+xi = 0
+q0 = 200
+beta1 = 110
+alpha = 16
+gamma_v = 0.98
+qprecip = 1e-4
+gamma_r = 1e-3
+
+# ----------------------------------------------------------------- #
+# Set up model objects
+# ----------------------------------------------------------------- #
+
+# Domain
+mesh = IcosahedralSphereMesh(radius=R, refinement_level=3, degree=2)
+degree = 1
+domain = Domain(mesh, dt, 'BDM', degree)
+x = SpatialCoordinate(mesh)
+
+# Equations
+parameters = ShallowWaterParameters(H=H, g=g)
+Omega = parameters.Omega
+fexpr = 2*Omega*x[2]/R
+
+tracers = [WaterVapour(space='DG'), CloudWater(space='DG'), Rain(space='DG')]
+
+eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr,
+                             u_transport_option='vector_advection_form',
+                             active_tracers=tracers)
+
+# IO
+dirname = "moist_convective_williamson2"
+output = OutputParameters(dirname=dirname,
+                          dumpfreq=dumpfreq,
+                          dumplist_latlon=['D', 'D_error'],
+                          dump_nc=True,
+                          dump_vtus=True)
+
+diagnostic_fields = [CourantNumber(), RelativeVorticity(),
+                     PotentialVorticity(),
+                     ShallowWaterKineticEnergy(),
+                     ShallowWaterPotentialEnergy(parameters),
+                     ShallowWaterPotentialEnstrophy(),
+                     SteadyStateError('u'), SteadyStateError('D'),
+                     SteadyStateError('water_vapour'),
+                     SteadyStateError('cloud_water')]
+
+io = IO(domain, output, diagnostic_fields=diagnostic_fields)
+
+
+# define saturation function
+def sat_func(x_in):
+    h = x_in.split()[1]
+    lamda, phi, _ = lonlatr_from_xyz(x[0], x[1], x[2])
+    numerator = theta_0 + sigma*((cos(phi))**2) * ((w + sigma)*(cos(phi))**2 + 2*(phi_0 - w - sigma))
+    denominator = phi_0**2 + (w + sigma)**2*(sin(phi))**4 - 2*phi_0*(w + sigma)*(sin(phi))**2
+    theta = numerator/denominator
+    return q0/(g*h) * exp(20*(theta))
+
+
+transport_methods = [DGUpwind(eqns, field_name) for field_name in eqns.field_names]
+
+limiter = DG1Limiter(domain.spaces('DG'))
+
+transported_fields = [TrapeziumRule(domain, "u"),
+                      SSPRK3(domain, "D"),
+                      SSPRK3(domain, "water_vapour", limiter=limiter),
+                      SSPRK3(domain, "cloud_water", limiter=limiter),
+                      SSPRK3(domain, "rain", limiter=limiter)
+                      ]
+
+linear_solver = MoistConvectiveSWSolver(eqns)
+
+sat_adj = SWSaturationAdjustment(eqns, sat_func,
+                                 time_varying_saturation=True,
+                                 convective_feedback=True, beta1=beta1,
+                                 gamma_v=gamma_v, time_varying_gamma_v=False,
+                                 parameters=parameters)
+
+inst_rain = InstantRain(eqns, qprecip, vapour_name="cloud_water",
+                        rain_name="rain", gamma_r=gamma_r)
+
+physics_schemes = [(sat_adj, ForwardEuler(domain)),
+                   (inst_rain, ForwardEuler(domain))]
+
+stepper = SemiImplicitQuasiNewton(eqns, io,
+                                  transport_schemes=transported_fields,
+                                  spatial_methods=transport_methods,
+                                  linear_solver=linear_solver,
+                                  physics_schemes=physics_schemes)
+
+# ----------------------------------------------------------------- #
+# Initial conditions
+# ----------------------------------------------------------------- #
+
+u0 = stepper.fields("u")
+D0 = stepper.fields("D")
+v0 = stepper.fields("water_vapour")
+
+lamda, phi, _ = lonlatr_from_xyz(x[0], x[1], x[2])
+
+uexpr = xyz_vector_from_lonlatr(u_max*cos(phi), 0, 0, x)
+g = parameters.g
+w = Omega*R*u_max + (u_max**2)/2
+sigma = 0
+
+Dexpr = H - (1/g)*(w)*((sin(phi))**2)
+D_for_v = H - (1/g)*(w + sigma)*((sin(phi))**2)
+
+# though this set-up has no buoyancy, we use the expression for theta to set up
+# the initial vapour
+numerator = theta_0 + sigma*((cos(phi))**2) * ((w + sigma)*(cos(phi))**2 + 2*(phi_0 - w - sigma))
+denominator = phi_0**2 + (w + sigma)**2*(sin(phi))**4 - 2*phi_0*(w + sigma)*(sin(phi))**2
+theta = numerator/denominator
+
+initial_msat = q0/(g*Dexpr) * exp(20*theta)
+vexpr = (1 - xi) * initial_msat
+
+u0.project(uexpr)
+D0.interpolate(Dexpr)
+v0.interpolate(vexpr)
+
+# Set reference profiles
+Dbar = Function(D0.function_space()).assign(H)
+stepper.set_reference_profiles([('D', Dbar)])
+
+# ----------------------------------------------------------------- #
+# Run
+# ----------------------------------------------------------------- #
+
+stepper.run(t=0, tmax=tmax)

examples/shallow_water/moist_thermal_williamson5.py

@@ -34,8 +34,8 @@
 mu2 = 0.98
 q0 = 135  # chosen to give an initial max vapour of approx 0.02
 beta2 = 10
-qprecip = 10e-4
-gamma_r = 10e-3
+qprecip = 1e-4
+gamma_r = 1e-3
 # topography parameters
 R0 = pi/9.
 R0sq = R0**2
@@ -104,8 +104,10 @@ def gamma_v(x_in):
 InstantRain(eqns, qprecip, vapour_name="cloud_water", rain_name="rain",
             gamma_r=gamma_r)
 
+transport_methods = [DGUpwind(eqns, field_name) for field_name in eqns.field_names]
+
 # Timestepper
-stepper = Timestepper(eqns, RK4(domain), io)
+stepper = Timestepper(eqns, RK4(domain), io, spatial_methods=transport_methods)
 
 # ----------------------------------------------------------------- #
 # Initial conditions

gusto/configuration.py

@@ -7,7 +7,7 @@
 __all__ = [
     "IntegrateByParts", "TransportEquationType", "OutputParameters",
     "CompressibleParameters", "ShallowWaterParameters",
-    "EmbeddedDGOptions", "RecoveryOptions", "SUPGOptions",
+    "EmbeddedDGOptions", "RecoveryOptions", "SUPGOptions", "MixedFSOptions",
     "SpongeLayerParameters", "DiffusionParameters", "BoundaryLayerParameters"
 ]
 
@@ -172,6 +172,15 @@ class SUPGOptions(WrapperOptions):
     ibp = IntegrateByParts.TWICE
 
 
+class MixedFSOptions(WrapperOptions):
+    """Specifies options for a mixed finite element formulation
+    where different suboptions are applied to different
+    prognostic variables."""
+
+    name = "mixed_options"
+    suboptions = {}
+
+
 class SpongeLayerParameters(Configuration):
     """Specifies parameters describing a 'sponge' (damping) layer."""
 

gusto/equations.py

@@ -986,7 +986,8 @@ def __init__(self, domain, parameters, Omega=None, sponge=None,
         # -------------------------------------------------------------------- #
         # Gravitational Term
         # -------------------------------------------------------------------- #
-        gravity_form = subject(prognostic(Term(g*inner(domain.k, w)*dx), 'u'), self.X)
+        gravity_form = name_label(subject(prognostic(Term(g*inner(domain.k, w)*dx),
+                                                     'u'), self.X), "gravity")
 
         residual = (mass_form + adv_form + pressure_gradient_form + gravity_form)
 

gusto/kernels.py

@@ -75,3 +75,40 @@ def apply(self, field_hat, field_DG1, field_old):
                   "field_DG1": (field_DG1, READ),
                   "field_old": (field_old, READ)},
                  is_loopy_kernel=True)
+
+
+class ClipZero():
+    """Clips any negative field values to be zero."""
+
+    def __init__(self, V):
+        """
+        Args:
+            V (:class:`FunctionSpace`): The space of the field to be clipped.
+        """
+        shapes = {'nDOFs': V.finat_element.space_dimension()}
+        domain = "{{[i]: 0 <= i < {nDOFs}}}".format(**shapes)
+
+        instrs = ("""
+                  for i
+                      if field_in[i] < 0.0
+                          field[i] = 0.0
+                      else
+                          field[i] = field_in[i]
+                      end
+                  end
+                  """)
+
+        self._kernel = (domain, instrs)
+
+    def apply(self, field, field_in):
+        """
+        Performs the par loop.
+
+        Args:
+            field (:class:`Function`): The field to be written to.
+            field_in (:class:`Function`): The field to be clipped.
+        """
+        par_loop(self._kernel, dx,
+                 {"field": (field, WRITE),
+                  "field_in": (field_in, READ)},
+                 is_loopy_kernel=True)

gusto/limiters.py

@@ -8,11 +8,12 @@
 from firedrake import (BrokenElement, Function, FunctionSpace, interval,
                        FiniteElement, TensorProductElement)
 from firedrake.slope_limiter.vertex_based_limiter import VertexBasedLimiter
-from gusto.kernels import LimitMidpoints
+from gusto.kernels import LimitMidpoints, ClipZero
 
 import numpy as np
 
-__all__ = ["DG1Limiter", "ThetaLimiter", "NoLimiter", "MixedFSLimiter"]
+__all__ = ["DG1Limiter", "ThetaLimiter", "NoLimiter", "ZeroLimiter",
+           "MixedFSLimiter"]
 
 
 class DG1Limiter(object):
@@ -163,6 +164,52 @@ def apply(self, field):
         field.interpolate(self.field_hat)
 
 
+class ZeroLimiter(object):
+    """
+    A simple limiter to enforce non-negativity of a field pointwise.
+
+    Negative values are simply clipped to be zero. There is also the option to
+    project the field to another function space to enforce non-negativity there.
+    """
+
+    def __init__(self, space, clipping_space=None):
+        """
+        Args:
+            space (:class:`FunctionSpace`): the space of the incoming field to
+                clip.
+            clipping_space (:class:`FunctionSpace`, optional): the space in
+                which to clip the field. If not specified, the space of the
+                input field is used.
+        """
+
+        self.space = space
+        if clipping_space is not None:
+            self.clipping_space = clipping_space
+            self.map_to_clip = True
+            self.field_to_clip = Function(self.clipping_space)
+        else:
+            self.clipping_space = space
+            self.map_to_clip = False
+
+        self._kernel = ClipZero(self.clipping_space)
+
+    def apply(self, field):
+        """
+        The application of the limiter to the field.
+
+        Args:
+            field (:class:`Function`): the field to apply the limiter to.
+         """
+
+        # Obtain field in clipping space
+        if self.map_to_clip:
+            self.field_to_clip.interpolate(field)
+            self._kernel.apply(self.field_to_clip, self.field_to_clip)
+            field.interpolate(self.field_to_clip)
+        else:
+            self._kernel.apply(field, field)
+
+
 class NoLimiter(object):
     """A blank limiter that does nothing."""