Compressible vorticity

gusto/diagnostics/balance.py

@@ -0,0 +1,180 @@
+from firedrake import (Constant, Function, dx, TestFunction, TrialFunction, as_vector,
+                       inner, dot, cross, div, jump, LinearVariationalProblem,
+                       LinearVariationalSolver, SpatialCoordinate, tan, FacetNormal,
+                       sqrt, avg, dS_v)
+
+from gusto.core.coord_transforms import lonlatr_from_xyz
+from gusto.diagnostics.diagnostics import DiagnosticField
+from gusto.equations.compressible_euler_equations import CompressibleEulerEquations
+from gusto.equations.thermodynamics import exner_pressure
+
+class GeostrophicImbalance(DiagnosticField):
+    """Geostrophic imbalance diagnostic field."""
+    name = "GeostrophicImbalance"
+
+    def __init__(self, equations, space=None, method='interpolate'):
+        """
+        Args:
+            equations (:class:`PrognosticEquationSet`): the equation set being
+                solved by the model.
+            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 'interpolate'.
+        """
+        # Work out required fields
+        if isinstance(equations, CompressibleEulerEquations):
+            required_fields = ['rho', 'theta', 'u']
+            self.equations = equations
+            self.parameters = equations.parameters
+        else:
+            raise NotImplementedError(f'Geostrophic Imbalance not implemented for {type(equations)}')
+        super().__init__(space=space, method=method, required_fields=required_fields)
+
+    def setup(self, domain, state_fields):
+        """
+        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.
+        """
+        Vcurl = domain.spaces("HCurl")
+        u = state_fields('u')
+        rho = state_fields('rho')
+        theta = state_fields('theta')
+
+        exner = exner_pressure(self.parameters, rho, theta)
+        cp = Constant(self.parameters.cp)
+        n = FacetNormal(domain.mesh)
+        # TODO: Generilise this for cases that aren't solid body rotation case 
+        omega = Constant(7.292e-5)
+        Omega = as_vector((0., 0., omega))
+        k = domain.k
+
+        F = TrialFunction(Vcurl)
+        w = TestFunction(Vcurl)
+
+        imbalance = Function(Vcurl)
+        a = inner(w, F)*dx
+
+        L = (- cp*div(theta*w)*exner*dx + cp*jump(theta*w, n)*avg(exner)*dS_v # exner pressure grad discretisation
+             - inner(w, cross(2*Omega, u))*dx # coriolis
+             + inner(w, dot(k, cross(2*Omega, u) )*k)*dx #vertical part of coriolis
+             + cp*div((theta*k)*dot(k,w))*exner*dx  # removes vertical part of the pressure divergence
+             - cp*jump((theta*k)*dot(k,w), n)*avg(exner)*dS_v) # removes vertical part of pressure jump condition
+
+
+        bcs = self.equations.bcs['u']
+
+        imbalanceproblem = LinearVariationalProblem(a, L, imbalance, bcs=bcs)
+        self.imbalance_solver = LinearVariationalSolver(imbalanceproblem)
+        self.expr = dot(imbalance, domain.k)
+        super().setup(domain, state_fields)
+
+    def compute(self):
+        """Compute and return the diagnostic field from the current state.
+        """
+        self.imbalance_solver.solve()
+        super().compute()
+
+
+class SolidBodyImbalance(DiagnosticField):
+    """Solid Body imbalance diagnostic field."""
+    name = "SolidBodyImbalance"
+
+    def __init__(self, equations, space=None, method='interpolate'):
+        """
+        Args:
+            equations (:class:`PrognosticEquationSet`): the equation set being
+                solved by the model.
+            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 'interpolate'.
+        """
+        # Work out required fields
+        if isinstance(equations, CompressibleEulerEquations):
+            required_fields = ['rho', 'theta', 'u']
+            self.equations = equations
+            self.parameters = equations.parameters
+        else:
+            raise NotImplementedError(f'Geostrophic Imbalance not implemented for {type(equations)}')
+        super().__init__(space=space, method=method, required_fields=required_fields)
+
+    def setup(self, domain, state_fields):
+        """
+        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.
+        """
+        Vu = domain.spaces("HDiv")
+        u = state_fields('u')
+        rho = state_fields('rho')
+        theta = state_fields('theta')
+
+        exner = tde.exner_pressure(self.parameters, rho, theta)
+        cp = Constant(self.parameters.cp)
+        n = FacetNormal(domain.mesh)
+        # TODO: Generilise this for cases that aren't solid body rotation case 
+        omega = Constant(7.292e-5)
+        Omega = as_vector((0., 0., omega))
+
+        # generating the spherical co-ords, and spherical components of velocity
+        x, y, z = SpatialCoordinate(domain.mesh)
+        x_hat = Constant(as_vector([1.0, 0.0, 0.0]))
+        y_hat = Constant(as_vector([0.0, 1.0, 0.0]))
+        z_hat = Constant(as_vector([0.0, 0.0, 1.0]))
+        R = sqrt(x**2 + y**2)  # distance from z axis
+        r = sqrt(x**2 + y**2 + z**2)  # distance from origin
+        lambda_hat = (x * y_hat - y * x_hat) / R
+        lon_dot = inner(u, lambda_hat)
+        phi_hat = (-x*z/R * x_hat - y*z/R * y_hat + R * z_hat) / r
+        lat_dot = inner(u, phi_hat)
+        r_hat = (x * x_hat + y * y_hat + z * z_hat) / r 
+        r_dot = inner(u, r_hat)     
+        mesh = domain.mesh
+        k=domain.k
+        #HACK loweing the quadrature manually
+        dx_low_quadrature = dx(degree=3)
+        lat = latlon_coords(mesh)[0]
+        F = TrialFunction(Vu)
+        w = TestFunction(Vu)
+
+        imbalance = Function(Vu)
+        a = inner(w, F)*dx
+
+        #TODO Segmentaion error somehwere here
+
+        L = (- cp*div((theta)*w)*exner*dx + cp*jump((theta)*w, n)*avg(exner)*dS_v # exner pressure grad discretisation
+             - inner(w, cross(2*Omega, u))*dx # coriolis
+             + inner(w, dot(k, cross(2*Omega, u) )*k)*dx # vertical part of coriolis
+             + cp*div((theta*k)*dot(k,w))*exner*dx  # vertical part of the pressure divergence
+             - cp*jump((theta*k)*dot(k,w), n)*avg(exner)*dS_v # vertical part of pressure jump condition
+            # BUG it is these terms arising from the non-linear that are the problem
+             - (lat_dot * lon_dot * tan(lat) / r)*inner(w, lambda_hat)*dx_low_quadrature 
+             + (lon_dot * r_dot / r)*dot(w, lambda_hat)*dx_low_quadrature # lambda component of non linear term            
+             + (lat_dot**2 * tan(lat) / r)*inner(w, phi_hat)*dx_low_quadrature   # phi component 1
+             + (lat_dot * r_dot / r)*inner(w, phi_hat)*dx_low_quadrature # phi component 1
+            )
+
+
+        bcs = self.equations.bcs['u']
+
+        imbalanceproblem = LinearVariationalProblem(a, L, imbalance, bcs=bcs)
+        self.imbalance_solver = LinearVariationalSolver(imbalanceproblem)
+        self.expr = dot(imbalance, domain.k)
+        super().setup(domain, state_fields)
+
+
+    def compute(self):
+        """Compute and return the diagnostic field from the current state.
+        """
+        self.imbalance_solver.solve()
+        super().compute()

gusto/diagnostics/compressible_euler_diagnostics.py

@@ -1,9 +1,9 @@
 """Common diagnostic fields for the compressible Euler equations."""
 
-from firedrake import (dot, dx, Function, ln, TestFunction, TrialFunction,
-                       Constant, grad, inner, LinearVariationalProblem,
+from firedrake import (dot, dx, Function, ln, TestFunction, TrialFunction, 
+                       Constant, grad, inner, LinearVariationalProblem, dS,
                        LinearVariationalSolver, FacetNormal, ds_b, dS_v, div,
-                       avg, jump, SpatialCoordinate)
+                       avg, jump, SpatialCoordinate, curl, as_vector, cross)
 
 from gusto.diagnostics.diagnostics import (
     DiagnosticField, Energy, IterativeDiagnosticField
@@ -18,7 +18,7 @@
            "PotentialEnergy", "ThermodynamicKineticEnergy",
            "Temperature", "Theta_d", "RelativeHumidity", "Pressure", "Exner_Vt",
            "HydrostaticImbalance", "Precipitation", "BruntVaisalaFrequencySquared",
-           "WetBulbTemperature", "DewpointTemperature"]
+           "WetBulbTemperature", "DewpointTemperature", "RelativeVorticity", "AbsoluteVorticity"]
 
 
 class RichardsonNumber(DiagnosticField):
@@ -801,3 +801,120 @@ def compute(self):
         """Increment the precipitation diagnostic."""
         self.solver.solve()
         self.field.assign(self.field + self.flux)
+
+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)
+                gamma = TestFunction(space)
+                f = Coefficient(space)
+                breakpoint()
+                n = FacetNormal(domain.mesh)
+                a = inner(vort, gamma) * dx
+                L = inner(curl(gamma), u) * dx 
+                if vorticity_type == 'absolute':
+                    Omega = as_vector((0, 0, self.parameters.Omega))
+                    L += inner(2*Omega, gamma) * dx - jump(inner(cross(u, n), gamma), n)*dS
+            else:
+                vort = TrialFunction(space)
+                gamma = TestFunction(space)
+                a = inner(vort, gamma) * dx
+                L = (inner(domain.perp(grad(gamma)), u)) * dx - jump(inner(cross(u, n), gamma, n))*dS
+                # TODO implement absolute version, unsure 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')

gusto/diagnostics/shallow_water_diagnostics.py

@@ -6,8 +6,8 @@
 from gusto.diagnostics.diagnostics import DiagnosticField, Energy
 
 __all__ = ["ShallowWaterKineticEnergy", "ShallowWaterPotentialEnergy",
-           "ShallowWaterPotentialEnstrophy", "PotentialVorticity",
-           "RelativeVorticity", "AbsoluteVorticity"]
+           "ShallowWaterPotentialEnstrophy", "ShallowWaterPotentialVorticity",
+           "ShallowWaterRelativeVorticity", "ShallowWaterAbsoluteVorticity"]
 
 
 class ShallowWaterKineticEnergy(Energy):
@@ -136,7 +136,7 @@ def setup(self, domain, state_fields):
         super().setup(domain, state_fields)
 
 
-class Vorticity(DiagnosticField):
+class ShallowWaterVorticity(DiagnosticField):
     """Base diagnostic field class for shallow-water vorticity variables."""
 
     def setup(self, domain, state_fields, vorticity_type=None):
@@ -191,9 +191,9 @@ def setup(self, domain, state_fields, vorticity_type=None):
             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'):
         """
@@ -220,9 +220,9 @@ def setup(self, domain, state_fields):
         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'):
         """
@@ -248,9 +248,9 @@ def setup(self, domain, state_fields):
         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'):
         """

gusto/equations/compressible_euler_equations.py

@@ -215,12 +215,13 @@ def __init__(self, domain, parameters, sponge_options=None,
         # Extra Terms (Coriolis, Sponge, Diffusion and others)
         # -------------------------------------------------------------------- #
         if parameters.Omega is not None:
-            # TODO: add linearisation
+            # 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))
+                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)
             x = SpatialCoordinate(domain.mesh)