Merge branch 'main' into wrapper_of_options

docs/Gemfile.lock

@@ -210,15 +210,15 @@ GEM
       jekyll-feed (~> 0.9)
       jekyll-seo-tag (~> 2.1)
     minitest (5.19.0)
-    nokogiri (1.14.3-x86_64-linux)
+    nokogiri (1.16.2-x86_64-linux)
       racc (~> 1.4)
     octokit (4.25.1)
       faraday (>= 1, < 3)
       sawyer (~> 0.9)
     pathutil (0.16.2)
       forwardable-extended (~> 2.6)
     public_suffix (4.0.7)
-    racc (1.6.2)
+    racc (1.7.3)
     rb-fsevent (0.11.2)
     rb-inotify (0.10.1)
       ffi (~> 1.0)

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

@@ -32,11 +32,10 @@
 NP = -20*epsilon
 mu1 = 0.05
 mu2 = 0.98
-L_v = 10
 q0 = 135  # chosen to give an initial max vapour of approx 0.02
-beta2 = 1
-qprecip = 10e-4
-gamma_r = 10e-3
+beta2 = 10
+qprecip = 1e-4
+gamma_r = 1e-3
 # topography parameters
 R0 = pi/9.
 R0sq = R0**2
@@ -94,19 +93,21 @@ def sat_func(x_in):
 def gamma_v(x_in):
     h = x_in.split()[1]
     b = x_in.split()[2]
-    return (1 + L_v*(20*q0/(g*h + g*tpexpr) * exp(20*(1 - b/g))))**(-1)
+    return (1 + beta2*(20*q0/(g*h + g*tpexpr) * exp(20*(1 - b/g))))**(-1)
 
 
-SWSaturationAdjustment(eqns, sat_func, L_v, time_varying_saturation=True,
+SWSaturationAdjustment(eqns, sat_func, time_varying_saturation=True,
                        parameters=parameters, thermal_feedback=True,
                        beta2=beta2, gamma_v=gamma_v,
                        time_varying_gamma_v=True)
 
 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/equations.py

@@ -645,7 +645,7 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None,
             # Default linearisation is time derivatives, pressure gradient and
             # transport term from depth equation. Don't include active tracers
             linearisation_map = lambda t: \
-                t.get(prognostic) in ['u', 'D'] \
+                t.get(prognostic) in ['u', 'D', 'b'] \
                 and (any(t.has_label(time_derivative, pressure_gradient))
                      or (t.get(prognostic) in ['D', 'b'] and t.has_label(transport)))
         super().__init__(field_names, domain, space_names,
@@ -985,7 +985,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."""