NEMO idealized basin example¶
This example shows how to use xgcm to perform finite-volume calculations on the
output of the NEMO ocean model. The output has been
saved to netCDF with grid-aware coordinate names using the
xnemogcm package, which reads NEMO datasets
so that they are directly understandable by xgcm.
The configuration is BASIN, an idealized, closed rectangular ocean basin on a
Mercator grid. Below we build an xgcm.Grid from the dataset and use it to compute
a sequence of increasingly involved diagnostics: horizontal gradients, the flow
divergence, the vertical velocity, the relative and potential vorticity, the
barotropic streamfunction, the kinetic energy, and finally vertical coordinate
transformations with grid.transform.
We start by importing xarray, numpy and xgcm:
import numpy as np
import xarray as xr
import xgcm
from matplotlib import pyplot as plt
# Compact repr for xarray objects
xr.set_options(display_expand_attrs=False, display_expand_data=False)
%matplotlib inline
plt.rcParams["figure.figsize"] = (5, 6)
print("xgcm version:", xgcm.__version__)
xgcm version: 0.1.dev638+g7e2d5db5e
Loading the dataset¶
We use a dataset produced for this documentation: the time-averaged output of the reference run from Caneill et al. (2022, JPO), The Polar Transition from Alpha to Beta Regions Set by a Surface Buoyancy Flux Inversion. It is distributed as a single netCDF file on Zenodo:
The file is downloaded once into a local nemo_idealized_data/ folder. Because the
dataset is already time-averaged, it has no time dimension. We promote the geographic
coordinate variables (those starting with g, e.g. glamt, gphit) to coordinates
so that they can be used directly as plot axes.
# download the data (only the first time)
import urllib.request
import shutil
from pathlib import Path
url = "https://zenodo.org/record/7795560/files/nemo_dataset.nc"
name = "nemo_dataset.nc"
path = Path("nemo_idealized_data/")
path.mkdir(parents=True, exist_ok=True)
if not (path / name).exists():
with urllib.request.urlopen(url) as response, open(path / name, "wb") as out_file:
shutil.copyfileobj(response, out_file)
# open the data
ds = xr.open_dataset(path / name)
# promote the geographic coordinates (glamt, gphit, ...) to coordinates
for v in ds.variables:
if v.startswith("g"):
ds.coords[v] = ds[v]
ds
<xarray.Dataset> Size: 14MB
Dimensions: (z_c: 36, axis_nbounds: 2, y_c: 79, x_c: 42, x_f: 42,
y_f: 79, z_f: 36)
Coordinates: (12/18)
* z_c (z_c) int64 288B 0 1 2 3 4 5 6 7 ... 28 29 30 31 32 33 34 35
gdept_1d (z_c) float64 288B ...
* y_c (y_c) int64 632B 0 1 2 3 4 5 6 7 ... 71 72 73 74 75 76 77 78
* x_c (x_c) int64 336B 0 1 2 3 4 5 6 7 ... 34 35 36 37 38 39 40 41
glamt (y_c, x_c) float64 27kB ...
glamv (y_f, x_c) float64 27kB ...
... ...
glamf (y_f, x_f) float64 27kB ...
gphiu (y_c, x_f) float64 27kB ...
gphif (y_f, x_f) float64 27kB ...
* y_f (y_f) float64 632B 0.5 1.5 2.5 3.5 ... 75.5 76.5 77.5 78.5
* z_f (z_f) float64 288B -0.5 0.5 1.5 2.5 ... 31.5 32.5 33.5 34.5
gdepw_1d (z_f) float64 288B ...
Dimensions without coordinates: axis_nbounds
Data variables: (12/46)
deptht_bounds (z_c, axis_nbounds) float32 288B ...
e3t (z_c, y_c, x_c) float32 478kB ...
thetao (z_c, y_c, x_c) float32 478kB ...
so (z_c, y_c, x_c) float32 478kB ...
tos (y_c, x_c) float32 13kB ...
zos (y_c, x_c) float32 13kB ...
... ...
hv_0 (y_f, x_c) float64 27kB ...
tmask (z_c, y_c, x_c) int8 119kB ...
umask (z_c, y_c, x_f) int8 119kB ...
vmask (z_c, y_f, x_c) int8 119kB ...
fmask (z_c, y_f, x_f) int8 119kB ...
mbathy (y_c, x_c) int32 13kB ...This configuration is symmetric about the equator. For simplicity we restrict the example to the northern half of the basin by masking out everything south of the equator on each grid point:
ds["tmask"] = ds.tmask.where(ds.y_c > 0, 0)
ds["umask"] = ds.umask.where(ds.y_c > 0, 0)
ds["vmask"] = ds.vmask.where(ds.y_f > 0, 0)
ds["fmask"] = ds.fmask.where(ds.y_f > 0, 0)
In the dimension names, the _c suffix means center while the _f suffix means
face. The Arakawa C-grid points are therefore:
| point | dimensions |
|---|---|
| T (tracer) | (z_c, y_c, x_c) |
| U (zonal velocity) | (z_c, y_c, x_f) |
| V (meridional velocity) | (z_c, y_f, x_c) |
| F (vorticity) | (z_c, y_f, x_f) |
| W (vertical velocity) | (z_f, y_c, x_c) |
The scale factors e1*, e2*, e3* are the grid cell widths in the X, Y and Z
directions at each of these points.
Geometry of the basin¶
The domain is a closed basin with a bottom bathymetry deepening from 2000 m at the coasts to 4000 m in the interior. Terrain-following coordinates are used with a linear free surface (fixed vertical levels), on a 1° Mercator grid.
fig, ax = plt.subplots(figsize=(4.5, 6))
ds.ht_0.where(ds.tmask.isel(z_c=0)).plot.contourf(
x="glamt", y="gphit", vmin=2000, vmax=4000, levels=21,
cmap="Blues", ax=ax,
cbar_kwargs={"label": "Bottom depth [m]"},
)
ax.set(xlabel="Longitude [°E]", ylabel="Latitude [°N]", title="Basin bathymetry")
plt.show()
Creating the grid object¶
We now create an xgcm.Grid from the dataset. All axes are non-periodic (closed
basin). The metrics dictionary tells xgcm which scale factors measure distances
along each axis, so that operators such as derivative, integrate and average
know which cell widths to use.
metrics = {
("X",): ["e1t", "e1u", "e1v", "e1f"], # X cell widths
("Y",): ["e2t", "e2u", "e2v", "e2f"], # Y cell widths
("Z",): ["e3t", "e3u", "e3v", "e3w"], # Z cell widths
}
grid = xgcm.Grid(ds, metrics=metrics, periodic=False)
print(grid)
<xgcm.Grid> X Axis (not periodic, boundary='fill'): * center x_c --> right * right x_f --> center Z Axis (not periodic, boundary='fill'): * center z_c --> left * left z_f --> center Y Axis (not periodic, boundary='fill'): * center y_c --> right * right y_f --> center
xgcm has identified three axes: X (longitude), Y (latitude) and Z (depth). There
is no time axis because the dataset has already been time-averaged.
Computation examples¶
Horizontal gradient of SST¶
As a first example we compute the zonal component of the gradient of sea surface
temperature (SST), $\frac{\partial \mathrm{SST}}{\partial x}$. The SST is the variable
tos (Temperature Ocean Surface).
In discrete form, using NEMO notation, this is [1]
$$\frac{\partial \mathrm{SST}}{\partial x} = \frac{1}{e_{1u}}\,\delta_{i+1/2}\,\mathrm{SST} = \frac{1}{e_{1u}}\left(\mathrm{SST}_{i+1} - \mathrm{SST}_i\right).$$
The last T point along X is a land point (as are the last two U points), so we set
boundary='fill' with fill_value=0; the precise fill value is unimportant because
it only affects land points. We compute the gradient two ways: explicitly with diff
divided by the scale factor e1u, and with derivative, which knows which scale
factor to use. The two give the same result, located on the U point.
[1] NEMO book v4.0.1, p. 22
sst = ds.tos
grad_T_lon0 = grid.diff(sst, axis="X", boundary="fill", fill_value=0) / ds.e1u
grad_T_lon1 = grid.derivative(sst, axis="X", boundary="fill", fill_value=0)
print("output dimensions:", grad_T_lon1.dims)
assert np.allclose(grad_T_lon0, grad_T_lon1)
print("diff/e1u and derivative agree")
output dimensions: ('y_c', 'x_f')
diff/e1u and derivative agree
Divergence¶
Next we compute the divergence of the flow. The details depend strongly on the model configuration (free surface vs. rigid lid, etc.). Here the flow is incompressible with a linear free surface and satisfies the continuity equation
$$\vec{\nabla}\cdot\vec{u} = \frac{\partial u}{\partial x} + \frac{\partial v}{\partial y} + \frac{\partial w}{\partial z} = 0.$$
(With a non-linear free surface the divergence of $\vec{u}$ is no longer zero; it is balanced by the time variation of the $e_3$ scale factors.)
In discrete form, using NEMO notation [1]
$$\vec{\nabla}\cdot\vec{u} = \frac{1}{e_{1t}e_{2t}e_{3t}}\left[ \delta_i(u\,e_{2u}\,e_{3u}) + \delta_j(v\,e_{1v}\,e_{3v})\right] + \frac{1}{e_{3t}}\,\delta_k(w).$$
The equation could be simplified given the geometry of this basin, but we keep the general form.
[1] NEMO book v4.0.1, p. 22
The first T point along X is a land point, so a 'fill' boundary with fill_value=0
is appropriate; the same holds along Y and Z.
⚠ The vertical term carries a negative sign because the k-axis increases downward in NEMO.
⚠ With a non-linear free surface the e3 scale factors vary in time, so the
thickness-weighted, time-varying e3t, e3u, e3v (rather than the reference
e3t_0, ...) must be used. To stay general we always use the e3 factors from the
output here.
bd = {"boundary": "fill", "fill_value": 0}
div_uv = (
grid.diff(ds.uo * ds.e2u * ds.e3u, "X", **bd) / (ds.e1t * ds.e2t * ds.e3t)
+ grid.diff(ds.vo * ds.e1v * ds.e3v, "Y", **bd) / (ds.e1t * ds.e2t * ds.e3t)
).where(ds.tmask)
div_w = -grid.diff(ds.woce, "Z", **bd) / ds.e3t
div_uvw = div_uv + div_w
print("max |divergence| =", float(np.abs(div_uvw).max()))
max |divergence| = 6.281258392047755e-11
As expected the total divergence is zero up to truncation error.
Vertical velocity¶
In NEMO the vertical velocity is diagnosed from the horizontal divergence. With the linear free surface used here it can be reconstructed offline as
$$w(z) = \int_{\text{bottom}}^{z}\vec{\nabla}_h\cdot\vec{u}\,\mathrm{d}z' = \int_{\text{surf}}^{z}\vec{\nabla}_h\cdot\vec{u}\,\mathrm{d}z' - \int_{\text{surf}}^{\text{bottom}}\vec{\nabla}_h\cdot\vec{u}\,\mathrm{d}z',$$
i.e. in discrete form
$$w(n) = \sum_{0}^{n}\left(\vec{\nabla}_h\cdot\vec{u}(k)\right)e_{3t}(k) - \sum_{0}^{n_{\text{bot}}}\left(\vec{\nabla}_h\cdot\vec{u}(k)\right)e_{3t}(k).$$
We use grid.cumint to integrate the horizontal divergence downward from the surface,
then subtract the full-column integral so that $w=0$ at the bottom.
w = grid.cumint(div_uv, axis="Z", boundary="fill", fill_value=0) # integral from the top
w = w - w.isel({"z_f": -1}) # reference so that w = 0 at the bottom
We compare the reconstructed vertical velocity with the one output by the model, at
the bottom of the uppermost grid cell (z_f=1). Both panels share a common, symmetric
color scale.
wlevel = 1
w_model = ds.woce.where(ds.tmask.isel(z_c=0)).isel(z_f=wlevel)
w_recon = w.where(ds.tmask.isel(z_c=0)).isel(z_f=wlevel)
vmax = float(np.abs(w_model).max())
fig, ax = plt.subplots(1, 2, figsize=(9, 5), sharey=True)
for a, field, title in zip(ax, [w_model, w_recon], ["model output", "reconstructed"]):
field.plot(
x="glamt", y="gphit", ax=a, cmap="RdBu_r", vmin=-vmax, vmax=vmax,
cbar_kwargs={"label": "w [m s$^{-1}$]"},
)
a.set(xlabel="Longitude [°E]", ylabel="Latitude [°N]", title=f"Vertical velocity ({title})")
fig.tight_layout()
plt.show()
The two fields agree closely, as confirmed by the maximum absolute difference:
print("max |w_reconstructed - w_model| =",
float(abs((w - ds.woce).where(ds.tmask.isel(z_c=0))).max()))
max |w_reconstructed - w_model| = 4.189313207592088e-09
Relative vorticity¶
The vertical component of the relative vorticity is a fundamental quantity in ocean circulation theory,
$$\zeta = \frac{\partial v}{\partial x} - \frac{\partial u}{\partial y}.$$
The NEMO discretization is [1]
$$\zeta = \frac{1}{e_{1f}e_{2f}}\left[\delta_{i+1/2}(v\,e_{2v}) - \delta_{j+1/2}(u\,e_{1u})\right].$$
[1] NEMO book v4.0.1, p. 22
zeta = (
1 / (ds.e1f * ds.e2f)
* (grid.diff(ds.vo * ds.e2v, "X", **bd) - grid.diff(ds.uo * ds.e1u, "Y", **bd))
* ds.fmask
)
print("zeta is located on the F point:", zeta.dims)
zeta is located on the F point: ('y_f', 'x_f', 'z_c')
$\zeta$ lives on the F point (z_c, y_f, x_f). To form the vertically integrated
(barotropic) vorticity we need the F-point cell thickness e3f, which we estimate by
interpolating e3u onto the F point:
ds["e3f"] = grid.interp(ds.e3u, "Y", boundary="extend")
zeta_bt = (zeta * ds.e3f).sum(dim="z_c")
vmax = float(np.abs(zeta_bt).max())
fig, ax = plt.subplots(figsize=(4.5, 6))
zeta_bt.plot.contourf(
x="glamf", y="gphif", levels=21, cmap="RdBu_r", vmin=-vmax, vmax=vmax, ax=ax,
cbar_kwargs={"label": r"$\int \zeta\, \mathrm{d}z$ [m s$^{-1}$]"},
)
ax.set(xlabel="Longitude [°E]", ylabel="Latitude [°N]", title="Barotropic relative vorticity")
plt.show()
Barotropic transport streamfunction¶
The barotropic transport streamfunction $\Psi$ is defined via
$$u_{bt} = -\frac{\partial\Psi}{\partial y}, \qquad v_{bt} = \frac{\partial\Psi}{\partial x},$$
$$\Psi(x,y) = \int_0^x\int_{\text{bottom}}^{\text{surface}} v(x,y,z)\,\mathrm{d}z\,\mathrm{d}x.$$
We integrate the meridional velocity over depth with grid.integrate, then
cumulatively integrate along X with grid.cumint, and convert to sverdrups
(1 Sv = $10^6\ \mathrm{m^3\,s^{-1}}$).
ds["psi"] = grid.cumint(grid.integrate(ds.vo, "Z"), "X", **bd) * 1e-6
print("psi is located on:", ds.psi.dims)
fig, ax = plt.subplots(figsize=(4.5, 6))
vmax = float(np.abs(ds.psi).max())
ds.psi.plot.contourf(
x="glamf", y="gphif", levels=25, cmap="RdBu_r", vmin=-vmax, vmax=vmax, ax=ax,
cbar_kwargs={"label": r"$\Psi$ [Sv]"},
)
ax.set(xlabel="Longitude [°E]", ylabel="Latitude [°N]", title="Barotropic streamfunction")
plt.show()
psi is located on: ('y_f', 'x_f')
By construction $\Psi = 0$ at the western boundary.
Kinetic energy¶
Finally we map the surface kinetic energy $\tfrac{1}{2}(u^2 + v^2)$, interpolating both velocity components onto the T (cell-center) point before adding them.
ke = 0.5 * (
grid.interp(ds.uo ** 2, "X", **bd) + grid.interp(ds.vo ** 2, "Y", **bd)
)
fig, ax = plt.subplots(figsize=(4.5, 6))
ke.where(ds.tmask).isel(z_c=0).plot(
x="glamt", y="gphit", robust=True, cmap="viridis", ax=ax,
cbar_kwargs={"label": "KE [m$^2$ s$^{-2}$]"},
)
ax.set(xlabel="Longitude [°E]", ylabel="Latitude [°N]", title="Surface kinetic energy")
plt.show()
Potential vorticity¶
We now compute the Ertel-like potential vorticity on the F point,
$$q = (\zeta + f)\,N^2,$$
where $f$ is the planetary vorticity (output by NEMO as ff_f) and
$N^2 = -\frac{g}{\rho_0}\frac{\partial\rho}{\partial z}$ is the buoyancy frequency.
We first build the potential density $\sigma_0$ from the model's simplified equation of
state, then differentiate it vertically and interpolate $N^2$ onto the F point.
# This NEMO run uses a simplified non-linear equation of state (ln_seos = .true.
# in &nameos). For runs with the full TEOS-10 equation of state, use gsw.sigma0 instead.
rn_a0, rn_b0, rn_lambda1 = 1.655e-1, 7.6554e-1, 0.06
def sigma0(T, S):
"""Potential density anomaly from the NEMO simplified equation of state."""
Ta, Sa = T - 10, S - 35
return 26 - rn_a0 * (1 + 0.5 * rn_lambda1 * Ta) * Ta + rn_b0 * Sa
ds["sigma0"] = sigma0(ds.thetao, ds.so)
g = 9.81 # gravitational acceleration [m s-2]
rho0 = 1026 # reference density [kg m-3]
# Extra minus sign because the Z axis is oriented downward in NEMO
N2 = -g / rho0 * (-grid.derivative(ds.sigma0, "Z", boundary="extend"))
qv = (zeta + ds.ff_f) * grid.interp(N2, ["X", "Y", "Z"], boundary="extend")
print("potential vorticity is located on:", qv.dims)
potential vorticity is located on: ('y_f', 'x_f', 'z_c')
klevel = 5
fig, ax = plt.subplots(figsize=(4.5, 6))
qv.isel(z_c=klevel).plot(
x="glamf", y="gphif", robust=True, cmap="RdBu_r", ax=ax,
cbar_kwargs={"label": "PV [s$^{-3}$]"},
)
ax.set(xlabel="Longitude [°E]", ylabel="Latitude [°N]",
title=f"Potential vorticity (z_c = {klevel})")
plt.show()
The transform function¶
xgcm can remap a variable from one vertical coordinate onto another with
grid.transform (for example, to remap velocities into density coordinates to compute
an overturning streamfunction). We show two uses: a conservative remapping from the
terrain-following levels onto uniform depth levels, and a linear remapping to sample
the temperature a fixed distance below the mixed-layer depth.
Conservative remapping onto uniform depth levels¶
The conservative remapping requires an outer/center axis. The NEMO vertical coordinate is a left/center axis, so we first drop the last (land) T point to turn the Z axis into an outer/center axis, and rebuild the grid.
ds = ds.isel(z_c=slice(None, -1))
grid = xgcm.Grid(ds, metrics=metrics, periodic=False)
print(grid)
<xgcm.Grid> X Axis (not periodic, boundary='fill'): * center x_c --> right * right x_f --> center Z Axis (not periodic, boundary='fill'): * center z_c --> outer * outer z_f --> center Y Axis (not periodic, boundary='fill'): * center y_c --> right * right y_f --> center
The Z axis is now an outer/center axis, as required.
Conservative remapping operates on extensive quantities, so the variable must be
weighted by the cell thickness before transforming. (An extensive quantity scales with
the size of the control volume — e.g. heat content or volume transport — unlike an
intensive one such as temperature.) Here we remap the meridional volume flux
$F_v = v\,e_{3v}$; we cast it to float64 so we can verify conservation precisely
afterwards.
ds["Fv"] = ds.vo.astype(np.float64) * ds.e3v # meridional volume flux per unit width
ds.Fv
<xarray.DataArray 'Fv' (z_c: 35, y_f: 79, x_c: 42)> Size: 929kB
0.0 0.8437 -0.004764 -0.6617 -0.8594 -0.7288 -0.5642 ... 0.0 0.0 0.0 0.0 0.0 0.0
Coordinates:
* z_c (z_c) int64 280B 0 1 2 3 4 5 6 7 8 ... 26 27 28 29 30 31 32 33 34
gdept_1d (z_c) float64 280B ...
* y_f (y_f) float64 632B 0.5 1.5 2.5 3.5 4.5 ... 75.5 76.5 77.5 78.5
* x_c (x_c) int64 336B 0 1 2 3 4 5 6 7 8 ... 33 34 35 36 37 38 39 40 41
glamv (y_f, x_c) float64 27kB ...
gphiv (y_f, x_c) float64 27kB ...Conservative remapping works with the bounds of the target coordinate (the depths of
the W points), not the cell-center depths. We remap onto uniform W-point depths spaced
every h = 40 m, providing the original W-point depths (interpolated onto the V point)
as target_data.
h = 40 # new uniform vertical resolution [m]
ds["Fv_transformed"] = grid.transform(
ds.Fv,
"Z",
target=xr.DataArray(np.arange(0, 4000 + h, h), dims=["depth_uniform"]),
method="conservative",
target_data=grid.interp(ds.gdepw_0, "Y", boundary="extend"),
)
# divide the remapped flux by the (uniform) layer thickness to recover a velocity
ds["vo_transformed"] = ds.Fv_transformed / h
We compare a meridional section of the original meridional velocity (on the native terrain-following levels) with its conservatively remapped counterpart (on uniform depth levels). The two are visually indistinguishable.
ds.coords["gdepv_0"] = grid.interp(ds.gdept_0, "Y", **bd)
ysec = 20
vmin, vmax = -0.02, 0.02
fig, ax = plt.subplots(2, 1, figsize=(8, 7), sharex=True, sharey=True)
ds.vo.where(ds.vmask).isel(y_f=ysec).plot(
x="glamv", y="gdepv_0", yincrease=False, ax=ax[0], cmap="RdBu_r",
vmin=vmin, vmax=vmax, cbar_kwargs={"label": "v [m s$^{-1}$]"},
)
ax[0].set(ylabel="Depth [m]", title="Meridional velocity on native terrain-following levels")
ds.vo_transformed.where(ds.depth_uniform <= ds.gdepv_0.isel(z_c=-1)).isel(y_f=ysec).plot(
x="glamv", y="depth_uniform", yincrease=False, ax=ax[1], cmap="RdBu_r",
vmin=vmin, vmax=vmax, cbar_kwargs={"label": "v [m s$^{-1}$]"},
)
ax[1].set(xlabel="Longitude [°E]", ylabel="Depth [m]",
title="Meridional velocity remapped onto uniform depth levels")
fig.tight_layout()
plt.show()
We verify that the remapping conserves the depth-integrated flux: the difference between the native and remapped column integrals is at the level of floating-point round-off.
residual = grid.integrate(ds.vo.astype(np.float64), "Z") - (ds.vo_transformed * h).sum("depth_uniform")
fig, ax = plt.subplots(figsize=(4.5, 6))
residual.plot(
x="glamv", y="gphiv", robust=True, cmap="RdBu_r", ax=ax,
cbar_kwargs={"label": "transport residual [m$^2$ s$^{-1}$]"},
)
ax.set(xlabel="Longitude [°E]", ylabel="Latitude [°N]",
title="Conservation residual (native minus remapped)")
plt.show()
Temperature 10 m below the mixed-layer depth¶
As a second example we sample the temperature a fixed distance (10 m) below the
mixed-layer depth (MLD). Because the MLD varies horizontally we cannot use a simple
xarray interpolation; instead we use grid.transform with the default linear method.
For linear remapping the target_data are the cell-center depths, here shifted so
that the depth measured relative to "10 m below the MLD" is gdept_0 - mldr10_1.
theta_below_MLD = grid.transform(
ds.thetao.where(ds.tmask),
"Z",
target=xr.DataArray([10.0], dims=["depth_below_MLD"]),
target_data=ds.gdept_0 - ds.mldr10_1,
).squeeze("depth_below_MLD")
fig, ax = plt.subplots(figsize=(4.5, 6))
theta_below_MLD.plot(
x="glamt", y="gphit", robust=True, cmap="magma", ax=ax,
cbar_kwargs={"label": "Temperature [°C]"},
)
ax.set(xlabel="Longitude [°E]", ylabel="Latitude [°N]",
title="Temperature 10 m below the mixed-layer depth")
plt.show()
