diff --git a/gusto/diagnostics/balance.py b/gusto/diagnostics/balance.py new file mode 100644 index 000000000..6062ac53d --- /dev/null +++ b/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() \ No newline at end of file diff --git a/gusto/diagnostics/compressible_euler_diagnostics.py b/gusto/diagnostics/compressible_euler_diagnostics.py index b625e6846..9f80ed4c2 100644 --- a/gusto/diagnostics/compressible_euler_diagnostics.py +++ b/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') diff --git a/gusto/diagnostics/shallow_water_diagnostics.py b/gusto/diagnostics/shallow_water_diagnostics.py index dce6fac69..1ce4d58fc 100644 --- a/gusto/diagnostics/shallow_water_diagnostics.py +++ b/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'): """ diff --git a/gusto/equations/compressible_euler_equations.py b/gusto/equations/compressible_euler_equations.py index 73b80b7d6..2b1c8d2d5 100644 --- a/gusto/equations/compressible_euler_equations.py +++ b/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) diff --git a/unit-tests/test_vorticity_diagnostic.py b/unit-tests/test_vorticity_diagnostic.py new file mode 100644 index 000000000..158799c88 --- /dev/null +++ b/unit-tests/test_vorticity_diagnostic.py @@ -0,0 +1,134 @@ + +from gusto.diagnostics import RelativeVorticity, ZonalComponent, MeridionalComponent, RadialComponent +from gusto.core.fields import StateFields, PrescribedFields, TimeLevelFields +from gusto import (Domain, CompressibleParameters, CompressibleEulerEquations, + GeneralCubedSphereMesh, lonlatr_from_xyz, xyz_vector_from_lonlatr) +from firedrake import (PeriodicIntervalMesh, ExtrudedMesh, Function, sin, cos, File, + SpatialCoordinate, pi, as_vector, errornorm, norm) +import pytest + + +@pytest.mark.parametrize("topology", ["slice", "sphere"]) +def test_relative_vorticity(topology): + if topology=='slice': + vertical_slice_test() + + if topology=='sphere': + sphere_test() + + +def sphere_test(): # try at a higher resolution + R = 1 # radius of ball + H = 5 # height of model top + nlayers=8 + c=16 + # Building model and state object + m = GeneralCubedSphereMesh(R, num_cells_per_edge_of_panel=c, degree=2) + mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers, extrusion_type='radial') + domain=Domain(mesh, 0.1, 'RTCF', degree=1) + params = CompressibleParameters(g=0, cp=0) + eqn = CompressibleEulerEquations(domain, params) + prog_field = TimeLevelFields(eqn) + HDiv = domain.spaces('HDiv') + HCurl = domain.spaces('HCurl') + pres_field = PrescribedFields() + pres_field('u', HDiv) + state = StateFields(prog_field, pres_field) + + # Getting spherical co-ordinates + xyz = SpatialCoordinate(mesh) + _, lat, r = lonlatr_from_xyz(xyz[0], xyz[1], xyz[2]) + e_zonal = xyz_vector_from_lonlatr(1, 0, 0, xyz) + e_rad = xyz_vector_from_lonlatr(0, 0, 1, xyz) + + + # Initlising vorticity field + zonal_u = sin(lat) / r + merid_u = 0.0 + radial_u = 0.0 + #u_expr = xyz_vector_from_lonlatr(zonal_u, merid_u, radial_u, xyz) + print('Projecting u') + state.u.project(e_zonal*zonal_u) + + # Analytic relative vorticity + radial_vort = 2*cos(lat) / r**2 + #analytical_vort_expr = xyz_vector_from_lonlatr(0, 0, radial_vort, xyz) + print('Projecting analytical vorticity') + vorticity_analytic = Function(HCurl, name='exact_vorticity').project(e_rad*radial_vort) + + # Setup and compute diagnostic vorticity and zonal components + Vorticity = RelativeVorticity() + Vorticity.setup(domain, state) + print('diagnosing vorticity') + Vorticity.compute() + + Zonal_u = ZonalComponent('u') + Zonal_u.setup(domain, state) + Zonal_u.compute() + + Meridional_u = MeridionalComponent('u') + Meridional_u.setup(domain, state) + Meridional_u.compute() + + Radial_u = RadialComponent('u') + Radial_u.setup(domain, state) + Radial_u.compute() + + + # zonal_diff = Function(HDiv, name='zonal_diff').project(state.u - state.u_zonal) + diff = Function(HCurl, name='difference').project(vorticity_analytic - state.RelativeVorticity) + + + radial_diagnostic_vort = RadialComponent('RelativeVorticity') + radial_diagnostic_vort.setup(domain, state) + radial_diagnostic_vort.compute() + + # Compare analytic vorticity expression to diagnostic + print('calculating error') + error = errornorm(vorticity_analytic, state.RelativeVorticity) / norm(vorticity_analytic) + print(error) + + outfile = File('spherical_vorticity_high.pvd') + outfile.write(state.u, vorticity_analytic, state.RelativeVorticity, diff, state.u_zonal, state.u_meridional, state.u_radial, state.RelativeVorticity_radial) + # We dont expect it to be zero as the discrete vorticity is not equal to analytic and dependent on resolution + assert error < 1e-6, \ + 'Relative Vorticity not in error tolerence' + +def vertical_slice_test(): + L = 10 + H = 10 + ncol = 100 + nlayers = 100 + + m = PeriodicIntervalMesh(ncol, L) + mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) + _, z = SpatialCoordinate(mesh) + + domain = Domain(mesh, 0.1, 'CG', 1) + params = CompressibleParameters() + eqn = CompressibleEulerEquations(domain, params) + prog_field = TimeLevelFields(eqn) + + H1 = domain.spaces('H1') + HDiv = domain.spaces('HDiv') + + u_expr = 3 * sin(2*pi*z/H) + vort_exact_expr = -6*pi/H * cos(2*pi*z/H) + vorticity_analytic = Function(H1, name='analytic_vort').interpolate(vort_exact_expr) + + # Setting up test field for the diagnostic to use + prescribed_fields = PrescribedFields() + prescribed_fields('u', HDiv) + state = StateFields(prog_field, prescribed_fields) + state.u.project(as_vector([u_expr, 0])) + + Vorticity = RelativeVorticity() + Vorticity.setup(domain, state) + Vorticity.compute() + # Compare analytic vorticity expression to diagnostic + error = errornorm(vorticity_analytic, state.RelativeVorticity) / norm(vorticity_analytic) + # We dont expect it to be zero as the discrete vorticity is not equal to analytic and dependent on resolution + assert error < 1e-6, \ + 'Relative Vorticity not in error tolerence' + +sphere_test() \ No newline at end of file