Compressible vorticity

examples/compressible/skamarock_klemp_hydrostatic.py

@@ -39,11 +39,9 @@
 domain = Domain(mesh, dt, "RTCF", 1)
 
 # Equation
-parameters = CompressibleParameters()
-Omega = as_vector((0., 0., 0.5e-4))
+parameters = CompressibleParameters(Omega=0.5e-4)
 balanced_pg = as_vector((0., -1.0e-4*20, 0.))
-eqns = CompressibleEulerEquations(domain, parameters, Omega=Omega,
-                                  extra_terms=[("u", balanced_pg)])
+eqns = CompressibleEulerEquations(domain, parameters, extra_terms=[("u", balanced_pg)])
 
 # I/O
 dirname = 'skamarock_klemp_hydrostatic'

gusto/configuration.py

@@ -136,6 +136,7 @@ class CompressibleParameters(Configuration):
     w_sat2 = -17.27  # second const. in Teten's formula (no units)
     w_sat3 = 35.86  # third const. in Teten's formula (K)
     w_sat4 = 610.9  # fourth const. in Teten's formula (Pa)
+    Omega = None  # Rotation rate
 
 
 class ShallowWaterParameters(Configuration):

gusto/diagnostics.py

@@ -3,18 +3,18 @@
 from firedrake import assemble, dot, dx, Function, sqrt, \
     TestFunction, TrialFunction, Constant, grad, inner, curl, \
     LinearVariationalProblem, LinearVariationalSolver, FacetNormal, \
-    ds_b, ds_v, ds_t, dS_h, dS_v, ds, dS, div, avg, jump, pi, \
+    ds_b, ds_v, ds_t, dS_h, dS_v, ds, dS, div, cross, avg, jump, pi, \
     TensorFunctionSpace, SpatialCoordinate, as_vector, \
     Projector, Interpolator, FunctionSpace, FiniteElement, \
     TensorProductElement
 from firedrake.assign import Assigner
 from ufl.domain import extract_unique_domain
 
 from abc import ABCMeta, abstractmethod, abstractproperty
+from gusto.equations import CompressibleEulerEquations
 import gusto.thermodynamics as tde
 from gusto.coord_transforms import rotated_lonlatr_vectors
 from gusto.recovery import Recoverer, BoundaryMethod
-from gusto.equations import CompressibleEulerEquations
 from gusto.active_tracers import TracerVariableType, Phases
 from gusto.logging import logger
 from gusto.kernels import MinKernel, MaxKernel
@@ -28,7 +28,8 @@
            "Perturbation", "Theta_e", "InternalEnergy", "PotentialEnergy",
            "ThermodynamicKineticEnergy", "Dewpoint", "Temperature", "Theta_d",
            "RelativeHumidity", "Pressure", "Exner_Vt", "HydrostaticImbalance", "Precipitation",
-           "PotentialVorticity", "RelativeVorticity", "AbsoluteVorticity", "Divergence",
+           "ShallowWaterPotentialVorticity", "ShallowWaterRelativeVorticity", "ShallowWaterAbsoluteVorticity",
+           "Vorticity", "RelativeVorticity", "AbsoluteVorticity", "Divergence",
            "BruntVaisalaFrequencySquared", "TracerDensity"]
 
 
@@ -798,7 +799,7 @@
 
 class ShallowWaterPotentialEnstrophy(DiagnosticField):
     """Diagnostic (dry) compressible kinetic energy density."""
-    def __init__(self, base_field_name="PotentialVorticity", space=None,
+    def __init__(self, base_field_name="ShallowWaterPotentialVorticity", space=None,
                  method='interpolate'):
         """
         Args:
@@ -811,15 +812,17 @@
                 for this diagnostic. Valid options are 'interpolate', 'project',
                 'assign' and 'solve'. Defaults to 'interpolate'.
         """
-        base_enstrophy_names = ["PotentialVorticity", "RelativeVorticity", "AbsoluteVorticity"]
+        base_enstrophy_names = ["ShallowWaterPotentialVorticity",
+                                "ShallowWaterRelativeVorticity",
+                                "ShallowWaterAbsoluteVorticity"]
         if base_field_name not in base_enstrophy_names:
             raise ValueError(
                 f"Don't know how to compute enstrophy with base_field_name={base_field_name};"
                 + f"base_field_name should be one of {base_enstrophy_names}")
         # Work out required fields
-        if base_field_name in ["PotentialVorticity", "AbsoluteVorticity"]:
+        if base_field_name in ["ShallowWaterPotentialVorticity", "ShallowWaterAbsoluteVorticity"]:
             required_fields = (base_field_name, "D")
-        elif base_field_name == "RelativeVorticity":
+        elif base_field_name == "ShallowWaterRelativeVorticity":
             required_fields = (base_field_name, "D", "coriolis")
         else:
             raise NotImplementedError(f'Enstrophy with vorticity {base_field_name} not implemented')
@@ -841,17 +844,17 @@
             domain (:class:`Domain`): the model's domain object.
             state_fields (:class:`StateFields`): the model's field container.
         """
-        if self.base_field_name == "PotentialVorticity":
-            pv = state_fields("PotentialVorticity")
+        if self.base_field_name == "ShallowWaterPotentialVorticity":
+            pv = state_fields("ShallowWaterPotentialVorticity")
             D = state_fields("D")
             self.expr = 0.5*pv**2*D
-        elif self.base_field_name == "RelativeVorticity":
-            zeta = state_fields("RelativeVorticity")
+        elif self.base_field_name == "ShallowWaterRelativeVorticity":
+            zeta = state_fields("ShallowWaterRelativeVorticity")
             D = state_fields("D")
             f = state_fields("coriolis")
             self.expr = 0.5*(zeta + f)**2/D
-        elif self.base_field_name == "AbsoluteVorticity":
-            zeta_abs = state_fields("AbsoluteVorticity")
+        elif self.base_field_name == "ShallowWaterAbsoluteVorticity":
+            zeta_abs = state_fields("ShallowWaterAbsoluteVorticity")
             D = state_fields("D")
             self.expr = 0.5*(zeta_abs)**2/D
         else:
@@ -1549,7 +1552,7 @@
         self.field.assign(self.field + self.flux)
 
 
-class Vorticity(DiagnosticField):
+class ShallowWaterVorticity(DiagnosticField):
     """Base diagnostic field class for shallow-water vorticity variables."""
 
     def setup(self, domain, state_fields, vorticity_type=None):
@@ -1577,9 +1580,9 @@
 
         if self.method != 'solve':
             if vorticity_type == "potential":
-                self.expr = curl(u + f) / D
+                self.expr = (curl(u) + f) / D
             elif vorticity_type == "absolute":
-                self.expr = curl(u + f)
+                self.expr = curl(u) + f
             elif vorticity_type == "relative":
                 self.expr = curl(u)
 
@@ -1604,9 +1607,9 @@
             self.evaluator = LinearVariationalSolver(problem, solver_parameters={"ksp_type": "cg"})
 
 
-class PotentialVorticity(Vorticity):
+class ShallowWaterPotentialVorticity(ShallowWaterVorticity):
     u"""Diagnostic field for shallow-water potential vorticity, q=(∇×(u+f))/D"""
-    name = "PotentialVorticity"
+    name = "ShallowWaterPotentialVorticity"
 
     def __init__(self, space=None, method='solve'):
         """
@@ -1633,9 +1636,9 @@
         super().setup(domain, state_fields, vorticity_type="potential")
 
 
-class AbsoluteVorticity(Vorticity):
+class ShallowWaterAbsoluteVorticity(ShallowWaterVorticity):
     u"""Diagnostic field for absolute vorticity, ζ=∇×(u+f)"""
-    name = "AbsoluteVorticity"
+    name = "ShallowWaterAbsoluteVorticity"
 
     def __init__(self, space=None, method='solve'):
         """
@@ -1661,9 +1664,9 @@
         super().setup(domain, state_fields, vorticity_type="absolute")
 
 
-class RelativeVorticity(Vorticity):
+class ShallowWaterRelativeVorticity(ShallowWaterVorticity):
     u"""Diagnostic field for relative vorticity, ζ=∇×u"""
-    name = "RelativeVorticity"
+    name = "ShallowWaterRelativeVorticity"
 
     def __init__(self, space=None, method='solve'):
         """
@@ -1689,6 +1692,121 @@
         super().setup(domain, state_fields, vorticity_type="relative")
 
 
+class Vorticity(DiagnosticField):
+    u""" Base diagnostic Field for three dimensional Vorticity """
+
+    def setup(self, domain, state_fields, vorticity_type=None):
+        """
+        Sets up the :class:`Function` for the diagnostic field.
+
+        Args:
+            domain (:class:`Domain`): the model's domain object.
+            state_fields (:class:`StateFields`): the model's field container.
+            vorticity_type (str, optional): denotes which type of vorticity to
+                be computed ('relative', 'absolute'). Defaults to
+                None.
+        """
+        # TODO Do we eventually want a potential voriticy?
+        vorticity_types = ['relative', 'absolute']
+        if vorticity_type not in vorticity_types:
+            if vorticity_type == 'potential':
+                raise ValueError('Potential vorticity has not yet been implemented')
+            else:
+                raise ValueError(f'vorticity type must be one of {vorticity_types}, not {vorticity_type}')
+        if domain.mesh.topological_dimension() == 3:
+            space = domain.spaces('HCurl')
+        else:
+            space = domain.spaces('H1')
+
+        u = state_fields('u')
+        if self.method != 'solve':
+            if vorticity_type == 'relative':
+                self.expression = curl(u)
+            elif vorticity_type == 'absolute':
+                Omega = as_vector((0, 0, self.parameters.Omega))
+                self.expression = curl(u) + 2*Omega
+
+        super().setup(domain, state_fields, space=space)
+
+        if self.method == 'solve':
+            if domain.mesh.topological_dimension() == 3:
+                vort = TrialFunction(space)
+                w = TestFunction(space)
+                n = FacetNormal(domain.mesh)
+                a = inner(vort, w) * dx
+                L = inner(u, curl(w)) * dx - jump(cross(w, u), n) * dS_h
+                if vorticity_type == 'absolute':
+                    Omega = as_vector((0, 0, self.parameters.Omega))
+                    L += inner(2*Omega, w) * dx
+            else:
+                vort = TrialFunction(space)
+                gamma = TestFunction(space)
+                n = FacetNormal(domain.mesh)
+                a = inner(vort, gamma) * dx
+                L = ( inner(domain.perp(grad(gamma)), u))*dx - jump(inner(domain.perp(n), u)*gamma)*ds_h
+                # TODO implement absolute version, unsure atm how to get corioilis in vertical slice smartly
+            problem = LinearVariationalProblem(a, L, self.field)
+            self.evaluator = LinearVariationalSolver(problem, solver_parameters={'ksp_type': 'cg'})
+
+class RelativeVorticity(Vorticity):
+    u""" Diagnostic field for compressible relative vorticity  """
+    name = 'RelativeVorticity'
+
+    def __init__(self, space=None, method='solve'):
+        u"""
+        Args:
+            space (:class:`FunctionSpace`, optional): the function space to
+                evaluate the diagnostic field in. Defaults to None, in which
+                case a default space will be chosen for this diagnostic.
+            method (str, optional): a string specifying the method of evaluation
+                for this diagnostic. Valid options are 'interpolate', 'project',
+                'assign' and 'solve'. Defaults to 'solve'.
+        """
+        self.solve_implemented = True
+        super().__init__(space=space, method=method, required_fields=('u'))
+
+    def setup(self, domain, state_fields):
+        u"""
+        Sets up the :class:`Function` for the diagnostic field.
+
+        Args:
+            domain (:class:`Domain`): the model's domain object.
+            state_fields (:class:`StateFields`): the model's field container.
+        """
+        super().setup(domain, state_fields, vorticity_type='relative')
+
+
+class AbsoluteVorticity(Vorticity):
+    u""" Diagnostic field for compressible euler absolute vorticity  """
+    name = 'AbsoluteVorticity'
+
+    def __init__(self, parameters, space=None, method='solve'):
+        u"""
+        Args:
+            parameters (:class:`CompressibleEulerParameters`): the configuration
+                object containing the physical parameters for this equation.
+            space (:class:`FunctionSpace`, optional): the function space to
+                evaluate the diagnostic field in. Defaults to None, in which
+                case a default space will be chosen for this diagnostic.
+            method (str, optional): a string specifying the method of evaluation
+                for this diagnostic. Valid options are 'interpolate', 'project',
+                'assign' and 'solve'. Defaults to 'solve'.
+        """
+        self.solve_implemented = True
+        self.parameters = parameters
+        super().__init__(space=space, method=method, required_fields=('u'))
+
+    def setup(self, domain, state_fields):
+        u"""
+        Sets up the :class:`Function` for the diagnostic field.
+
+        Args:
+            domain (:class:`Domain`): the model's domain object.
+            state_fields (:class:`StateFields`): the model's field container.
+        """
+        super().setup(domain, state_fields, vorticity_type='absolute')
+
+
 class TracerDensity(DiagnosticField):
     """Diagnostic for computing the density of a tracer. This is
     computed as the product of a mixing ratio and dry density"""

gusto/equations.py

@@ -5,7 +5,7 @@
     TestFunction, Function, sin, pi, inner, dx, div, cross,
     FunctionSpace, MixedFunctionSpace, TestFunctions, TrialFunction,
     FacetNormal, jump, avg, dS_v, dS, DirichletBC, conditional,
-    SpatialCoordinate, split, Constant, action
+    SpatialCoordinate, split, Constant, action, as_vector
 )
 from firedrake.fml import (
     Term, all_terms, keep, drop, Label, subject,
@@ -1074,7 +1074,7 @@ class CompressibleEulerEquations(PrognosticEquationSet):
     pressure.
     """
 
-    def __init__(self, domain, parameters, Omega=None, sponge_options=None,
+    def __init__(self, domain, parameters, sponge_options=None,
                  extra_terms=None, space_names=None,
                  linearisation_map='default',
                  u_transport_option="vector_invariant_form",
@@ -1256,10 +1256,12 @@ def __init__(self, domain, parameters, Omega=None, sponge_options=None,
         # -------------------------------------------------------------------- #
         # Extra Terms (Coriolis, Sponge, Diffusion and others)
         # -------------------------------------------------------------------- #
-        if Omega is not None:
-            # TODO: add linearisation
-            residual += coriolis(subject(prognostic(
+        if parameters.Omega is not None:
+            # TODO add linerisation and label for this
+            Omega = as_vector((0, 0, parameters.Omega))
+            coriolis_form = coriolis(subject(prognostic(
                 inner(w, cross(2*Omega, u))*dx, "u"), self.X))
+            residual += coriolis_form
 
         if sponge_options is not None:
             W_DG = FunctionSpace(domain.mesh, "DG", 2)