Relative Vorticity

This notebook shows a simple example of calculation the vertical component of relative vorticity,

\[\zeta = \partial_x v - \partial_y u.\]

For demonstration purposes we will compute the vorticity near the surface, but the method can be applied to any depth.

We will demonstrate three methods for computing relative voritcity:

  1. A naïve method of simply converting degrees of longitude/latitude into metres and differentiating. This method works if you simply wanna have a look at how the vorticity field looks like.

  2. A much more careful way of replicating exactly what the model does for computing the vertical component of relative vorticity.

  3. A simpler but accurate method leveraging the functionality of the xgcm.

Caveat: Both methods 2 and 3 automatically extend land masks into the domain, so needs extreme care if you care about coastlines! see https://github.com/xgcm/xgcm/issues/324

(To make a long story short: best method is probably method 3.)

Requirements: The conda/analysis3-24.04 (or later) module on the VDI/gadi (or your own up-to-date cookbook installation).

Firstly, load in the requisite libraries:

[1]:

import intake

import xarray as xr
import cf_xarray as cfxr
import numpy as np
from dask.distributed import Client

import matplotlib.pyplot as plt
import cmocean as cm

from matplotlib import rc

rc("font", **{"family": "sans-serif"})
rc("text", usetex=False)

Load a dask client.

[2]:

client = Client(threads_per_worker = 1)
client

[2]:

Client

Client-8385b172-ac68-11ef-a2a2-0000007dfe80

Connection method: Cluster object Cluster type: distributed.LocalCluster
Dashboard: /proxy/8787/status

Cluster Info

Use the ACCESS-NRI default catalog.

[3]:

catalog = intake.cat.access_nri

Various parameters that will be used for computing quantities.

[4]:

# values used by MOM5

Ω = 7.292e-5  # Earth's rotation rate; in radians/s
Rearth = 6371e3  # Earth's radius; in m

Load lon and lat used in the experiment. These datasets can be used for plotting.

[5]:

experiment = "025deg_jra55_iaf_omip2_cycle2"

datastore = catalog[experiment]

var_search = datastore.search(variable=["geolon_c", "geolat_c"])
var_search = var_search.search(path=var_search.df["path"][0])
grid_ds = var_search.to_dask(xarray_open_kwargs={"chunks": -1})
grid_ds = grid_ds.load()
lon = grid_ds["geolon_c"]
lat = grid_ds["geolat_c"]

Calculate the Coriolis parameter \(f = 2\Omega \sin(\texttt{lat})\).

[6]:

f = 2 * Ω * np.sin(np.deg2rad(lat))  # convert lat in radians
f = f.rename("Coriolis")
f.attrs["long_name"] = "Coriolis parameter"
f.attrs["units"] = "s-1"
f.attrs["coordinates"] = "geolon_c geolat_c"

Load the \(u\) and \(v\) velocity snapshots. Then we pick a depth value below the Ekman layer, namely the closest to \(z=-30\,m\).

The code .isel(time=-1) selects the final snapshot of u or v. Remove .isel(time=-1) to load all available snapshots of the flow fields.

[7]:

depth = 30  # m; avoid the surface Ekman layer by taking the "surface values" at some depth close to the surface

var_search = datastore.search(variable=["u", "v"])
ds = var_search.to_dask(xarray_open_kwargs={"chunks": {"yu_ocean": -1, "xu_ocean": -1}})

u = ds["u"].isel(time=-1).cf.sel(vertical=depth, method="nearest")
v = ds["v"].isel(time=-1).cf.sel(vertical=depth, method="nearest")

[8]:

ds = xr.merge([u, v, grid_ds], compat="override")

Method 1 (naïve computation)

To compute relative vorticity \(\zeta = \partial_x v - \partial_y u\) we simply differentiate the velocity components with respect of lon (here xu_ocean in degrees) and lat (here yu_ocean in degrees). We then convert the derivatives from units of degrees\(^{-1}\) to m\(^{-1}\). To do so, we use the value of the radius of the Earth Rearth and also take into account that as we go towards the poles the lon-grid spacing is scaled by \(\cos(\) lat \()\).

(Note the unicode characters like ζ can be used in python.)

[9]:

ζ_naive = v.differentiate("xu_ocean") / np.deg2rad(
    Rearth * np.cos(np.deg2rad(lat))
) - u.differentiate("yu_ocean") / (np.pi / 180 * Rearth)
ζ_naive = ζ_naive.rename("Relative Vorticity")
ζ_naive.attrs["long_name"] = "Relative Vorticity, ∂v/∂x-∂u/∂y"
ζ_naive.attrs["units"] = "s-1"

[10]:

ζ_naive

[10]:

<xarray.DataArray 'Relative Vorticity' (yu_ocean: 1080, xu_ocean: 1440)> Size: 6MB
dask.array<sub, shape=(1080, 1440), dtype=float32, chunksize=(1080, 1440), chunktype=numpy.ndarray>
Coordinates:
    time      datetime64[ns] 8B 2018-12-14T12:00:00
    st_ocean  float64 8B 30.36
  * xu_ocean  (xu_ocean) float64 12kB -279.8 -279.5 -279.2 ... 79.5 79.75 80.0
  * yu_ocean  (yu_ocean) float64 9kB -81.02 -80.92 -80.81 ... 89.79 89.89 90.0
    geolat_c  (yu_ocean, xu_ocean) float32 6MB nan nan nan nan ... nan nan nan
    geolon_c  (yu_ocean, xu_ocean) float32 6MB nan nan nan nan ... nan nan nan
Attributes:
    long_name:  Relative Vorticity, ∂v/∂x-∂u/∂y
    units:      s-1

We now plot \(\zeta\) in the North Atlantic.

[11]:

maxvalue = 6e-6
levels = np.linspace(-maxvalue, maxvalue, 24)

[12]:

plt.figure(figsize=(10, 5))

ζ_naive.plot.contourf(
    levels=levels,
    x="geolon_c",
    y="geolat_c",
    cmap="RdBu_r",
    add_labels=False,
)
plt.title(
    "snapshot of relative vorticity, $\zeta = \partial_x v - \partial_y u$ [naïve]"
)
plt.xlabel("longitude")
plt.ylabel("latitude")
plt.xlim(-120, 20)
plt.ylim(20, 90);
../_images/03-Advanced-Recipes_Relative_Vorticity_21_0.png

This method is sort-of-OK below 65N where the complications of the tripolar grid do not enter. But even south of 65N this method gives a rough estimate. It’s OK for visualising but it should not be used in computing, e.g., vorticity budgets.

Method 2: replicating how MOM5 computes vorticity_z

Looking at MOM5 code and by doing some translation from Fortran code to “english” we can see that the models computes vorticity_z via:

\[\begin{split} \zeta(i, j, k) = \frac1{2}\Big[ \underbrace{ \frac{v(i,j,k) - v(i-1,j,k)}{\Delta x_{N}(i,j,k)} }_{\approx \partial_x v(i,j,k)} + \underbrace{\frac{v(i,j-1,k) - v(i-1,j-1,k)}{\Delta x_{N}(i,j-1,k)}}_{\approx \partial_x v(i,j-1,k)} \Big] \\ \qquad \qquad \qquad - \frac1{2}\Big[ \underbrace{ \frac{u(i,j,k) - u(i,j-1,k)}{\Delta y_{E}(i,j,k)} }_{\approx \partial_y u(i,j,k)} + \underbrace{\frac{u(i-1,j,k) - u(i-1,j-1,k)}{\Delta y_{E}(i-1,j,k)}}_{\approx \partial_y u(i,j-1,k)} \Big].\end{split}\]

Above, \((i, j, k)\) refers to the grid-point indices for directions \(x, y, z\) respectively.

The distances \(\Delta x_N\) and \(\Delta y_E\) correspond to the North and East faces of the corresponding \(T\)-cell; see the diagram below. (T-points are those in the cells’s centres where tracers are evaluated and U-points are in the cells’s corners where velocities are evaluated.)

../_images/Bgrid.png

After a lot of fiddling with the indices you can confirm that the right way to compute \(\delta x_N\) and \(\delta y_E\) from above is:

[13]:

inverse_dxtn = 0.5 * (1 / ds.dxu + np.roll(1 / ds.dxu, 1, axis=1))
inverse_dyte = 0.5 * (1 / ds.dyu + np.roll(1 / ds.dyu, 1, axis=0))

… and with that we can now compute ζ_mom5 which is the vorticity as exactly the way is computed by the model:

[14]:

vx_ijk = (ds.v - np.roll(ds.v, 1, axis=1)) * inverse_dxtn
uy_ijk = (ds.u - np.roll(ds.u, 1, axis=0)) * inverse_dyte

vx = 0.5 * (vx_ijk + np.roll(vx_ijk, 1, axis=0))
uy = 0.5 * (uy_ijk + np.roll(uy_ijk, 1, axis=1))

ζ_mom5 = vx - uy

/g/data/hh5/public/apps/miniconda3/envs/analysis3-24.07/lib/python3.10/site-packages/distributed/client.py:3371: UserWarning: Sending large graph of size 11.87 MiB.
This may cause some slowdown.
Consider loading the data with Dask directly
 or using futures or delayed objects to embed the data into the graph without repetition.
See also https://docs.dask.org/en/stable/best-practices.html#load-data-with-dask for more information.
  warnings.warn(
/g/data/hh5/public/apps/miniconda3/envs/analysis3-24.07/lib/python3.10/site-packages/distributed/client.py:3371: UserWarning: Sending large graph of size 11.87 MiB.
This may cause some slowdown.
Consider loading the data with Dask directly
 or using futures or delayed objects to embed the data into the graph without repetition.
See also https://docs.dask.org/en/stable/best-practices.html#load-data-with-dask for more information.
  warnings.warn(

The way we computed ζ_mom5 above using numpy.roll is not really ideal. One would like to do the above using xarray’s functionality… But, as we will argue below, you shouldn’t be using this way anyway because it’s cumbersome and it requires you to know the specifics of the staggered grid used in the model (Arakawa B-grid for MOM5). A way to avoid getting into the nitty-gritty of staggered grids is to use xgcm packaged (described in Method 3).

With this and that let’s plot ζ_mom5.

[15]:

plt.figure(figsize=(10, 5))

ζ_mom5.plot.contourf(
    levels=levels,
    x="geolon_c",
    y="geolat_c",
    cmap="RdBu_r",
    add_labels=False,
)

plt.title(
    "snapshot of relative vorticity, $\zeta = \partial_x v - \partial_y u$ [as in MOM5]"
)
plt.xlabel("longitude")
plt.ylabel("latitude")
plt.xlim(-120, 20)
plt.ylim(20, 90);

/g/data/hh5/public/apps/miniconda3/envs/analysis3-24.07/lib/python3.10/site-packages/distributed/client.py:3371: UserWarning: Sending large graph of size 35.60 MiB.
This may cause some slowdown.
Consider loading the data with Dask directly
 or using futures or delayed objects to embed the data into the graph without repetition.
See also https://docs.dask.org/en/stable/best-practices.html#load-data-with-dask for more information.
  warnings.warn(
../_images/03-Advanced-Recipes_Relative_Vorticity_31_1.png

Note that the Arctic region now looks more “reasonable”.

Method 3: Using xgcm to replicate MOM5’s calculation for vorticity_z

`xgcm <https://xgcm.readthedocs.io/en/stable/>`__ is a package that deals with staggered grids that are typically used in ocean models. An excerpt from xgcm’s docs mentions:

“(in model output datasets), different variables are located at different positions with respect to a volume or area element (e.g. cell center, cell face, etc.) xgcm solves the problem of how to interpolate and difference these variables from one position to another.”

[16]:

import xgcm

print("xgcm version ", xgcm.__version__)

xgcm version  0.8.1

The way xgcm works is that we first create a grid object that has all the information regarding our staggered grid. For our case, grid needs to know the location of the xt_ocean, xu_ocean points (and same for \(y\)) and their relative orientation to one another, i.e., that xu_ocean is shifted to the right of xt_ocean by \(\frac1{2}\) grid-cell.

[17]:

ds.coords["xt_ocean"].attrs.update(axis="X")
ds.coords["xu_ocean"].attrs.update(axis="X", c_grid_axis_shift=0.5)
ds.coords["yt_ocean"].attrs.update(axis="Y")
ds.coords["yu_ocean"].attrs.update(axis="Y", c_grid_axis_shift=0.5)

grid = xgcm.Grid(ds, periodic=["X"])

Then, xgcm give you a way to interpolate between grids (with the .interp function) and a way to compute differences (.diff function). For example, the exrpession \(v(i,j,k) - v(i-1,j,k)\) is obtained via grid.diff(ds.v, 'X').

Using xgcm’s functionality we can replicate the MOM5 vertical vorticity computation as:

[18]:

vx = grid.interp(grid.diff(ds.v, "X") / grid.interp(ds.dxu, "X"), "Y", boundary="extend")
uy = grid.interp(grid.diff(ds.u, "Y", boundary="extend") / grid.interp(ds.dyt, "X"), "X")

ζ_xgcm = vx - uy

ζ_xgcm = ζ_xgcm.rename("Relative Vorticity")
ζ_xgcm.attrs["long_name"] = "Relative Vorticity, ∂v/∂x-∂u/∂y"
ζ_xgcm.attrs["units"] = "s-1"

/g/data/hh5/public/apps/miniconda3/envs/analysis3-24.07/lib/python3.10/site-packages/xgcm/grid_ufunc.py:832: FutureWarning: The return type of `Dataset.dims` will be changed to return a set of dimension names in future, in order to be more consistent with `DataArray.dims`. To access a mapping from dimension names to lengths, please use `Dataset.sizes`.
  out_dim: grid._ds.dims[out_dim] for arg in out_core_dims for out_dim in arg
/g/data/hh5/public/apps/miniconda3/envs/analysis3-24.07/lib/python3.10/site-packages/xgcm/grid_ufunc.py:832: FutureWarning: The return type of `Dataset.dims` will be changed to return a set of dimension names in future, in order to be more consistent with `DataArray.dims`. To access a mapping from dimension names to lengths, please use `Dataset.sizes`.
  out_dim: grid._ds.dims[out_dim] for arg in out_core_dims for out_dim in arg
/g/data/hh5/public/apps/miniconda3/envs/analysis3-24.07/lib/python3.10/site-packages/xgcm/grid_ufunc.py:832: FutureWarning: The return type of `Dataset.dims` will be changed to return a set of dimension names in future, in order to be more consistent with `DataArray.dims`. To access a mapping from dimension names to lengths, please use `Dataset.sizes`.
  out_dim: grid._ds.dims[out_dim] for arg in out_core_dims for out_dim in arg
/g/data/hh5/public/apps/miniconda3/envs/analysis3-24.07/lib/python3.10/site-packages/xgcm/grid_ufunc.py:832: FutureWarning: The return type of `Dataset.dims` will be changed to return a set of dimension names in future, in order to be more consistent with `DataArray.dims`. To access a mapping from dimension names to lengths, please use `Dataset.sizes`.
  out_dim: grid._ds.dims[out_dim] for arg in out_core_dims for out_dim in arg
/g/data/hh5/public/apps/miniconda3/envs/analysis3-24.07/lib/python3.10/site-packages/xgcm/grid_ufunc.py:832: FutureWarning: The return type of `Dataset.dims` will be changed to return a set of dimension names in future, in order to be more consistent with `DataArray.dims`. To access a mapping from dimension names to lengths, please use `Dataset.sizes`.
  out_dim: grid._ds.dims[out_dim] for arg in out_core_dims for out_dim in arg
/g/data/hh5/public/apps/miniconda3/envs/analysis3-24.07/lib/python3.10/site-packages/xgcm/grid_ufunc.py:832: FutureWarning: The return type of `Dataset.dims` will be changed to return a set of dimension names in future, in order to be more consistent with `DataArray.dims`. To access a mapping from dimension names to lengths, please use `Dataset.sizes`.
  out_dim: grid._ds.dims[out_dim] for arg in out_core_dims for out_dim in arg

(Note: A simpler expression for ζ_xgcm could be:

ζ_xgcm = (+ grid.interp(grid.diff(ds.v, 'X'), 'Y', boundary='extend') / ds.dxt
          - grid.interp(grid.diff(ds.u, 'Y', boundary='extend'), 'X') / ds.dyt)

However, the above does not replicate precisely the way MOM5 computes vorticity_z.)

Now, let’s plot ζ_xgcm…

[19]:

fig = plt.figure(figsize=(10, 5))
plt.contourf(lon, lat, ζ_xgcm.values, levels=levels, cmap="RdBu_r", extend="both")

plt.title(
    "snapshot of relative vorticity, $\zeta = \partial_x v - \partial_y u$ [by xgcm]"
)
plt.colorbar()
plt.xlim(-120, 20)
plt.ylim(20, 90);

/g/data/hh5/public/apps/miniconda3/envs/analysis3-24.07/lib/python3.10/site-packages/distributed/client.py:3371: UserWarning: Sending large graph of size 11.87 MiB.
This may cause some slowdown.
Consider loading the data with Dask directly
 or using futures or delayed objects to embed the data into the graph without repetition.
See also https://docs.dask.org/en/stable/best-practices.html#load-data-with-dask for more information.
  warnings.warn(
../_images/03-Advanced-Recipes_Relative_Vorticity_42_1.png

Looks OK and it was definitely easier than method 2.

Comparison of the three methods

Now let’s be thorough and compare the vorticities computed via the three methods.

[20]:

fig = plt.figure(figsize=(10, 5))
plt.contourf(
    lon,
    lat,
    np.abs(ζ_mom5.values - ζ_naive.values),
    levels=np.linspace(0, maxvalue, 22),
    cmap="Oranges",
    extend="max",
)

plt.title("Difference $|$MOM5$-$naive$|$")
plt.colorbar()
plt.xlim(-120, 20)
plt.ylim(20, 90);

/g/data/hh5/public/apps/miniconda3/envs/analysis3-24.07/lib/python3.10/site-packages/distributed/client.py:3371: UserWarning: Sending large graph of size 35.60 MiB.
This may cause some slowdown.
Consider loading the data with Dask directly
 or using futures or delayed objects to embed the data into the graph without repetition.
See also https://docs.dask.org/en/stable/best-practices.html#load-data-with-dask for more information.
  warnings.warn(
../_images/03-Advanced-Recipes_Relative_Vorticity_46_1.png 
[21]:

plt.figure(figsize=(10, 5))
plt.contourf(
    lon,
    lat,
    np.abs(ζ_mom5.values - ζ_xgcm.values),
    levels=np.linspace(0, maxvalue, 22),
    cmap="Oranges",
    vmin=0,
    vmax=maxvalue,
    extend="max",
)

plt.title("Difference $|$MOM5$-$xgcm$|$")
plt.colorbar()
plt.xlim(-120, 20)
plt.ylim(20, 90);
../_images/03-Advanced-Recipes_Relative_Vorticity_47_0.png

Indeed the xgcm method gives the same as the MOM5 code.

We have confirmed that indeed ζ_mom5 corresponds exactly to the vorticity_z diagnostic that the model outputs. To test this, the model was run for 1 day, save u, v and vorticity_z and then confirm that the ζ_mom5 computation from method 2 is equal to vorticity_z (up to machine precision). For brevity this comparison is not included here.

Further Potential Improvements on Method 3

xgcm has a functionality to perform derivatives and interpolations using the grid metrics. Implementing that would simplify the vorticity calculations even more, e.g.,

ζ_xgcm = grid.derivative(ds.v, 'X') - grid.derivative(ds.u, 'Y')

Something like that would be terrific; the user will only have to type up the algebraic expression of what they want to compute without any reference to dxt or dxu distances…

The Rossby number

To conclude, let’s visualize the Rossby number

\[\mathrm{Ro} = \frac{\zeta }{f},\]

where \(f=2\Omega\sin(\theta)\) is the Coriolis parameter.

[22]:

f = grid.interp(grid.interp(f, "X"), "Y", boundary="extend")

Ro = ζ_xgcm / f
Ro = Ro.rename("Rossby number ζ/f")

/g/data/hh5/public/apps/miniconda3/envs/analysis3-24.07/lib/python3.10/site-packages/xgcm/grid_ufunc.py:832: FutureWarning: The return type of `Dataset.dims` will be changed to return a set of dimension names in future, in order to be more consistent with `DataArray.dims`. To access a mapping from dimension names to lengths, please use `Dataset.sizes`.
  out_dim: grid._ds.dims[out_dim] for arg in out_core_dims for out_dim in arg

[23]:

plt.figure(figsize=(10, 5))
plt.contourf(
    lon,
    lat,
    Ro,
    levels=np.linspace(-0.1, 0.1, 51),
    cmap="RdBu_r",
    extend="both",
)

plt.title("Rossby number $\zeta/f$")
plt.colorbar()
plt.xlim(-120, 20)
plt.ylim(20, 90);

/g/data/hh5/public/apps/miniconda3/envs/analysis3-24.07/lib/python3.10/site-packages/distributed/client.py:3371: UserWarning: Sending large graph of size 17.81 MiB.
This may cause some slowdown.
Consider loading the data with Dask directly
 or using futures or delayed objects to embed the data into the graph without repetition.
See also https://docs.dask.org/en/stable/best-practices.html#load-data-with-dask for more information.
  warnings.warn(
../_images/03-Advanced-Recipes_Relative_Vorticity_52_1.png 
[24]:

client.close()