From ac3df29e55d51cfec52d5a20cd963782f1e846a2 Mon Sep 17 00:00:00 2001 From: 12shughes Date: Thu, 21 Aug 2025 16:46:21 +0100 Subject: [PATCH 1/6] create diagnostic to calculate shallow water avaliable potential energy --- .../diagnostics/shallow_water_diagnostics.py | 36 ++++++++++++++++++- 1 file changed, 35 insertions(+), 1 deletion(-) diff --git a/gusto/diagnostics/shallow_water_diagnostics.py b/gusto/diagnostics/shallow_water_diagnostics.py index 024e58ecf..96d136fe3 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) \ No newline at end of file From 730f0d90264d47ea669ad53abf873c43a9b9c809 Mon Sep 17 00:00:00 2001 From: 12shughes Date: Tue, 9 Sep 2025 15:51:41 +0100 Subject: [PATCH 2/6] fix lint by adding newline --- gusto/diagnostics/shallow_water_diagnostics.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/gusto/diagnostics/shallow_water_diagnostics.py b/gusto/diagnostics/shallow_water_diagnostics.py index 96d136fe3..4f22e58e2 100644 --- a/gusto/diagnostics/shallow_water_diagnostics.py +++ b/gusto/diagnostics/shallow_water_diagnostics.py @@ -415,4 +415,5 @@ def setup(self, domain, state_fields): H = self.parameters.H D = state_fields("D") self.expr = 0.5*g*(D-H)**2 - super().setup(domain, state_fields) \ No newline at end of file + super().setup(domain, state_fields) + From 3021f4b81d61ad10e6b1846f545cf8c1dccc37cc Mon Sep 17 00:00:00 2001 From: 12shughes Date: Tue, 9 Sep 2025 15:55:58 +0100 Subject: [PATCH 3/6] try and fix linter again --- gusto/diagnostics/shallow_water_diagnostics.py | 1 - 1 file changed, 1 deletion(-) diff --git a/gusto/diagnostics/shallow_water_diagnostics.py b/gusto/diagnostics/shallow_water_diagnostics.py index 4f22e58e2..e873f53ca 100644 --- a/gusto/diagnostics/shallow_water_diagnostics.py +++ b/gusto/diagnostics/shallow_water_diagnostics.py @@ -416,4 +416,3 @@ def setup(self, domain, state_fields): D = state_fields("D") self.expr = 0.5*g*(D-H)**2 super().setup(domain, state_fields) - From c87496d52e0d116d45ffa68754c1cb7bf9adbf90 Mon Sep 17 00:00:00 2001 From: 12shughes Date: Thu, 18 Sep 2025 10:30:16 +0100 Subject: [PATCH 4/6] start making unit test for diagnostic --- .../diagnostic_tests/test_sw-available-pe.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 unit-tests/diagnostic_tests/test_sw-available-pe.py 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..3c7c89ec0 --- /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' \ No newline at end of file From 7084a151ae601b6baba8be600a1e8ce5f90175fd Mon Sep 17 00:00:00 2001 From: 12shughes Date: Thu, 18 Sep 2025 10:39:43 +0100 Subject: [PATCH 5/6] correct lint --- unit-tests/diagnostic_tests/test_sw-available-pe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unit-tests/diagnostic_tests/test_sw-available-pe.py b/unit-tests/diagnostic_tests/test_sw-available-pe.py index 3c7c89ec0..bb8a21f19 100644 --- a/unit-tests/diagnostic_tests/test_sw-available-pe.py +++ b/unit-tests/diagnostic_tests/test_sw-available-pe.py @@ -30,4 +30,4 @@ def test_swavlbpe(): 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' \ No newline at end of file + 'The sw available potential energy diagnostic does not seem to be correct' From f59ffdbf9ed72d1e26dee5e246c7f039dc130858 Mon Sep 17 00:00:00 2001 From: 12shughes Date: Thu, 18 Sep 2025 10:44:34 +0100 Subject: [PATCH 6/6] actually correct lint --- unit-tests/diagnostic_tests/test_sw-available-pe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unit-tests/diagnostic_tests/test_sw-available-pe.py b/unit-tests/diagnostic_tests/test_sw-available-pe.py index bb8a21f19..79e4b5955 100644 --- a/unit-tests/diagnostic_tests/test_sw-available-pe.py +++ b/unit-tests/diagnostic_tests/test_sw-available-pe.py @@ -10,7 +10,7 @@ def test_swavlbpe(): nx = 10 Lx = 1 - H=0 + H = 0 mesh = PeriodicSquareMesh(nx=nx, ny=nx, L=Lx, quadrilateral=True)