From 555d65c492e25013449ebe911599bf2a17a6405b Mon Sep 17 00:00:00 2001 From: 12shughes Date: Wed, 25 Feb 2026 15:20:39 +0000 Subject: [PATCH 1/8] update files to allow r-theta coords centred off 0,0 --- gusto/core/coord_transforms.py | 14 ++++++++++++-- gusto/core/equation_configuration.py | 9 ++++++--- gusto/equations/shallow_water_equations.py | 2 +- 3 files changed, 19 insertions(+), 6 deletions(-) diff --git a/gusto/core/coord_transforms.py b/gusto/core/coord_transforms.py index 19726fe97..b1bcb4f0d 100644 --- a/gusto/core/coord_transforms.py +++ b/gusto/core/coord_transforms.py @@ -534,13 +534,15 @@ def great_arc_angle(lon1, lat1, lon2, lat2, units='rad'): return arc_length -def xy_from_rtheta(r, theta, angle_units='rad'): +def xy_from_rtheta(r, theta, x0, y0, angle_units='rad'): """ Returns the planar cartsian x, y coordinates from the r, theta coordinates Args: r (:class:`np.ndarray` or :class:`ufl.Expr`): r-coordinate. theta (:class:`np.ndarray` or :class:`ufl.Expr`): theta-coordinate. + x0: central x-coordinate for r-theta coords + y0: central y-coordinate for r-theta cooreds angle_units (str, optional): the units to use for the angle. Valid options are 'rad' (radians) or 'deg' (degrees). Defaults to 'rad'. Returns: @@ -568,16 +570,21 @@ def xy_from_rtheta(r, theta, angle_units='rad'): x = r * cos(theta) y = r * sin(theta) + x += x0 + y += y0 + return x, y -def rtheta_from_xy(x, y, angle_units='rad'): +def rtheta_from_xy(x, y, x0, y0, angle_units='rad'): """ Returns the r, theta coordinates (where theta is measured anticlockwise from horizontal) from the planar Cartesian x, y coordinates. Args: x (:class:`np.ndarray` or :class:`ufl.Expr`): x-coordinate. y (:class:`np.ndarray` or :class:`ufl.Expr`): y-coordinate. + x0: central x-coordinate for r-theta coords + y0: central y-coordinate for r-theta coords angle_units (str, optional): the units to use for the angle. Valid options are 'rad' (radians) or 'deg' (degrees). Defaults to 'rad'. Returns: @@ -600,6 +607,9 @@ def rtheta_from_xy(x, y, angle_units='rad'): if angle_units == 'rad': unit_factor = 1.0 + x -= x0 + y -= y0 + theta = atan2(y, x) r = sqrt(x**2 + y**2) diff --git a/gusto/core/equation_configuration.py b/gusto/core/equation_configuration.py index 1a3b552e2..a482ee37e 100644 --- a/gusto/core/equation_configuration.py +++ b/gusto/core/equation_configuration.py @@ -1,13 +1,13 @@ """Some simple tools for configuring the model.""" -from enum import Enum from firedrake import Function, FunctionSpace, Constant import inspect +from enum import Enum __all__ = [ "BoussinesqParameters", "CompressibleParameters", "ShallowWaterParameters", "CoriolisOptions", - "SpongeLayerParameters", "DiffusionParameters", "BoundaryLayerParameters" + "SpongeLayerParameters", "DiffusionParameters", "BoundaryLayerParameters", ] @@ -111,7 +111,8 @@ class ShallowWaterParameters(EquationParameters): Omega = 7.292e-5 # rotation rate f0 = None # f-, beta- and gamma-plane Coriolis parameter beta = None # beta-plane y-variation parameter - y0 = None # beta-plane y-centre + y0 = None # beta- or gamma-plane y-centre + x0 = None # gamma-plane y-centre R = None # Radius of planet used to compute gamma in gamma-plane approx topog_expr = None # topography expression H = None # mean depth @@ -124,6 +125,8 @@ class ShallowWaterParameters(EquationParameters): # Scaling factor for the saturation function in the equivalent buoyancy # formulation of the thermal shallow water equations q0 = None + # Scaling factor for evaporation term when present + cD = None class SpongeLayerParameters(EquationParameters): diff --git a/gusto/equations/shallow_water_equations.py b/gusto/equations/shallow_water_equations.py index 638207b44..7aaec6d2e 100644 --- a/gusto/equations/shallow_water_equations.py +++ b/gusto/equations/shallow_water_equations.py @@ -215,7 +215,7 @@ def _setup_residual(self, u_transport_option): xyz[1] - self.parameters.y0) elif rotation is CoriolisOptions.gammaplane: x, y = SpatialCoordinate(self.domain.mesh) - r, _ = rtheta_from_xy(x, y) + r, _ = rtheta_from_xy(x, y, self.parameters.x0, self.parameters.y0) Rsq = self.parameters.R**2 fexpr = 2*self.parameters.Omega * (1 - 0.5 * r**2 / Rsq) if self.coriolis_trap is not None: From e3112ad1a9c2d8a31d8623b93563f3299e274c9b Mon Sep 17 00:00:00 2001 From: 12shughes Date: Mon, 2 Mar 2026 10:35:09 +0000 Subject: [PATCH 2/8] correct typo --- gusto/core/coord_transforms.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gusto/core/coord_transforms.py b/gusto/core/coord_transforms.py index b1bcb4f0d..43e00d34e 100644 --- a/gusto/core/coord_transforms.py +++ b/gusto/core/coord_transforms.py @@ -542,7 +542,7 @@ def xy_from_rtheta(r, theta, x0, y0, angle_units='rad'): r (:class:`np.ndarray` or :class:`ufl.Expr`): r-coordinate. theta (:class:`np.ndarray` or :class:`ufl.Expr`): theta-coordinate. x0: central x-coordinate for r-theta coords - y0: central y-coordinate for r-theta cooreds + y0: central y-coordinate for r-theta coords angle_units (str, optional): the units to use for the angle. Valid options are 'rad' (radians) or 'deg' (degrees). Defaults to 'rad'. Returns: From 5da099c7304da13520009deb48f3c8c885df54e8 Mon Sep 17 00:00:00 2001 From: 12shughes Date: Mon, 2 Mar 2026 10:35:28 +0000 Subject: [PATCH 3/8] correct typos --- gusto/core/equation_configuration.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/gusto/core/equation_configuration.py b/gusto/core/equation_configuration.py index a482ee37e..c16982c1b 100644 --- a/gusto/core/equation_configuration.py +++ b/gusto/core/equation_configuration.py @@ -7,7 +7,7 @@ __all__ = [ "BoussinesqParameters", "CompressibleParameters", "ShallowWaterParameters", "CoriolisOptions", - "SpongeLayerParameters", "DiffusionParameters", "BoundaryLayerParameters", + "SpongeLayerParameters", "DiffusionParameters", "BoundaryLayerParameters" ] @@ -125,8 +125,6 @@ class ShallowWaterParameters(EquationParameters): # Scaling factor for the saturation function in the equivalent buoyancy # formulation of the thermal shallow water equations q0 = None - # Scaling factor for evaporation term when present - cD = None class SpongeLayerParameters(EquationParameters): From 0e9e3257b671545d0b22c65d8ec8e652e3c35c03 Mon Sep 17 00:00:00 2001 From: 12shughes Date: Mon, 2 Mar 2026 10:35:50 +0000 Subject: [PATCH 4/8] change default behaviour --- gusto/equations/shallow_water_equations.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/gusto/equations/shallow_water_equations.py b/gusto/equations/shallow_water_equations.py index 7aaec6d2e..3b7479dfe 100644 --- a/gusto/equations/shallow_water_equations.py +++ b/gusto/equations/shallow_water_equations.py @@ -215,7 +215,12 @@ def _setup_residual(self, u_transport_option): xyz[1] - self.parameters.y0) elif rotation is CoriolisOptions.gammaplane: x, y = SpatialCoordinate(self.domain.mesh) - r, _ = rtheta_from_xy(x, y, self.parameters.x0, self.parameters.y0) + Lx = self.domain.mesh.coordinates.dat.data[:, 0].max() + Ly = self.domain.mesh.coordinates.dat.data[:, 1].max() + if not self.parameters.x0 and not self.parameters.y0: + r, _ = rtheta_from_xy(x, y, Lx/2, Ly/2) + else: + r, _ = rtheta_from_xy(x, y, self.parameters.x0, self.parameters.y0) Rsq = self.parameters.R**2 fexpr = 2*self.parameters.Omega * (1 - 0.5 * r**2 / Rsq) if self.coriolis_trap is not None: From 6a2879ad0a717ad093f5243d571c13332be87c3e Mon Sep 17 00:00:00 2001 From: 12shughes Date: Mon, 2 Mar 2026 10:46:53 +0000 Subject: [PATCH 5/8] add comments about gamma-plane defaults --- gusto/core/equation_configuration.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gusto/core/equation_configuration.py b/gusto/core/equation_configuration.py index c16982c1b..e8611edc9 100644 --- a/gusto/core/equation_configuration.py +++ b/gusto/core/equation_configuration.py @@ -111,8 +111,8 @@ class ShallowWaterParameters(EquationParameters): Omega = 7.292e-5 # rotation rate f0 = None # f-, beta- and gamma-plane Coriolis parameter beta = None # beta-plane y-variation parameter - y0 = None # beta- or gamma-plane y-centre - x0 = None # gamma-plane y-centre + y0 = None # beta- or gamma-plane y-centre. Gamma-plane default is Ly/2 + x0 = None # gamma-plane x-centre. Gamma-plane default is Lx/2 R = None # Radius of planet used to compute gamma in gamma-plane approx topog_expr = None # topography expression H = None # mean depth From 5133b47d71a7e213a4a05764342b0c1475052a9b Mon Sep 17 00:00:00 2001 From: 12shughes Date: Mon, 2 Mar 2026 11:52:35 +0000 Subject: [PATCH 6/8] remove option to specify centre of gamma-plane --- gusto/core/equation_configuration.py | 3 +-- gusto/equations/shallow_water_equations.py | 5 +---- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/gusto/core/equation_configuration.py b/gusto/core/equation_configuration.py index e8611edc9..d8d9dfce1 100644 --- a/gusto/core/equation_configuration.py +++ b/gusto/core/equation_configuration.py @@ -111,8 +111,7 @@ class ShallowWaterParameters(EquationParameters): Omega = 7.292e-5 # rotation rate f0 = None # f-, beta- and gamma-plane Coriolis parameter beta = None # beta-plane y-variation parameter - y0 = None # beta- or gamma-plane y-centre. Gamma-plane default is Ly/2 - x0 = None # gamma-plane x-centre. Gamma-plane default is Lx/2 + y0 = None # beta-plane y-centre R = None # Radius of planet used to compute gamma in gamma-plane approx topog_expr = None # topography expression H = None # mean depth diff --git a/gusto/equations/shallow_water_equations.py b/gusto/equations/shallow_water_equations.py index 3b7479dfe..fdb33397b 100644 --- a/gusto/equations/shallow_water_equations.py +++ b/gusto/equations/shallow_water_equations.py @@ -217,10 +217,7 @@ def _setup_residual(self, u_transport_option): x, y = SpatialCoordinate(self.domain.mesh) Lx = self.domain.mesh.coordinates.dat.data[:, 0].max() Ly = self.domain.mesh.coordinates.dat.data[:, 1].max() - if not self.parameters.x0 and not self.parameters.y0: - r, _ = rtheta_from_xy(x, y, Lx/2, Ly/2) - else: - r, _ = rtheta_from_xy(x, y, self.parameters.x0, self.parameters.y0) + r, _ = rtheta_from_xy(x, y, Lx/2, Ly/2) Rsq = self.parameters.R**2 fexpr = 2*self.parameters.Omega * (1 - 0.5 * r**2 / Rsq) if self.coriolis_trap is not None: From 84af116a03a6a947c5018786bbd7e9b45b8b55cf Mon Sep 17 00:00:00 2001 From: 12shughes Date: Tue, 3 Mar 2026 11:09:25 +0000 Subject: [PATCH 7/8] add checking script --- examples/shallow_water/check_gamma_plane.py | 81 +++++++++++++++++++++ 1 file changed, 81 insertions(+) create mode 100644 examples/shallow_water/check_gamma_plane.py diff --git a/examples/shallow_water/check_gamma_plane.py b/examples/shallow_water/check_gamma_plane.py new file mode 100644 index 000000000..7978ab60f --- /dev/null +++ b/examples/shallow_water/check_gamma_plane.py @@ -0,0 +1,81 @@ +from gusto import( + OutputParameters, ShallowWaterParameters, Domain, ShallowWaterEquations, + logger, PotentialVorticity, IO, SubcyclingOptions, TrapeziumRule, + SSPRK3, DGUpwind, SemiImplicitQuasiNewton, Function, + xy_from_rtheta, rtheta_from_xy, rtheta_from_lonlat, lonlat_from_rtheta, + CoriolisOptions +) +from firedrake import ( + SpatialCoordinate, VertexOnlyMesh, as_vector, pi, interpolate, exp, sqrt, + PeriodicRectangleMesh, conditional +) +import scipy +import numpy as np +import time +import os +import shutil +import sympy as sp + +folder_name_suffix = 'no_trap' + +nx = 256 +ny = nx +Lx = 7e7 +Ly = Lx + +rstar = Lx/2-3*Lx/nx +smooth_delta = 2 + +Bu = 1 +g=24.79 +Omega = 1.74e-4 +R = 71.4e6 +f0 = 2 * Omega +rm = 1e6 +phi0 = Bu * (f0*rm)**2 +H = phi0/g + +dt = 250 +tmax = 5*dt + +dirname=f'/data/home/sh1293/results/jupiter_sw/check_gamma_plane_{folder_name_suffix}' + +mesh = PeriodicRectangleMesh(nx=nx, ny=ny, Lx=Lx, Ly=Ly, quadrilateral=True) +output = OutputParameters(dirname=f'{dirname}', dumpfreq=1, dump_nc=True) + +parameters = ShallowWaterParameters(mesh, H=H, Omega=Omega, R=R, + rotation=CoriolisOptions.gammaplane) + +domain = Domain(mesh, dt, "RTCF", 1) + + +### rstar and smooth_delta are so the trap can be in the same place as my big script + +eqns = ShallowWaterEquations(domain, parameters)#, coriolis_trap=(0.5*rstar-smooth_delta*Lx/nx, 2*Omega)) + +diagnostic_fields = [PotentialVorticity()] + +io = IO(domain, output=output, diagnostic_fields=diagnostic_fields) + +subcycling_options = SubcyclingOptions(subcycle_by_courant=0.33) + +transport_methods = [ + DGUpwind(eqns, field_name) for field_name in eqns.field_names +] +transported_fields = [TrapeziumRule(domain, "u"), + SSPRK3(domain, "D", subcycling_options=subcycling_options)] +stepper = SemiImplicitQuasiNewton( + eqns, io, transported_fields, transport_methods +) + +u0 = stepper.fields("u") +D0 = stepper.fields("D") + +u0.assign(0.) +D0.assign(H) +Dbar = Function(D0.function_space()).assign(H) +stepper.set_reference_profiles([('D', Dbar)]) + +stepper.run(t=0, tmax=tmax) + +logger.info(f'File produced:\ncheck_gamma_plane_{folder_name_suffix}') \ No newline at end of file From bed7434c4aff8cf1601530aa23ed0744fcffbf6c Mon Sep 17 00:00:00 2001 From: 12shughes Date: Tue, 3 Mar 2026 12:19:04 +0000 Subject: [PATCH 8/8] update file --- examples/shallow_water/check_gamma_plane.py | 23 +++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/examples/shallow_water/check_gamma_plane.py b/examples/shallow_water/check_gamma_plane.py index 7978ab60f..ec1a4c7bb 100644 --- a/examples/shallow_water/check_gamma_plane.py +++ b/examples/shallow_water/check_gamma_plane.py @@ -16,7 +16,18 @@ import shutil import sympy as sp -folder_name_suffix = 'no_trap' + +# Currently trap=True works, and trap=False doesn't. + +# trap=False should give a gamma plane centred in the middle of the mesh, i.e. +# gradually decreasing PV from the centre. trap=True should be the same inside +# the trap radius, but be constant at the polar PV value outside that. +trap = True + +if trap: + folder_name_suffix = 'trap' +else: + folder_name_suffix = 'no_trap' nx = 256 ny = nx @@ -38,7 +49,7 @@ dt = 250 tmax = 5*dt -dirname=f'/data/home/sh1293/results/jupiter_sw/check_gamma_plane_{folder_name_suffix}' +dirname=f'check_gamma_plane_{folder_name_suffix}' mesh = PeriodicRectangleMesh(nx=nx, ny=ny, Lx=Lx, Ly=Ly, quadrilateral=True) output = OutputParameters(dirname=f'{dirname}', dumpfreq=1, dump_nc=True) @@ -48,10 +59,10 @@ domain = Domain(mesh, dt, "RTCF", 1) - -### rstar and smooth_delta are so the trap can be in the same place as my big script - -eqns = ShallowWaterEquations(domain, parameters)#, coriolis_trap=(0.5*rstar-smooth_delta*Lx/nx, 2*Omega)) +if trap: + eqns = ShallowWaterEquations(domain, parameters, coriolis_trap=(rstar-smooth_delta*Lx/nx, 2*Omega)) +else: + eqns = ShallowWaterEquations(domain, parameters) diagnostic_fields = [PotentialVorticity()]