update files to allow r-theta coords centred off 0,0

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}')

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)
 

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__ = [

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: