diff --git a/gusto/core/coord_transforms.py b/gusto/core/coord_transforms.py index f40456819..19726fe97 100644 --- a/gusto/core/coord_transforms.py +++ b/gusto/core/coord_transforms.py @@ -11,7 +11,8 @@ __all__ = ["xyz_from_lonlatr", "lonlatr_from_xyz", "xyz_vector_from_lonlatr", "lonlatr_components_from_xyz", "rodrigues_rotation", "pole_rotation", "rotated_lonlatr_vectors", "rotated_lonlatr_coords", - "periodic_distance", "great_arc_angle"] + "periodic_distance", "great_arc_angle", "xy_from_rtheta", + "rtheta_from_xy", "rtheta_from_lonlat", "lonlat_from_rtheta"] def firedrake_or_numpy(variable): @@ -531,3 +532,155 @@ def great_arc_angle(lon1, lat1, lon2, lat2, units='rad'): arc_length *= 180.0 / pi return arc_length + + +def xy_from_rtheta(r, theta, 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. + angle_units (str, optional): the units to use for the angle. Valid + options are 'rad' (radians) or 'deg' (degrees). Defaults to 'rad'. + Returns: + tuple of :class`np.ndarray` or tuple of :class:`ufl.Expr`: the tuple + of (x, y) coordinates in the appropriate form given the + provided arguments. + """ + + # Import routines + module, _ = firedrake_or_numpy(r) + cos = module.cos + sin = module.sin + pi = module.pi + + if angle_units not in ['rad', 'deg']: + raise ValueError(f'angle_units arg {angle_units} not valid') + + if angle_units == 'deg': + unit_factor = 180./pi + if angle_units == 'rad': + unit_factor = 1.0 + + theta = theta/unit_factor + + x = r * cos(theta) + y = r * sin(theta) + + return x, y + + +def rtheta_from_xy(x, y, 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. + angle_units (str, optional): the units to use for the angle. Valid + options are 'rad' (radians) or 'deg' (degrees). Defaults to 'rad'. + Returns: + tuple of :class`np.ndarray` or tuple of :class:`ufl.Expr`: the tuple + of (r, theta) coordinates in the appropriate form given the + provided arguments. + """ + + # Import routines + module, _ = firedrake_or_numpy(x) + atan2 = module.atan2 if hasattr(module, "atan2") else module.arctan2 + sqrt = module.sqrt + pi = module.pi + + if angle_units not in ['rad', 'deg']: + raise ValueError(f'angle_units arg {angle_units} not valid') + + if angle_units == 'deg': + unit_factor = 180./pi + if angle_units == 'rad': + unit_factor = 1.0 + + theta = atan2(y, x) + r = sqrt(x**2 + y**2) + + return r, theta*unit_factor + + +def rtheta_from_lonlat(lon, lat, R, angle_units='rad', pole='north'): + """ + Returns the polar r theta coordinates from the lon lat coordinates + Args: + lon (:class:`np.ndarray` or :class:`ufl.Expr`): lon-coordinate. + lat (:class:`np.ndarray` or :class:`ufl.Expr`): lat-coordinate. + R (float): Planetary radius. + angle_units (str, optional): the units to use for the angle. Valid + options are 'rad' (radians) or 'deg' (degrees). Defaults to 'rad'. + pole (str, optional): which pole? Valid options are 'north' or 'south'. + Defaults to 'north'. + Returns: + tuple of :class`np.ndarray` or tuple of :class:`ufl.Expr`: the tuple + of (r, theta) coordinates in the appropriate form given the + provided arguments. + """ + + # Import routines + module, _ = firedrake_or_numpy(lon) + pi = module.pi + + if angle_units not in ['rad', 'deg']: + raise ValueError(f'angle_units arg {angle_units} not valid') + + if angle_units == 'deg': + unit_factor = 180./pi + if angle_units == 'rad': + unit_factor = 1.0 + + if pole == 'south': + lat *= -1 + + lon_scaled = lon/unit_factor + lat_scaled = lat/unit_factor + + theta = pi/2-lon_scaled + r = R*(pi/2-lat_scaled) + + return r, theta*unit_factor + + +def lonlat_from_rtheta(r, theta, R, angle_units='rad', pole='north'): + """ + Returns the lon lat coordinates from the polar r theta coordinates + Args: + r (:class:`np.ndarray` or :class:`ufl.Expr`): r-coordinate. + theta (:class:`np.ndarray` or :class:`ufl.Expr`): theta-coordinate. + R (float): Planetary radius. + angle_units (str, optional): the units to use for the angle. Valid + options are 'rad' (radians) or 'deg' (degrees). Defaults to 'rad'. + pole (str, optional): which pole? Valid options are 'north' or 'south'. + Defaults to 'north'. + Returns: + tuple of :class`np.ndarray` or tuple of :class:`ufl.Expr`: the tuple + of (lon, lat) coordinates in the appropriate form given the + provided arguments. + """ + + # Import routines + module, _ = firedrake_or_numpy(r) + pi = module.pi + + if angle_units not in ['rad', 'deg']: + raise ValueError(f'angle_units arg {angle_units} not valid') + + if angle_units == 'deg': + unit_factor = 180./pi + if angle_units == 'rad': + unit_factor = 1.0 + + theta_scaled = theta/unit_factor + + lon = pi/2-theta_scaled + lat = pi/2 - r/R + if pole == 'south': + lat *= -1 + + return lon*unit_factor, lat*unit_factor diff --git a/gusto/core/equation_configuration.py b/gusto/core/equation_configuration.py index 7f7a87f8f..1a3b552e2 100644 --- a/gusto/core/equation_configuration.py +++ b/gusto/core/equation_configuration.py @@ -13,10 +13,11 @@ class CoriolisOptions(Enum): - nonrotating = 11 # no Coriolis term - fplane = 23 # constant Coriolis term equal to f0 - betaplane = 31 # Coriolis term equal to f0 + beta(y-y0) - sphere = 42 # Coriolis term equal to 2 Omega sin(lat) + nonrotating = 11 # no Coriolis term + fplane = 23 # constant Coriolis term f0 + betaplane = 31 # Coriolis term f0 + beta(y-y0) + gammaplane = 37 # Coriolis term f0 - gamma r**2 with gamma=Omega/R**2 and f0=2Omega + sphere = 42 # Coriolis term 2 Omega sin(lat) class EquationParameters(object): @@ -108,9 +109,10 @@ class ShallowWaterParameters(EquationParameters): # Coriolis options: controlled by rotation parameter rotation = CoriolisOptions.sphere # type of Coriolis term Omega = 7.292e-5 # rotation rate - f0 = None # f- and beta- plane Coriolis parameter + f0 = None # f-, beta- and gamma-plane Coriolis parameter beta = None # beta-plane y-variation parameter 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 # Factor that multiplies the vapour in the equivalent buoyancy diff --git a/gusto/equations/shallow_water_equations.py b/gusto/equations/shallow_water_equations.py index e075f25f2..638207b44 100644 --- a/gusto/equations/shallow_water_equations.py +++ b/gusto/equations/shallow_water_equations.py @@ -5,6 +5,7 @@ SpatialCoordinate, Constant, TrialFunctions, TestFunctions, dot, grad, interpolate, assemble) from firedrake.fml import subject, drop, all_terms +from gusto.core.coord_transforms import rtheta_from_xy from gusto.core.labels import ( linearisation, pressure_gradient, coriolis, prognostic ) @@ -39,6 +40,7 @@ class ShallowWaterEquations(PrognosticEquationSet): def __init__(self, domain, parameters, space_names=None, linearisation_map=all_terms, u_transport_option='vector_invariant_form', + coriolis_trap=None, no_normal_flow_bc_ids=None, active_tracers=None): """ Args: @@ -60,6 +62,10 @@ def __init__(self, domain, parameters, 'vector_invariant_form', 'vector_advection_form', and 'circulation_form'. Defaults to 'vector_invariant_form'. + coriolis_trap (tuple, optional): (r_edge, f_trap) specifies a + distance r_edge and `trap' function f_trap such that the + Coriolis parameter is set to f_trap(r) for r>r_edge. For use + with gamma-plane setup only. Defaults to None. no_normal_flow_bc_ids (list, optional): a list of IDs of domain boundaries at which no normal flow will be enforced. Defaults to None. @@ -83,6 +89,10 @@ def __init__(self, domain, parameters, self.domain = domain self.active_tracers = active_tracers + if coriolis_trap is not None: + assert self.parameters.rotation is CoriolisOptions.gammaplane, "coriolis_trap option is only for use with gamma-plane setup" + self.coriolis_trap = coriolis_trap + self._setup_residual(u_transport_option) # -------------------------------------------------------------------- # @@ -201,7 +211,16 @@ def _setup_residual(self, u_transport_option): fexpr = self.parameters.f0 elif rotation is CoriolisOptions.betaplane: xyz = SpatialCoordinate(self.domain.mesh) - fexpr = self.parameters.f0 + self.parameters.beta * xyz[1] + fexpr = self.parameters.f0 + self.parameters.beta * ( + xyz[1] - self.parameters.y0) + elif rotation is CoriolisOptions.gammaplane: + x, y = SpatialCoordinate(self.domain.mesh) + r, _ = rtheta_from_xy(x, y) + Rsq = self.parameters.R**2 + fexpr = 2*self.parameters.Omega * (1 - 0.5 * r**2 / Rsq) + if self.coriolis_trap is not None: + r_edge, f_trap = self.coriolis_trap + fexpr = conditional(r < r_edge, fexpr, f_trap) else: raise NotImplementedError('Coriolis option is not implemented') self.prescribed_fields('coriolis', CG1).interpolate(fexpr)