🖥️ Sigma Coordinate Vertical Grids#
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation
import parcels
from parcels._datasets.structured.generated import (
simple_UV_dataset,
simple_UV_sigma_dataset,
)
from parcels._datasets.unstructured.generated import sigma_coordinate_lattice_dataset
plt.rcParams["animation.html"] = "jshtml"
Building the Dataset#
For this tutorial, we construct our dataset using the parcels built in simple_UV_sigma_dataset, which makes a structured horizontally gridded dataset that uses a linear vertical \(\sigma\) grid. Its horizontal grid is a flat square that spans from -1000 km to 1000 km in both \(x\) and \(y\). Here, we assume \(\frac{\partial \sigma}{\partial z}\) to be constant, so for any given z coordinate, the equivalent \(\sigma\) coordinate is a dimensionless fraction of column height, i.e. \(\sigma = \frac{z + \eta}{H +\eta}\), where \(H\) is the depth of the water column, and \(\eta\) is the sea surface height. It is common in \(\sigma\) coordinate models to use more advanced vertical grids for which \(\frac{\partial \sigma}{\partial z}\) is not constant, this allows the concentration of vertical layers near the surface. Nonetheless, our simple example allows for a vertical coordinate system that adapts to changes in bathymetry and variations in sea surface height. We will demonstrate this shortly but begin by constructing the dataset itself. At the end of the tutorial, we show the same setup on an unstructured grid.
n_days = 10
n_snapshots = 24 * n_days + 1 # hourly snapshots from hour 0 to the end of day 10
n_vertical_layers = 10
ny, nx = (
5,
81,
) # the number of grid nodes in y and x, so that the cells are 25km wide in x
lon = simple_UV_dataset(
n_times=n_snapshots, nz=n_vertical_layers + 1, ny=ny, nx=nx, mesh_str="flat"
)["lon"].values
dx = lon[1] - lon[0]
x_centres = lon - 0.5 * dx
# we construct a depth profile that is 100m deep in the west and logarithmically increases to 1000m depth in the east
domain_length = lon[-1] - lon[0]
bottom_depth = 100 + 900 * np.log10(1 + 9 * (x_centres - lon[0]) / domain_length)
bottom_depth = np.broadcast_to(bottom_depth, (ny, nx))
# for illustrative purposes we create a 50m surface wave with a 1000km wavelength that propagates east with time
t = np.arange(n_snapshots) * 3600.0
wavelength, wave_period = 1000e3, 2 * 86400.0
amplitude = 50
wave = amplitude * np.sin(
2.0
* np.pi
* (x_centres[np.newaxis, :] / wavelength - t[:, np.newaxis] / wave_period)
)
wave = np.broadcast_to(wave[:, np.newaxis, :], (n_snapshots, ny, nx))
# construct the dataset itself
ds = simple_UV_sigma_dataset(
n_times=n_snapshots,
nz=n_vertical_layers + 1,
ny=ny,
nx=nx,
bottom_depth=bottom_depth,
eta=wave,
)
ds
<xarray.Dataset> Size: 34MB
Dimensions: (time: 241, ZF: 11, YG: 5, XG: 81, YC: 5, XC: 81)
Coordinates:
* time (time) datetime64[ns] 2kB 2000-01-01 ... 2000-01-11
* ZF (ZF) float64 88B 0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0
* YG (YG) int64 40B 0 1 2 3 4
lat (YG) float64 40B -1e+06 -5e+05 0.0 5e+05 1e+06
* XG (XG) int64 648B 0 1 2 3 4 5 6 7 8 9 ... 72 73 74 75 76 77 78 79 80
lon (XG) float64 648B -1e+06 -9.75e+05 -9.5e+05 ... 9.75e+05 1e+06
* YC (YC) float64 40B 0.5 1.5 2.5 3.5 4.5
* XC (XC) float64 648B 0.5 1.5 2.5 3.5 4.5 ... 76.5 77.5 78.5 79.5 80.5
depth (time, ZF, YC, XC) float64 9MB 3.923 -3.923 -11.67 ... 993.3 997.8
Data variables:
U (time, ZF, YG, XG) float64 9MB 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
V (time, ZF, YG, XG) float64 9MB 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
grid int64 8B 0
W (time, ZF, YG, XG) float64 9MB 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
Attributes:
Conventions: SGRIDWe can now compare the evolution of the vertical grid in \(\sigma\) space to the vertical grid in z space. In z space, the vertical grid spacing is much smaller in the shallow west than in the deep east, this demonstrates the strength of a terrain following vertical grid. We can also see that the vertical grid has time dependency in z space. Currently, for structured grids, Parcels uses piecewise constant horizontal interpolation of z values for time-varying z.
time = ds["time"].values.astype("datetime64[s]")
z_moving = ds["depth"].values # interface depths, shape (time, depth, YC, XC)
sigma = np.linspace(0.0, 1.0, n_vertical_layers + 1)
centre_row = ny // 2 # the row of cells in the middle of the domain
xs = x_centres / 1e3
fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)
ax_sigma.plot(xs, np.broadcast_to(sigma, (xs.size, sigma.size)), "k", linewidth=0.8)
ax_sigma.set_ylabel(r"$\sigma$")
ax_sigma.set_title(r"Layer interfaces in $\sigma$ space")
z_lines = ax_z.plot(
xs, z_moving[0][:, centre_row].T, "k", linewidth=0.8, drawstyle="steps-mid"
)
ax_z.set_ylabel("depth [m]")
z_title = ax_z.set_title(f"Layer interfaces in z space at {time[0]}")
for ax in (ax_sigma, ax_z):
ax.invert_yaxis()
ax.set_xlabel("x [km]")
plt.tight_layout()
def animate(snapshot):
for line, interface_depths in zip(
z_lines, z_moving[snapshot][:, centre_row], strict=True
):
line.set_ydata(interface_depths)
z_title.set_text(f"Layer interfaces in z space at {time[snapshot]}")
anim = FuncAnimation(fig, animate, frames=range(0, time.size, 6), interval=300)
plt.close()
anim
Building the FieldSet#
We prescribe a uniform velocity of \(u_0 = 1\) m/s in the \(x\) direction, with \(v = w = 0\). All three components sit on the grid nodes at the layer interfaces (depth), so Parcels interpolates them linearly in all three directions.
To advect the particles on a sigma grid, for which the z grid varies in time, we use a 4D depth coordinate with dimensions of (time, depth, YC, XC). For each particle, Parcels takes the column of the cell that the particle is in, interpolates it linearly in time between the two surrounding snapshots, and then searches that column for the particle’s layer. We pass mesh="flat" because the coordinates are Cartesian meters.
u0 = 1.0
ds["U"].values[:] = u0
ds
<xarray.Dataset> Size: 34MB
Dimensions: (time: 241, ZF: 11, YG: 5, XG: 81, YC: 5, XC: 81)
Coordinates:
* time (time) datetime64[ns] 2kB 2000-01-01 ... 2000-01-11
* ZF (ZF) float64 88B 0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0
* YG (YG) int64 40B 0 1 2 3 4
lat (YG) float64 40B -1e+06 -5e+05 0.0 5e+05 1e+06
* XG (XG) int64 648B 0 1 2 3 4 5 6 7 8 9 ... 72 73 74 75 76 77 78 79 80
lon (XG) float64 648B -1e+06 -9.75e+05 -9.5e+05 ... 9.75e+05 1e+06
* YC (YC) float64 40B 0.5 1.5 2.5 3.5 4.5
* XC (XC) float64 648B 0.5 1.5 2.5 3.5 4.5 ... 76.5 77.5 78.5 79.5 80.5
depth (time, ZF, YC, XC) float64 9MB 3.923 -3.923 -11.67 ... 993.3 997.8
Data variables:
U (time, ZF, YG, XG) float64 9MB 1.0 1.0 1.0 1.0 ... 1.0 1.0 1.0 1.0
V (time, ZF, YG, XG) float64 9MB 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
grid int64 8B 0
W (time, ZF, YG, XG) float64 9MB 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
Attributes:
Conventions: SGRIDfieldset = parcels.FieldSet.from_sgrid_conventions(ds)
fieldset.describe()
| Name | Type | Grid number | Interp method / value | Parcels backend |
|:-------|:------------|--------------:|:------------------------|:------------------|
| U | Field | 0 | XLinear(...) | NumPy |
| V | Field | 0 | XLinear(...) | NumPy |
| W | Field | 0 | XLinear(...) | NumPy |
| UV | VectorField | 0 | XLinear_Velocity(...) | - |
| UVW | VectorField | 0 | XLinear_Velocity(...) | - |
mesh: FlatMesh()
time interval: (np.datetime64('2000-01-01T00:00:00.000000000'), np.datetime64('2000-01-11T00:00:00.000000000'))
/tmp/ipykernel_4182/1955122603.py:1: FieldSetWarning: Time-varying (4D) z coordinates are experimental and may cause significant memory overhead that leads to OOM errors. This z coordinate has sizes {'time': 241, 'ZF': 11, 'YC': 5, 'XC': 81} (0.00859 GB). Assumptions: z has dims ('time', 'ZF', 'YC', 'XC'), i.e. it is defined at the layer interfaces on the cell centres, and is strictly increasing along 'ZF'; each particle uses the z column at the centre of its grid cell, so z is constant within a cell and jumps between cells, it is linearly interpolated in time between z snapshots.
fieldset = parcels.FieldSet.from_sgrid_conventions(ds)
Release the particles#
We release 10 particles at the first snapshot near the shallow western end of the domain, at \(x = -890\) km and \(y = 0\), at depths from 60 m to 240 m. This puts the shallowest particle just below the deepest wave trough (50 m), so every particle will remain in the water column throughout the simulation.
n_particles = 10
x0 = np.full(n_particles, -890e3)
y0 = np.zeros(n_particles)
z0 = np.linspace(60.0, 240.0, n_particles) # depths in m
pset = parcels.ParticleSet(
fieldset=fieldset,
pclass=parcels.Particle,
t=np.full(n_particles, time[0]),
z=z0,
y=y0,
x=x0,
)
H = z_moving[0, -1, centre_row]
eta = -z_moving[:, 0, centre_row]
cell = np.searchsorted(
lon, x0
) # index of the cell centre whose z column Parcels uses for each particle
sigma0 = (z0 + eta[0, cell]) / (
H[cell] + eta[0, cell]
) # initial sigma coordinate of each particle
colors = [f"C{p}" for p in range(n_particles)]
fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)
ax_sigma.plot(xs, np.broadcast_to(sigma, (xs.size, sigma.size)), "k", linewidth=0.8)
ax_sigma.scatter(x0 / 1e3, sigma0, color=colors, zorder=3)
ax_sigma.set_ylabel(r"$\sigma$")
ax_sigma.set_title(r"Initial particle positions in $\sigma$ space")
ax_z.plot(xs, z_moving[0][:, centre_row].T, "k", linewidth=0.8, drawstyle="steps-mid")
ax_z.scatter(x0 / 1e3, z0, color=colors, zorder=3)
ax_z.set_ylabel("depth [m]")
ax_z.set_title(f"Initial particle positions in z space at {time[0]}")
for ax in (ax_sigma, ax_z):
ax.invert_yaxis()
ax.set_xlabel("x [km]")
plt.tight_layout()
plt.show()
Advect the particles#
We advect the particles for 10 days using the 2D second-order Runge-Kutta kernel parcels.kernels.AdvectionRK2() with a 30-minute time step, and write their positions every 6 hours. We chose the 2D version of RK2 since the vertical velocity is zero across the domain.
outputfile = parcels.ParticleFile(
path="sigma_particles.parquet", outputdt=np.timedelta64(6, "h"), mode="w"
)
pset.execute(
[parcels.kernels.AdvectionRK2],
runtime=np.timedelta64(n_days, "D"),
dt=np.timedelta64(30, "m"),
output_file=outputfile,
)
df = parcels.read_particlefile("sigma_particles.parquet").sort(["t", "particle_id"])
output_times = df["t"].unique().sort()
n_frames = len(output_times)
px = df["x"].to_numpy().reshape(n_frames, n_particles)
pz = df["z"].to_numpy().reshape(n_frames, n_particles)
# calculate the sigma coordinate of every particle at every output time, using the z column of its cell
snapshots = np.searchsorted(time, output_times.to_numpy(), side="right") - 1
cells = np.searchsorted(lon, px)
eta_p = eta[snapshots[:, np.newaxis], cells]
psigma = (pz + eta_p) / (H[cells] + eta_p)
colors = [f"C{p}" for p in range(n_particles)]
fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)
ax_sigma.plot(xs, np.broadcast_to(sigma, (xs.size, sigma.size)), "k", linewidth=0.8)
z_lines = ax_z.plot(
xs, z_moving[0][:, centre_row].T, "k", linewidth=0.8, drawstyle="steps-mid"
)
sigma_trails = [ax_sigma.plot([], [], color=c, linewidth=1.5)[0] for c in colors]
z_trails = [ax_z.plot([], [], color=c, linewidth=1.5)[0] for c in colors]
sigma_dots = ax_sigma.scatter(px[0] / 1e3, psigma[0], c=colors, zorder=3)
z_dots = ax_z.scatter(px[0] / 1e3, pz[0], c=colors, zorder=3)
ax_sigma.set_ylabel(r"$\sigma$")
ax_sigma.set_title(r"Particles in $\sigma$ space")
ax_z.set_ylabel("depth [m]")
z_title = ax_z.set_title("")
for ax in (ax_sigma, ax_z):
ax.invert_yaxis()
ax.set_xlabel("x [km]")
plt.tight_layout()
def animate(frame):
for p in range(n_particles):
sigma_trails[p].set_data(px[: frame + 1, p] / 1e3, psigma[: frame + 1, p])
z_trails[p].set_data(px[: frame + 1, p] / 1e3, pz[: frame + 1, p])
sigma_dots.set_offsets(np.c_[px[frame] / 1e3, psigma[frame]])
z_dots.set_offsets(np.c_[px[frame] / 1e3, pz[frame]])
for line, interface_depths in zip(
z_lines, z_moving[snapshots[frame]][:, centre_row], strict=True
):
line.set_ydata(interface_depths)
z_title.set_text(f"Particles in z space at {output_times[frame]}")
animate(0)
anim = FuncAnimation(fig, animate, frames=n_frames, interval=300)
plt.close()
anim
We can see that, as expected, the particles move horizontally in x and cross the terrain following depth levels. However, in z space, they maintain their exact height as expected.
Sigma coordinates on unstructured grids#
The same approach works on unstructured grids. On an unstructured grid, z is a 3D zf coordinate with dimensions (time, zf, n_node), defined on the corners (nodes) of the triangular faces, rather than a 4D depth coordinate with dimensions (time, depth, YC, XC) defined on the cell centres. For each particle, Parcels interpolates the z columns of the three corners of its triangle barycentrically in space and linearly in time, so on unstructured grids the z grid varies continuously in space instead of being constant across each cell.
Below we build the same basin on a triangulated lattice with sigma_coordinate_lattice_dataset, evaluating the bathymetry and sea surface height at the lattice nodes, and compare the two z coordinates.
n_lattice_nodes = (
21 # the number of lattice nodes in x and y, so that the nodes are 100km apart
)
x_nodes = np.linspace(-1000e3, 1000e3, n_lattice_nodes)
# the same bathymetry and surface wave as above, evaluated at the lattice nodes, with shapes (x, y) and (time, x, y)
bottom_depth_nodes = 100 + 900 * np.log10(
1 + 9 * (x_nodes - x_nodes[0]) / domain_length
)
bottom_depth_nodes = np.broadcast_to(
bottom_depth_nodes[:, np.newaxis], (n_lattice_nodes, n_lattice_nodes)
)
wave_nodes = amplitude * np.sin(
2.0 * np.pi * (x_nodes[np.newaxis, :] / wavelength - t[:, np.newaxis] / wave_period)
)
wave_nodes = np.broadcast_to(
wave_nodes[:, :, np.newaxis], (n_snapshots, n_lattice_nodes, n_lattice_nodes)
)
ux_ds = sigma_coordinate_lattice_dataset(
n_lattice_nodes,
(-1000e3, 1000e3),
(-1000e3, 1000e3),
n_vertical_layers + 1,
bottom_depth_nodes,
eta=wave_nodes,
)
print("structured z: ", ds["depth"].dims, ds["depth"].shape)
print("unstructured z:", ux_ds["zf"].dims, ux_ds["zf"].shape)
structured z: ('time', 'ZF', 'YC', 'XC') (241, 11, 5, 81)
unstructured z: ('time', 'zf', 'n_node') (241, 11, 441)
ux_ds["U"].values[:] = u0
ux_fieldset = parcels.FieldSet.from_ugrid_conventions(ux_ds)
ux_fieldset.describe()
| Name | Type | Grid number | Interp method / value | Parcels backend |
|:-------|:------------|--------------:|:--------------------------|:------------------|
| U | Field | 0 | UxLinearNodeLinearZF(...) | NumPy |
| V | Field | 0 | UxLinearNodeLinearZF(...) | NumPy |
| W | Field | 0 | UxLinearNodeLinearZF(...) | NumPy |
| UV | VectorField | 0 | Ux_Velocity(...) | - |
| UVW | VectorField | 0 | Ux_Velocity(...) | - |
mesh: FlatMesh()
time interval: (np.datetime64('2000-01-01T00:00:00.000000000'), np.datetime64('2000-01-11T00:00:00.000000000'))
/tmp/ipykernel_4182/1845485040.py:2: FieldSetWarning: Time-varying (3D) z coordinates are experimental and may cause significant memory overhead that leads to OOM errors. This z coordinate has sizes {'time': 241, 'zf': 11, 'n_node': 441} (0.00935 GB). Assumptions: z is defined at the layer interfaces ('zf') on the mesh nodes ('n_node') and is strictly increasing along 'zf'; each particle's z column is interpolated barycentrically from its face's nodes and linearly in time between z snapshots.
ux_fieldset = parcels.FieldSet.from_ugrid_conventions(ux_ds)
We release the same 10 particles at the same positions on the unstructured grid, and advect them with the same kernel, time step and output frequency as before.
ux_pset = parcels.ParticleSet(
fieldset=ux_fieldset,
pclass=parcels.Particle,
t=np.full(n_particles, time[0]),
z=z0,
y=y0,
x=x0,
)
ux_outputfile = parcels.ParticleFile(
path="ux_sigma_particles.parquet", outputdt=np.timedelta64(6, "h"), mode="w"
)
ux_pset.execute(
[parcels.kernels.AdvectionRK2],
runtime=np.timedelta64(n_days, "D"),
dt=np.timedelta64(30, "m"),
output_file=ux_outputfile,
)
ux_df = parcels.read_particlefile("ux_sigma_particles.parquet").sort(
["t", "particle_id"]
)
ux_px = ux_df["x"].to_numpy().reshape(n_frames, n_particles)
ux_pz = ux_df["z"].to_numpy().reshape(n_frames, n_particles)
# the particles move along the row of lattice nodes at y = 0, i.e. along triangle edges, so the z column that
# Parcels uses for each particle is the linear interpolation in x of the z columns at the nodes of that row
row = np.isclose(ux_ds.uxgrid.node_lat.values, 0.0)
z_nodes = ux_ds["zf"].values[
:, :, row
] # interface depths along y = 0, shape (time, zf, x)
H_nodes, eta_nodes = z_nodes[0, -1], -z_nodes[:, 0]
ux_psigma = np.empty_like(ux_pz)
for frame, snapshot in enumerate(snapshots):
ux_eta_p = np.interp(ux_px[frame], x_nodes, eta_nodes[snapshot])
ux_psigma[frame] = (ux_pz[frame] + ux_eta_p) / (
np.interp(ux_px[frame], x_nodes, H_nodes) + ux_eta_p
)
fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)
ax_sigma.plot(
x_nodes / 1e3,
np.broadcast_to(sigma, (x_nodes.size, sigma.size)),
"k",
linewidth=0.8,
)
z_lines = ax_z.plot(x_nodes / 1e3, z_nodes[0].T, "k", linewidth=0.8)
sigma_trails = [ax_sigma.plot([], [], color=c, linewidth=1.5)[0] for c in colors]
z_trails = [ax_z.plot([], [], color=c, linewidth=1.5)[0] for c in colors]
sigma_dots = ax_sigma.scatter(ux_px[0] / 1e3, ux_psigma[0], c=colors, zorder=3)
z_dots = ax_z.scatter(ux_px[0] / 1e3, ux_pz[0], c=colors, zorder=3)
ax_sigma.set_ylabel(r"$\sigma$")
ax_sigma.set_title(r"Particles in $\sigma$ space on the unstructured grid")
ax_z.set_ylabel("depth [m]")
z_title = ax_z.set_title("")
for ax in (ax_sigma, ax_z):
ax.invert_yaxis()
ax.set_xlabel("x [km]")
plt.tight_layout()
def animate(frame):
for p in range(n_particles):
sigma_trails[p].set_data(ux_px[: frame + 1, p] / 1e3, ux_psigma[: frame + 1, p])
z_trails[p].set_data(ux_px[: frame + 1, p] / 1e3, ux_pz[: frame + 1, p])
sigma_dots.set_offsets(np.c_[ux_px[frame] / 1e3, ux_psigma[frame]])
z_dots.set_offsets(np.c_[ux_px[frame] / 1e3, ux_pz[frame]])
for line, interface_depths in zip(z_lines, z_nodes[snapshots[frame]]):
line.set_ydata(interface_depths)
z_title.set_text(f"Particles in z space at {output_times[frame]}")
animate(0)
anim = FuncAnimation(fig, animate, frames=n_frames, interval=300)
plt.close()
anim
As on the structured grid, the particles keep their exact height in z space while crossing the terrain following depth levels. On the unstructured grid, however, the layer interfaces vary continuously between the nodes rather than being constant across each grid cell.