diff --git a/gusto/diagnostics/shallow_water_diagnostics.py b/gusto/diagnostics/shallow_water_diagnostics.py index 024e58ecf..e873f53ca 100644 --- a/gusto/diagnostics/shallow_water_diagnostics.py +++ b/gusto/diagnostics/shallow_water_diagnostics.py @@ -11,7 +11,7 @@ __all__ = ["ShallowWaterKineticEnergy", "ShallowWaterPotentialEnergy", "ShallowWaterPotentialEnstrophy", "PotentialVorticity", "RelativeVorticity", "AbsoluteVorticity", "PartitionedVapour", - "PartitionedCloud"] + "PartitionedCloud", "ShallowWaterAvailablePotentialEnergy"] class ShallowWaterKineticEnergy(Energy): @@ -382,3 +382,37 @@ def compute(self): """Performs the computation of the diagnostic field.""" self.qsat_func.assign(assemble(self.qsat_interpolate)) super().compute() + + +class ShallowWaterAvailablePotentialEnergy(Energy): + """Diagnostic shallow-water available potential energy density.""" + name = "ShallowWaterAvailablePotentialEnergy" + + def __init__(self, parameters, space=None, method='interpolate'): + """ + Args: + parameters (:class:`ShallowWaterParameters`): 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 'interpolate'. + """ + self.parameters = parameters + super().__init__(space=space, method=method, required_fields=("D")) + + 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. + """ + g = self.parameters.g + H = self.parameters.H + D = state_fields("D") + self.expr = 0.5*g*(D-H)**2 + super().setup(domain, state_fields) diff --git a/unit-tests/diagnostic_tests/test_sw-available-pe.py b/unit-tests/diagnostic_tests/test_sw-available-pe.py new file mode 100644 index 000000000..79e4b5955 --- /dev/null +++ b/unit-tests/diagnostic_tests/test_sw-available-pe.py @@ -0,0 +1,33 @@ + +from gusto.diagnostics import ShallowWaterAvailablePotentialEnergy +from gusto.core.fields import StateFields, PrescribedFields, TimeLevelFields +from gusto import (Domain, ShallowWaterParameters, ShallowWaterEquations) +from firedrake import PeriodicSquareMesh +import numpy as np + + +def test_swavlbpe(): + + nx = 10 + Lx = 1 + H = 0 + + mesh = PeriodicSquareMesh(nx=nx, ny=nx, L=Lx, quadrilateral=True) + + domain = Domain(mesh, 0.1, 'RTCF', 1) + params = ShallowWaterParameters(mesh, H=H) + eqn = ShallowWaterEquations(domain, params) + prog_fields = TimeLevelFields(eqn) + prescribed_fields = PrescribedFields() + state_fields = StateFields(prog_fields, prescribed_fields) + + diagnostic = ShallowWaterAvailablePotentialEnergy(params) + diagnostic.setup(domain, state_fields) + diagnostic.compute() + + print(diagnostic.field.dat.data) + + g = params.g + + assert np.allclose(diagnostic.field.dat.data, 0.5*g*H**2, atol=0.0), \ + 'The sw available potential energy diagnostic does not seem to be correct'