Gamma plane

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

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

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)