diff --git a/examples/shallow_water/check_gamma_plane.py b/examples/shallow_water/check_gamma_plane.py new file mode 100644 index 000000000..ec1a4c7bb --- /dev/null +++ b/examples/shallow_water/check_gamma_plane.py @@ -0,0 +1,92 @@ +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 + + +# 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 +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'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) + +if trap: + eqns = ShallowWaterEquations(domain, parameters, coriolis_trap=(rstar-smooth_delta*Lx/nx, 2*Omega)) +else: + eqns = ShallowWaterEquations(domain, parameters) + +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 diff --git a/gusto/core/coord_transforms.py b/gusto/core/coord_transforms.py index 19726fe97..43e00d34e 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 coords 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..d8d9dfce1 100644 --- a/gusto/core/equation_configuration.py +++ b/gusto/core/equation_configuration.py @@ -1,7 +1,7 @@ """Some simple tools for configuring the model.""" -from enum import Enum from firedrake import Function, FunctionSpace, Constant import inspect +from enum import Enum __all__ = [ diff --git a/gusto/equations/shallow_water_equations.py b/gusto/equations/shallow_water_equations.py index 638207b44..fdb33397b 100644 --- a/gusto/equations/shallow_water_equations.py +++ b/gusto/equations/shallow_water_equations.py @@ -215,7 +215,9 @@ 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) + Lx = self.domain.mesh.coordinates.dat.data[:, 0].max() + Ly = self.domain.mesh.coordinates.dat.data[:, 1].max() + 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: