From ba7fe37af406c32e399033622ce1d36babe5225c Mon Sep 17 00:00:00 2001
From: Wyatt Sieminski <wyatt@fluidnumerics.com>
Date: Mon, 5 Oct 2026 19:42:14 +0000
Subject: [PATCH 01/22] Initial 3d z coordinate implementation for unstructured
 grids

---
 ...orial_unstructured_sigma_coordinates.ipynb | 369 ++++++++++++++++++
 docs/user_guide/index.md                      |   1 +
 src/parcels/_core/basegrid.py                 |   4 +-
 src/parcels/_core/field.py                    |   5 +-
 src/parcels/_core/index_search.py             |  39 ++
 src/parcels/_core/model.py                    |   3 +
 src/parcels/_core/particleset.py              |   6 +
 src/parcels/_core/uxgrid.py                   |  45 ++-
 src/parcels/_core/xgrid.py                    |   2 +-
 .../_datasets/unstructured/generated.py       |  66 ++++
 src/parcels/interpolators/_uxinterpolators.py |  12 +-
 tests/test_uxadvection.py                     |  22 ++
 tests/test_uxarray_fieldset.py                |  83 ++++
 tests/test_uxgrid.py                          |  10 +
 tests/test_xgrid.py                           |  28 +-
 15 files changed, 668 insertions(+), 27 deletions(-)
 create mode 100644 docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb

diff --git a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
new file mode 100644
index 0000000000..ee32d8ede2
--- /dev/null
+++ b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
@@ -0,0 +1,369 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 🖥️ Sigma Coordinate Vertical Grids on Unstructured Grids"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import polars as pl\n",
+    "from matplotlib.animation import FuncAnimation\n",
+    "\n",
+    "import parcels\n",
+    "from parcels._datasets.unstructured.generated import sigma_coordinate_lattice_dataset\n",
+    "\n",
+    "# for interactive display of animations\n",
+    "plt.rcParams[\"animation.html\"] = \"jshtml\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Build the channel mesh\n",
+    "\n",
+    "We begin by constructing a 2D mesh that we will later extend into 3 dimensions."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "nx = 21  # nodes along and across the channel\n",
+    "n_layers = 10\n",
+    "x_nodes, _ = np.meshgrid(np.linspace(0.0, 40e3, nx), np.linspace(0.0, 10e3, nx), indexing=\"ij\")\n",
+    "bottom_depth = 10.0 + 90.0 * np.log10(1.0 + 9.0 * x_nodes / 40e3)  # (x, y) array, in meters\n",
+    "\n",
+    "ds_static = sigma_coordinate_lattice_dataset(nx, (0.0, 40e3), (0.0, 10e3), n_layers + 1, bottom_depth)\n",
+    "node_x, node_y = ds_static.uxgrid.node_lon.values, ds_static.uxgrid.node_lat.values\n",
+    "face_nodes = ds_static.uxgrid.face_node_connectivity.values\n",
+    "\n",
+    "print(f\"n_node={ds_static.uxgrid.n_node}, n_face={ds_static.uxgrid.n_face}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(figsize=(10, 3))\n",
+    "ax.triplot(node_x / 1e3, node_y / 1e3, face_nodes, color=\"k\", linewidth=0.5)\n",
+    "ax.set_aspect(\"equal\")\n",
+    "ax.set_xlabel(\"x [km]\")\n",
+    "ax.set_ylabel(\"y [km]\")\n",
+    "ax.set_title(\"Channel mesh\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Build the vertical grid\n",
+    "\n",
+    "Next we give the channel a terrain-following (sigma) vertical coordinate. The bottom depth $H$ increases logarithmically along the channel, from 10 m at the shallow left end to 100 m at the deep right end, so the bottom slopes steeply near the shallow end and flattens towards the deep end, typical of a coastal shelf. Following Parcels' convention that depth is positive downward, we define $\\sigma = \\frac{z + \\eta}{H + \\eta}$, so $\\sigma = 0$ at the surface and $\\sigma = 1$ at the bottom. Initially, we assume $\\eta$ to be zero. We place 11 layer interfaces (10 layers) equally spaced in $\\sigma$, so every node has its own column of interface depths $z = \\sigma H$."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sigma = np.linspace(0.0, 1.0, n_layers + 1)\n",
+    "z = ds_static[\"zf\"].values[0]  # interface depths at every node, shape (zf, n_node)\n",
+    "H = z[-1]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "centre_line = np.isclose(node_y, 5e3) \n",
+    "xs = node_x[centre_line] / 1e3\n",
+    "\n",
+    "fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)\n",
+    "ax_sigma.plot(xs, np.broadcast_to(sigma, (xs.size, sigma.size)), \"k\", linewidth=0.8)\n",
+    "ax_sigma.set_ylabel(r\"$\\sigma$\")\n",
+    "ax_sigma.set_title(r\"Layer interfaces in $\\sigma$ space\")\n",
+    "\n",
+    "ax_z.plot(xs, z[:, centre_line].T, \"k\", linewidth=0.8)\n",
+    "ax_z.set_ylabel(\"depth [m]\")\n",
+    "ax_z.set_title(\"Layer interfaces in z space\")\n",
+    "\n",
+    "for ax in (ax_sigma, ax_z):\n",
+    "    ax.invert_yaxis()\n",
+    "    ax.set_xlabel(\"x [km]\")\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Temporal Variation in $\\eta$\n",
+    "\n",
+    "Next, we add a fluctuating sea surface height field to the vertical grid. For illustrative purposes, we use a large amplitude sea surface height wave that's independent of the velocity field. In sigma space, the vertical grid remains exactly the same, since it accounts for $\\eta$ when calculating the fractional vertical coordinate. In z space however, the change in $\\eta$ causes fluctuations in the grid itself."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "t = np.arange(13) * 3600.0\n",
+    "\n",
+    "# for illustrative purposes we put a 2m surface wave that propogates with time\n",
+    "wavelength, wave_period = 20e3, 6 * 3600.0\n",
+    "wave = 2 * np.sin(2.0 * np.pi * (x_nodes[np.newaxis] / wavelength - t[:, np.newaxis, np.newaxis] / wave_period))\n",
+    "\n",
+    "ds = sigma_coordinate_lattice_dataset(nx, (0.0, 40e3), (0.0, 10e3), n_layers + 1, bottom_depth, eta=wave)\n",
+    "time = ds[\"time\"].values\n",
+    "z_moving = ds[\"zf\"].values  # (time, zf, n_node)\n",
+    "eta = -z_moving[:, 0, :]  # the surface interface sits at depth -eta"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)\n",
+    "ax_sigma.plot(xs, np.broadcast_to(sigma, (xs.size, sigma.size)), \"k\", linewidth=0.8)\n",
+    "ax_sigma.set_ylabel(r\"$\\sigma$\")\n",
+    "ax_sigma.set_title(r\"Layer interfaces in $\\sigma$ space\")\n",
+    "\n",
+    "z_lines = ax_z.plot(xs, z_moving[0][:, centre_line].T, \"k\", linewidth=0.8)\n",
+    "ax_z.set_ylabel(\"depth [m]\")\n",
+    "z_title = ax_z.set_title(f\"Layer interfaces in z space at {time[0]}\")\n",
+    "\n",
+    "for ax in (ax_sigma, ax_z):\n",
+    "    ax.invert_yaxis()\n",
+    "    ax.set_xlabel(\"x [km]\")\n",
+    "plt.tight_layout()\n",
+    "\n",
+    "def animate(snapshot):\n",
+    "    for line, interface_depths in zip(z_lines, z_moving[snapshot][:, centre_line]):\n",
+    "        line.set_ydata(interface_depths)\n",
+    "    z_title.set_text(f\"Layer interfaces in z space at {time[snapshot]}\")\n",
+    "\n",
+    "anim = FuncAnimation(fig, animate, frames=time.size, interval=500)\n",
+    "plt.close()\n",
+    "anim"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Building the dataset and `FieldSet`\n",
+    "\n",
+    "We prescribe a uniform along channel velocity of $u_0 = 0.5$ m/s, with $v = w = 0$. All three components sit on the mesh nodes at the layer interfaces (`zf`), so Parcels interpolates them barycentrically in the horizontal and linearly in the vertical.\n",
+    "\n",
+    "To advect the particles on sigma grid, for which the z grid varies in time, we use a 3D `zf` coordinate with dimensions of `(time, zf, n_node)`. Parcels then uses it to find each particle's own column at the current time, before searching that column for the particle's layer. `parcels.FieldSet.from_ugrid_conventions()` also expects a `zc` dimension for the layer centres, so we add the layer-centre depths as well. We pass `mesh=\"flat\"` because the coordinates are Cartesian meters."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "u0 = 0.5\n",
+    "ds[\"U\"].values[:] = u0\n",
+    "ds"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fieldset = parcels.FieldSet.from_ugrid_conventions(ds, mesh=\"flat\")\n",
+    "fieldset.describe()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Release the particles\n",
+    "\n",
+    "We release 10 particles at the first snapshot near the shallow end of the channel, at $x = 2$ km on the centre line, at depths from 3 m to 22 m. The water column there is about 22.5 m deep, and the shallowest particle starts below the deepest wave trough (2 m), so every particle will remain in the water column throughout the simulation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "n_particles = 10\n",
+    "x0 = np.full(n_particles, 2e3)\n",
+    "y0 = np.full(n_particles, 5e3)\n",
+    "z0 = np.linspace(3.0, 22.0, n_particles)  # depths in m\n",
+    "\n",
+    "pset = parcels.ParticleSet(\n",
+    "    fieldset=fieldset, pclass=parcels.Particle, t=np.full(n_particles, time[0]), z=z0, y=y0, x=x0\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x_line = node_x[centre_line]\n",
+    "H0 = np.interp(x0, x_line, H[centre_line])\n",
+    "eta0 = np.interp(x0, x_line, eta[0, centre_line])\n",
+    "sigma0 = (z0 + eta0) / (H0 + eta0)  # initial sigma coordinate of each particle\n",
+    "\n",
+    "colors = [f\"C{p}\" for p in range(n_particles)]\n",
+    "\n",
+    "fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)\n",
+    "ax_sigma.plot(xs, np.broadcast_to(sigma, (xs.size, sigma.size)), \"k\", linewidth=0.8)\n",
+    "ax_sigma.scatter(x0 / 1e3, sigma0, color=colors, zorder=3)\n",
+    "ax_sigma.set_ylabel(r\"$\\sigma$\")\n",
+    "ax_sigma.set_title(r\"Initial particle positions in $\\sigma$ space\")\n",
+    "\n",
+    "ax_z.plot(xs, z_moving[0][:, centre_line].T, \"k\", linewidth=0.8)\n",
+    "ax_z.scatter(x0 / 1e3, z0, color=colors, zorder=3)\n",
+    "ax_z.set_ylabel(\"depth [m]\")\n",
+    "ax_z.set_title(f\"Initial particle positions in z space at {time[0]}\")\n",
+    "\n",
+    "for ax in (ax_sigma, ax_z):\n",
+    "    ax.invert_yaxis()\n",
+    "    ax.set_xlabel(\"x [km]\")\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Advect the particles\n",
+    "\n",
+    "We advect the particles for 12 hours using the 3D fourth-order Runge-Kutta kernel `parcels.kernels.AdvectionRK4_3D` with a 5-minute time step, and write their positions every 30 minutes. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "outputfile = parcels.ParticleFile(path=\"sigma_particles.parquet\", outputdt=np.timedelta64(30, \"m\"), mode=\"w\")\n",
+    "\n",
+    "pset.execute(\n",
+    "    [parcels.kernels.AdvectionRK4_3D],\n",
+    "    runtime=np.timedelta64(12, \"h\"),\n",
+    "    dt=np.timedelta64(5, \"m\"),\n",
+    "    output_file=outputfile,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df = parcels.read_particlefile(\"sigma_particles.parquet\").sort([\"t\", \"particle_id\"])\n",
+    "output_times = df[\"t\"].unique().sort()\n",
+    "n_frames = len(output_times)\n",
+    "px = df[\"x\"].to_numpy().reshape(n_frames, n_particles)\n",
+    "pz = df[\"z\"].to_numpy().reshape(n_frames, n_particles)\n",
+    "\n",
+    "# sigma coordinate of every particle at every output time, on the latest grid snapshot at or before that time\n",
+    "snapshots = np.searchsorted(time, output_times.to_numpy(), side=\"right\") - 1\n",
+    "psigma = np.empty_like(pz)\n",
+    "for frame, snapshot in enumerate(snapshots):\n",
+    "    eta_p = np.interp(px[frame], x_line, eta[snapshot, centre_line])\n",
+    "    psigma[frame] = (pz[frame] + eta_p) / (np.interp(px[frame], x_line, H[centre_line]) + eta_p)\n",
+    "\n",
+    "colors = [f\"C{p}\" for p in range(n_particles)]\n",
+    "\n",
+    "fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)\n",
+    "ax_sigma.plot(xs, np.broadcast_to(sigma, (xs.size, sigma.size)), \"k\", linewidth=0.8)\n",
+    "z_lines = ax_z.plot(xs, z_moving[0][:, centre_line].T, \"k\", linewidth=0.8)\n",
+    "sigma_trails = [ax_sigma.plot([], [], color=c, linewidth=1.5)[0] for c in colors]\n",
+    "z_trails = [ax_z.plot([], [], color=c, linewidth=1.5)[0] for c in colors]\n",
+    "sigma_dots = ax_sigma.scatter(px[0] / 1e3, psigma[0], c=colors, zorder=3)\n",
+    "z_dots = ax_z.scatter(px[0] / 1e3, pz[0], c=colors, zorder=3)\n",
+    "\n",
+    "ax_sigma.set_ylabel(r\"$\\sigma$\")\n",
+    "ax_sigma.set_title(r\"Particles in $\\sigma$ space\")\n",
+    "ax_z.set_ylabel(\"depth [m]\")\n",
+    "z_title = ax_z.set_title(\"\")\n",
+    "for ax in (ax_sigma, ax_z):\n",
+    "    ax.invert_yaxis()\n",
+    "    ax.set_xlabel(\"x [km]\")\n",
+    "plt.tight_layout()\n",
+    "\n",
+    "def animate(frame):\n",
+    "    for p in range(n_particles):\n",
+    "        sigma_trails[p].set_data(px[: frame + 1, p] / 1e3, psigma[: frame + 1, p])\n",
+    "        z_trails[p].set_data(px[: frame + 1, p] / 1e3, pz[: frame + 1, p])\n",
+    "    sigma_dots.set_offsets(np.c_[px[frame] / 1e3, psigma[frame]])\n",
+    "    z_dots.set_offsets(np.c_[px[frame] / 1e3, pz[frame]])\n",
+    "    for line, interface_depths in zip(z_lines, z_moving[snapshots[frame]][:, centre_line]):\n",
+    "        line.set_ydata(interface_depths)\n",
+    "    z_title.set_text(f\"Particles in z space at {output_times[frame]}\")\n",
+    "\n",
+    "\n",
+    "animate(0)\n",
+    "anim = FuncAnimation(fig, animate, frames=n_frames, interval=300)\n",
+    "plt.close()\n",
+    "anim"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "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."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python (parcels-pixi)",
+   "language": "python",
+   "name": "parcels-pixi"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.14.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}
diff --git a/docs/user_guide/index.md b/docs/user_guide/index.md
index 533020cc1f..17107a5e4a 100644
--- a/docs/user_guide/index.md
+++ b/docs/user_guide/index.md
@@ -42,6 +42,7 @@ examples/tutorial_mitgcm.ipynb
 examples/tutorial_delft3d.ipynb
 examples/tutorial_fesom.ipynb
 examples/tutorial_schism.ipynb
+examples/tutorial_unstructured_sigma_coordinates.ipynb
 ```
 
 ## Work with FieldSets
diff --git a/src/parcels/_core/basegrid.py b/src/parcels/_core/basegrid.py
index 7b7c7e9c49..ac3ae9c664 100644
--- a/src/parcels/_core/basegrid.py
+++ b/src/parcels/_core/basegrid.py
@@ -29,7 +29,7 @@ class BaseGrid(ABC):
     _mesh: FlatMesh | SphericalMesh
 
     @abstractmethod
-    def search(self, z: float, y: float, x: float, ei=None) -> dict[str, tuple[int, float | np.ndarray]]:
+    def search(self, z: float, y: float, x: float, ei=None, ti=None) -> dict[str, tuple[int, float | np.ndarray]]:
         """
         Perform a spatial (and optionally vertical) search to locate the grid element
         that contains a given point (x, y, z).
@@ -49,6 +49,8 @@ def search(self, z: float, y: float, x: float, ei=None) -> dict[str, tuple[int,
             A previously computed encoded index (e.g., raveled face or cell index). If provided,
             the search will first attempt to validate and reuse it before falling back to
             a global or local search strategy.
+        ti : np.ndarray, optional
+            Time index of each query point, as returned by ``_search_time_index``.
         search2D : bool, default=False
             If True, perform only a 2D search (x, y), ignoring the vertical component z.
 
diff --git a/src/parcels/_core/field.py b/src/parcels/_core/field.py
index ded5993edd..c8a470248d 100644
--- a/src/parcels/_core/field.py
+++ b/src/parcels/_core/field.py
@@ -398,8 +398,9 @@ def _get_positions(field: Field, t, z, y, x, particles, _ei) -> tuple[dict, dict
         raise ValueError(f"Time values for particles with indices {nan_indices} cannot be NaN.")
     particle_positions = {"t": t, "z": z, "y": y, "x": x}
     grid_positions = {}
-    grid_positions.update(_search_time_index(field, t))
-    grid_positions.update(field.grid.search(z, y, x, ei=_ei))
+    time_positions = _search_time_index(field, t)
+    grid_positions.update(time_positions)
+    grid_positions.update(field.grid.search(z, y, x, ei=_ei, ti=time_positions["T"]["index"]))
     _update_particles_ei(particles, grid_positions, field)
     _update_particle_states_position(particles, grid_positions)
     return particle_positions, grid_positions
diff --git a/src/parcels/_core/index_search.py b/src/parcels/_core/index_search.py
index 8406b68651..ae9c57bf59 100644
--- a/src/parcels/_core/index_search.py
+++ b/src/parcels/_core/index_search.py
@@ -62,6 +62,45 @@ def _search_1d_array(
     return np.atleast_1d(index), np.atleast_1d(bcoord)
 
 
+def _search_1d_columns(columns: np.ndarray, x: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Searches for particle locations in per-particle 1D columns and returns barycentric coordinate along dimension.
+
+    Row-wise counterpart of ``_search_1d_array``: particle p searches only its own column ``columns[p]``.
+
+    Assumptions:
+    - each column is strictly monotonically increasing.
+
+    Parameters
+    ----------
+    columns : np.ndarray
+        2D array of shape (n_particles, n_levels), one column per particle.
+    x : np.ndarray
+        Position of each particle along its column, shape (n_particles,).
+
+    Returns
+    -------
+    array of int
+        Index of the element just before the position x in each column. Note that this index is -2 if the index is left out of bounds and -1 if the index is right out of bounds.
+    array of float
+        Barycentric coordinate.
+    """
+    n_levels = columns.shape[1]
+    if n_levels < 2:
+        return np.zeros(shape=x.shape, dtype=np.int32), np.zeros_like(x)
+    # The number of column entries strictly below x equals np.searchsorted(column, x, side="left")
+    index = np.clip((columns < x[:, None]).sum(axis=1) - 1, 0, n_levels - 2)
+    rows = np.arange(columns.shape[0])
+    lower = columns[rows, index]
+    upper = columns[rows, index + 1]
+    bcoord = (x - lower) / (upper - lower)
+
+    index = np.where(x < columns[:, 0], LEFT_OUT_OF_BOUNDS, index)
+    index = np.where(x > columns[:, -1], RIGHT_OUT_OF_BOUNDS, index)
+
+    return index, bcoord
+
+
 def _search_time_index(field: Field, time: np.ndarray):
     """Find and return the index and relative coordinate in the time array associated with a given time.
 
diff --git a/src/parcels/_core/model.py b/src/parcels/_core/model.py
index eeddc75ae4..c36c6fb2bc 100644
--- a/src/parcels/_core/model.py
+++ b/src/parcels/_core/model.py
@@ -359,6 +359,9 @@ def __init__(self, data: ux.UxDataset, grid: UxGrid, vector_field_components: pt
         if not isinstance(grid, UxGrid):
             raise ValueError(f"Expected `grid` to be a Parcels UxGrid object. Got {type(grid)}.")
 
+        if grid.z.ndim == 3 and not grid.z["time"].equals(data["time"]):
+            raise ValueError("A time-varying (3D) z must have the same time coordinate as `data`.")
+
         self.data = data
         self.grid = grid
         self.vector_field_components = vector_field_components
diff --git a/src/parcels/_core/particleset.py b/src/parcels/_core/particleset.py
index 90ddb0e858..a349eb0127 100644
--- a/src/parcels/_core/particleset.py
+++ b/src/parcels/_core/particleset.py
@@ -19,6 +19,7 @@
     float_to_datelike,
     timedelta_to_float,
 )
+from parcels._core.uxgrid import UxGrid
 from parcels._core.warnings import ParticleSetWarning
 from parcels._logger import logger
 
@@ -82,6 +83,11 @@ def __init__(
         if z is None:
             minz = None
             for field in self.fieldset.fields.values():
+                if isinstance(field.grid, UxGrid) and field.grid.z.ndim == 3:
+                    raise ValueError(
+                        f"Field {field.name!r} has a time-varying (3D) vertical grid, so there is no default "
+                        "particle depth. Pass the particle depths explicitly with `z`."
+                    )
                 for depth in field.grid.depth:
                     if minz is None or np.abs(depth) < np.abs(minz):
                         minz = depth
diff --git a/src/parcels/_core/uxgrid.py b/src/parcels/_core/uxgrid.py
index 3257785dbd..50600cbbc6 100644
--- a/src/parcels/_core/uxgrid.py
+++ b/src/parcels/_core/uxgrid.py
@@ -3,10 +3,11 @@
 from typing import TYPE_CHECKING, Literal
 
 import numpy as np
+import xarray as xr
 from dask import is_dask_collection
 
 from parcels._core.basegrid import BaseGrid
-from parcels._core.index_search import GRID_SEARCH_ERROR, _search_1d_array, uxgrid_point_in_cell
+from parcels._core.index_search import GRID_SEARCH_ERROR, _search_1d_columns, uxgrid_point_in_cell
 from parcels._core.mesh import SphericalMesh, get_mesh
 
 if TYPE_CHECKING:
@@ -32,9 +33,9 @@ def __init__(
         grid : ux.grid.Grid
             The uxarray grid object containing the unstructured grid data.
         z : ux.UxDataArray
-            A 1D array of vertical coordinates (depths) associated with the layer interface heights (not the mid-layer depths).
-            While uxarray allows nz to be spatially and temporally varying, the parcels.UxGrid class considers the case where
-            the vertical coordinate is constant in time and space. This implies flat bottom topography and no moving ALE vertical grid.
+            Vertical coordinates (depths) of the layer interface heights (not the mid-layer depths). Either a 1D array,
+            constant in time and space (flat bottom topography, no moving vertical grid), or a 3D array with dims
+            ("time", "zf", "n_node"), varying in time and across the mesh nodes.
         mesh : str
             The type of mesh used for the grid. Either "flat" or "spherical".
         """
@@ -45,8 +46,10 @@ def __init__(
         self.uxgrid = grid
         if not isinstance(z, ux.UxDataArray):
             raise TypeError("z must be an instance of ux.UxDataArray")
-        if z.ndim != 1:
-            raise ValueError("z must be a 1D array of vertical coordinates")
+        if z.ndim not in (1, 3):
+            raise ValueError(f"z must be a 1D or 3D array of vertical coordinates, got {z.ndim}D")
+        if z.ndim == 3 and z.dims != ("time", "zf", "n_node"):
+            raise ValueError(f"A 3D z must have dims ('time', 'zf', 'n_node'), got {z.dims}")
         self.z = z
         self._mesh = get_mesh(mesh)
         self._spatialhash = None
@@ -76,6 +79,8 @@ def get_axis_dim(self, axis: _UXGRID_AXES) -> int:
             raise ValueError(f"Axis {axis!r} is not part of this grid. Available axes: {self.axes}")
 
         if axis == "Z":
+            if self.z.ndim == 3:
+                return self.z.sizes["zf"]
             return len(self.z.values)
         elif axis == "FACE":
             return self.uxgrid.n_face
@@ -87,7 +92,7 @@ def deg2m(self) -> float:
             return self._mesh.deg2m
         return 1.0
 
-    def search(self, z, y, x, ei=None, tol=1e-6):
+    def search(self, z, y, x, ei=None, ti=None, tol=1e-6):
         """
         Search for the grid cell (face) and vertical layer that contains the given points.
 
@@ -105,6 +110,9 @@ def search(self, z, y, x, ei=None, tol=1e-6):
             TO BE IMPLEMENTED : If provided, we'll check
             if the points are within the faces specified by these indices. For cells where the particles
             are not found, a nearest neighbor search will be performed. As a last resort, the spatial hash will be used.
+        ti : np.ndarray, optional
+            Time index of each point, as returned by ``_search_time_index``. Required when z is 3D; selects the z
+            snapshot used for each particle's column.
         tol : float, optional
             Tolerance for barycentric coordinate checks. Default is 1e-6.
         """
@@ -112,8 +120,6 @@ def search(self, z, y, x, ei=None, tol=1e-6):
         y = np.asarray(y, dtype=np.float32)
         z = np.asarray(z, dtype=np.float32)
 
-        zi, zeta = _search_1d_array(self.z.values, z)
-
         if np.any(ei):
             indices = self.unravel_index(ei)
             fi = indices.get("FACE")
@@ -134,4 +140,25 @@ def search(self, z, y, x, ei=None, tol=1e-6):
             coords[zero_indices, :] = coords_q
             fi[zero_indices] = face_ids_q
 
+        found = fi >= 0
+        if self.z.ndim == 3:
+            if ti is None:
+                raise ValueError("Searching a UxGrid with a time-varying (3D) z requires the time index ti")
+            node_columns = self.z.isel(
+                time=xr.DataArray(np.broadcast_to(ti, fi.shape)[found], dims="points"),
+                n_node=xr.DataArray(self.uxgrid.face_node_connectivity[fi[found], :].values, dims=("points", "nodes")),
+                ignore_grid=True,
+            ).compute()
+            bcoords = xr.DataArray(coords[found], dims=("points", "nodes"))
+            # Particles outside the mesh are given a NaN z column for vertical searching
+            columns = np.full((fi.size, self.z.sizes["zf"]), np.nan)
+            columns[found] = xr.dot(node_columns, bcoords, dim="nodes").transpose("points", "zf").values
+        else:
+            columns = np.broadcast_to(self.z.values, (z.size, self.z.size))
+
+        zi, zeta = _search_1d_columns(columns, z)
+        # Particles outside the mesh are given a 0 vertical index and a NaN vertical barycentric coordinate
+        zi = np.where(found, zi, 0)
+        zeta = np.where(found, zeta, np.nan)
+
         return {"Z": {"index": zi, "bcoord": zeta}, "FACE": {"index": fi, "bcoord": coords}}
diff --git a/src/parcels/_core/xgrid.py b/src/parcels/_core/xgrid.py
index 70603d914a..c3f69b0f22 100644
--- a/src/parcels/_core/xgrid.py
+++ b/src/parcels/_core/xgrid.py
@@ -313,7 +313,7 @@ def _gtype(self):
             else:
                 return GridType.CurvilinearSGrid
 
-    def search(self, z, y, x, ei=None):
+    def search(self, z, y, x, ei=None, ti=None):
         ds = self._ds
 
         if "Z" in self.axes:
diff --git a/src/parcels/_datasets/unstructured/generated.py b/src/parcels/_datasets/unstructured/generated.py
index 35d435ec47..a31145de08 100644
--- a/src/parcels/_datasets/unstructured/generated.py
+++ b/src/parcels/_datasets/unstructured/generated.py
@@ -165,6 +165,28 @@ def _build_delaunay_grid(nx, lon_range, lat_range):
     return uxgrid
 
 
+def _build_triangulated_lattice_grid(nx, x_range, y_range):
+    """Build a flat UxGrid by splitting each cell of a regular nx-by-nx node lattice over x_range by y_range (in meters) into two triangles."""
+    X, Y = np.meshgrid(np.linspace(x_range[0], x_range[1], nx), np.linspace(y_range[0], y_range[1], nx))
+    node_x, node_y = X.ravel(), Y.ravel()
+    i, j = np.meshgrid(np.arange(nx - 1), np.arange(nx - 1))
+    i, j = i.ravel(), j.ravel()
+    sw = (j * nx) + i
+    se = sw + 1
+    nw = sw + nx
+    ne = nw + 1
+
+    tri1 = np.stack((sw, ne, nw), axis=-1)
+    tri2 = np.stack((sw, se, ne), axis=-1)
+    face_node_connectivity = np.concatenate((tri1, tri2))
+
+    uxgrid = ux.Grid.from_topology(node_lon=node_x, node_lat=node_y, face_node_connectivity=face_node_connectivity)
+    uxgrid.node_lon.values[:] = node_x
+    uxgrid.node_lat.values[:] = node_y
+    uxgrid.attrs["Conventions"] = "UGRID-1.0"
+    return uxgrid
+
+
 def _wrap_uvw_dataset(uxgrid, u, v, w, zc, zf, time, uv_dim, uv_location, description):
     """Wrap (u, v, w) numpy arrays into a UxDataset following Parcels' UGRID conventions.
 
@@ -342,3 +364,47 @@ def solid_body_rotation_3d_node_centered(nx=40, nz=10, omega=2.0 * math.pi / 360
     w = np.full((1, zf.size, uxgrid.n_node), w0, dtype=np.float64)
 
     return _wrap_uvw_dataset(uxgrid, u, v, w, zc, zf, time, "n_node", "node", "3D solid-body rotation")
+
+
+def sigma_coordinate_lattice_dataset(nx, x_range, y_range, nz, bottom_depth, eta=None, start_time="2000-01-01"):
+    """Zero-velocity UGRID dataset on an nx-by-nx lattice: nz sigma interfaces from the surface (depth -eta) to bottom_depth (scalar or (x, y), meters); eta is None or (t, x, y), hourly from start_time."""
+    uxgrid = _build_triangulated_lattice_grid(nx, x_range, y_range)
+    sigma = np.linspace(0.0, 1.0, nz)
+    if bottom_depth is None or np.ndim(bottom_depth) not in (0, 2):
+        raise ValueError(
+            f"bottom_depth must be a scalar or an (x, y) array, got {type(bottom_depth).__name__} with ndim {np.ndim(bottom_depth)}"
+        )
+
+    if eta is None:
+        if np.ndim(bottom_depth) == 0:
+            zf = sigma*bottom_depth # (nz,)
+        elif np.ndim(bottom_depth)== 2:
+            zf = (sigma*(bottom_depth.T.ravel()[:, np.newaxis])).T # (nz, n_node)
+            zf = zf[np.newaxis, :, :] # (1, nz, n_node)
+    else:
+        total_depth = np.swapaxes(bottom_depth + eta, 1, 2).reshape(eta.shape[0], -1)
+        eta_nodes = np.swapaxes(eta, 1, 2).reshape(eta.shape[0], -1)
+        zf = sigma[np.newaxis, :, np.newaxis] * total_depth[:, np.newaxis, :] - eta_nodes[:, np.newaxis, :]
+
+    n_times = 1 if eta is None else eta.shape[0]
+    time = xr.date_range(start_time, periods=n_times, freq="1h")
+    zc = 0.5 * (zf[:-1] + zf[1:]) if zf.ndim == 1 else 0.5 * (zf[:, :-1] + zf[:, 1:])
+
+    if zf.ndim == 1:
+        vertical_coords = {"zf": ("zf", zf), "zc": ("zc", zc)}
+    else:
+        vertical_coords = {"zf": (("time", "zf", "n_node"), zf), "zc": (("time", "zc", "n_node"), zc)}
+
+    velocity_dims = ("time", "zf", "n_node")
+    velocity_shape = (n_times, nz, uxgrid.n_node)
+    velocity_attrs = dict(units="m/s", location="node", mesh="triangulated_lattice", Conventions="UGRID-1.0")
+    velocities = {name: (velocity_dims, np.zeros(velocity_shape), velocity_attrs) for name in ("U", "V", "W")}
+
+    return ux.UxDataset(xr.Dataset(velocities, coords={"time": time, **vertical_coords}), uxgrid=uxgrid)
+
+
+
+
+
+
+
diff --git a/src/parcels/interpolators/_uxinterpolators.py b/src/parcels/interpolators/_uxinterpolators.py
index 17249bfdc5..baa3f5859b 100644
--- a/src/parcels/interpolators/_uxinterpolators.py
+++ b/src/parcels/interpolators/_uxinterpolators.py
@@ -54,7 +54,7 @@ def interp(
         ti, zi, fi = np.broadcast_arrays(
             grid_positions["T"]["index"], grid_positions["Z"]["index"], grid_positions["FACE"]["index"]
         )
-        z = particle_positions["z"]
+        zeta = grid_positions["Z"]["bcoord"]
 
         tdim, zdim, fdim = field.data.dims
 
@@ -75,9 +75,7 @@ def _zsample(z_index):
         fzkp1 = _zsample(zi + 1)
 
         # Then, do piecewise linear interpolation in the vertical direction
-        zk = field.grid.z.values[zi]
-        zkp1 = field.grid.z.values[zi + 1]
-        return (fzk * (zkp1 - z) + fzkp1 * (z - zk)) / (zkp1 - zk)  # Linear interpolation in the vertical direction
+        return fzk * (1 - zeta) + fzkp1 * zeta
 
 
 class UxLinearNodeConstantZC(ScalarInterpolator):
@@ -129,7 +127,7 @@ def interp(
         ti, zi, fi = np.broadcast_arrays(
             grid_positions["T"]["index"], grid_positions["Z"]["index"], grid_positions["FACE"]["index"]
         )
-        z = particle_positions["z"]
+        zeta = grid_positions["Z"]["bcoord"]
         bcoords = xr.DataArray(grid_positions["FACE"]["bcoord"], dims=("points", "nodes"))
         node_ids = field.grid.uxgrid.face_node_connectivity[fi, :].values
 
@@ -153,9 +151,7 @@ def _zsample(z_index):
         fzkp1 = _zsample(zi + 1)
 
         # Then, do piecewise linear interpolation in the vertical direction
-        zk = field.grid.z.values[zi]
-        zkp1 = field.grid.z.values[zi + 1]
-        value = (fzk * (zkp1 - z) + fzkp1 * (z - zk)) / (zkp1 - zk)  # Linear interpolation in the vertical direction
+        value = fzk * (1 - zeta) + fzkp1 * zeta
         return value.compute() if is_dask_collection(value) else value
 
 
diff --git a/tests/test_uxadvection.py b/tests/test_uxadvection.py
index 8ffd84b1a8..85ee920768 100644
--- a/tests/test_uxadvection.py
+++ b/tests/test_uxadvection.py
@@ -3,11 +3,13 @@
 import pytest
 
 import parcels
+from parcels._datasets.unstructured.generated import sigma_coordinate_lattice_dataset
 from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
 from parcels.kernels import (
     AdvectionEE,
     AdvectionRK2,
     AdvectionRK4,
+    AdvectionRK4_3D,
 )
 
 
@@ -26,3 +28,23 @@ def test_ux_constant_flow_face_centered_2D(integrator, tmp_parquet):
 
     df = pd.read_parquet(tmp_parquet)
     np.testing.assert_allclose(df["x"].iloc[-1], expected_lon, atol=1e-5)
+
+
+def test_ux_advection_on_moving_sigma_grid_keeps_depth():
+    """With uniform u and w = 0, particles keep their depth while the sigma levels move past them."""
+    nx, nz = 5, 6
+    x_nodes, y_nodes = np.meshgrid(np.linspace(0.0, 4e3, nx), np.linspace(0.0, 4e3, nx), indexing="ij")
+    bottom_depth = 20.0 + 1e-2 * x_nodes + 2e-3 * y_nodes
+    eta = np.stack([np.sin(snapshot + 1e-3 * x_nodes) for snapshot in range(3)])
+    ds = sigma_coordinate_lattice_dataset(nx, (0.0, 4e3), (0.0, 4e3), nz, bottom_depth, eta)
+    u0 = 0.5
+    ds["U"].values[:] = u0
+    fieldset = parcels.FieldSet.from_ugrid_conventions(ds, mesh="flat")
+
+    x0, z0 = 200.0, np.array([2.0, 8.0, 15.0])
+    pset = parcels.ParticleSet(fieldset, x=np.full(z0.size, x0), y=np.full(z0.size, 2e3), z=z0)
+    pset.execute(AdvectionRK4_3D, runtime=np.timedelta64(2, "h"), dt=np.timedelta64(300, "s"), verbose_progress=False)
+
+    np.testing.assert_array_equal(pset.z, z0)
+    np.testing.assert_allclose(pset.x, x0 + u0 * 7200.0)
+    assert np.all(pset.state < parcels.StatusCode.Error)
diff --git a/tests/test_uxarray_fieldset.py b/tests/test_uxarray_fieldset.py
index b080f52001..7535ed4a26 100644
--- a/tests/test_uxarray_fieldset.py
+++ b/tests/test_uxarray_fieldset.py
@@ -8,11 +8,15 @@
 import parcels.tutorial
 from parcels import (
     FieldSet,
+    UxGrid,
     convert,
 )
+from parcels._core.model import UnstructuredModelData
+from parcels._datasets.unstructured.generated import sigma_coordinate_lattice_dataset
 from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
 from parcels.interpolators import (
     UxConstantFaceConstantZC,
+    UxConstantFaceLinearZF,
     UxLinearNodeLinearZF,
 )
 from tests.utils import create_uxgrid_from_triangulation
@@ -204,3 +208,82 @@ def test_nestedgrids_triangulation_spherical_search():
         f"grid search failed for {n_lost} of {len(x)} particles; e.g. (lon, lat)="
         f"{np.column_stack((x, y))[face < 0][:3].tolist()}"
     )
+
+
+SLOPING_BOTTOM_DEPTH = np.broadcast_to(np.linspace(20.0, 60.0, 5)[:, np.newaxis], (5, 5))
+RISING_ETA = np.broadcast_to(np.linspace(-1.0, 1.0, 3)[:, np.newaxis, np.newaxis], (3, 5, 5))
+
+
+@pytest.mark.parametrize(
+    "eta, expected_n_snapshots",
+    [
+        pytest.param(None, 1, id="single_snapshot_3d_z"),
+        pytest.param(RISING_ETA, 3, id="time_varying_3d_z"),
+    ],
+)
+def test_fieldset_from_sigma_coordinate_dataset(eta, expected_n_snapshots):
+    nz = 4
+    ds = sigma_coordinate_lattice_dataset(5, (0.0, 4e3), (0.0, 4e3), nz, SLOPING_BOTTOM_DEPTH, eta=eta)
+    fieldset = FieldSet.from_ugrid_conventions(ds, mesh="flat")
+
+    assert fieldset.U.grid.z.dims == ("time", "zf", "n_node")
+    assert fieldset.U.grid.z.sizes["time"] == expected_n_snapshots
+    assert fieldset.U.grid.get_axis_dim("Z") == nz
+    for field in (fieldset.U, fieldset.V, fieldset.W):
+        assert isinstance(field.interp_method, UxLinearNodeLinearZF)
+
+
+@pytest.mark.parametrize(
+    "horizontal_dim, expected_interpolator",
+    [
+        pytest.param("n_node", UxLinearNodeLinearZF, id="node_registered"),
+        pytest.param("n_face", UxConstantFaceLinearZF, id="face_registered"),
+    ],
+)
+def test_sigma_linear_field_is_exact_on_moving_sigma_grid(horizontal_dim, expected_interpolator):
+    """A field equal to a + b * sigma on every interface must evaluate to exactly a + b * sigma at any particle.
+
+    Bottom depth and eta are linear in x and y, so barycentric interpolation of the node columns is exact and a
+    particle placed at sigma = s has a known depth.
+    """
+    nx, nz = 11, 8
+    n_snapshots = 5
+    a, b = 0.3, 1.7
+
+    def bottom_depth(x, y):
+        return 20.0 + 4e-3 * x + 2e-3 * y
+
+    def eta(snapshot, x, y):
+        return 0.3 * snapshot + (-1.0) ** snapshot * 1e-4 * x + 5e-5 * y
+
+    x_nodes, y_nodes = np.meshgrid(np.linspace(0.0, 10e3, nx), np.linspace(0.0, 10e3, nx), indexing="ij")
+    eta_nodes = np.stack([eta(snapshot, x_nodes, y_nodes) for snapshot in range(n_snapshots)])
+    ds = sigma_coordinate_lattice_dataset(nx, (0.0, 10e3), (0.0, 10e3), nz, bottom_depth(x_nodes, y_nodes), eta_nodes)
+    field_on_interfaces = a + b * np.linspace(0.0, 1.0, nz)
+    field_shape = (n_snapshots, nz, getattr(ds.uxgrid, horizontal_dim))
+    ds["F"] = (("time", "zf", horizontal_dim), np.broadcast_to(field_on_interfaces[None, :, None], field_shape).copy())
+    fieldset = FieldSet.from_ugrid_conventions(ds, mesh="flat")
+    assert isinstance(fieldset.F.interp_method, expected_interpolator)
+
+    rng = np.random.default_rng(0)
+    n_particles = 100
+    x = rng.uniform(1.0, 10e3 - 1.0, n_particles)
+    y = rng.uniform(1.0, 10e3 - 1.0, n_particles)
+    t = np.concatenate([[0.0, 3600.0, 7200.0], rng.uniform(0.0, (n_snapshots - 1) * 3600.0, n_particles - 3)])
+    snapshot = np.maximum(np.ceil(t / 3600.0).astype(int) - 1, 0)
+    sigma = rng.uniform(0.02, 0.98, n_particles)
+    eta_particles = eta(snapshot, x, y)
+    z = -eta_particles + sigma * (bottom_depth(x, y) + eta_particles)
+
+    np.testing.assert_allclose(fieldset.F.eval(t, z, y, x), a + b * sigma, rtol=1e-6)
+
+
+def test_unstructured_model_data_rejects_z_with_different_time_coordinate():
+    ds = sigma_coordinate_lattice_dataset(5, (0.0, 4e3), (0.0, 4e3), 4, SLOPING_BOTTOM_DEPTH, RISING_ETA)
+    ds_one_day_later = sigma_coordinate_lattice_dataset(
+        5, (0.0, 4e3), (0.0, 4e3), 4, SLOPING_BOTTOM_DEPTH, RISING_ETA, start_time="2000-01-02"
+    )
+    grid_with_later_z = UxGrid(ds.uxgrid, z=ds_one_day_later.coords["zf"], mesh="flat")
+
+    with pytest.raises(ValueError, match="same time coordinate"):
+        UnstructuredModelData(ds, grid_with_later_z, {})
diff --git a/tests/test_uxgrid.py b/tests/test_uxgrid.py
index 63df6a3cae..3d9800734b 100644
--- a/tests/test_uxgrid.py
+++ b/tests/test_uxgrid.py
@@ -1,6 +1,8 @@
+import numpy as np
 import pytest
 
 from parcels import UxGrid
+from parcels._datasets.unstructured.generated import sigma_coordinate_lattice_dataset
 from parcels._datasets.unstructured.generic import datasets as uxdatasets
 
 GENERIC_Z_COORDS = ["nz", "zf", "depth_2"]
@@ -32,3 +34,11 @@ def test_xgrid_get_axis_dim(uxds):
 
     assert grid.get_axis_dim("FACE") == 721
     assert grid.get_axis_dim("Z") == 2
+
+
+def test_uxgrid_search_3d_z_requires_ti():
+    ds = sigma_coordinate_lattice_dataset(5, (0.0, 4e3), (0.0, 4e3), 4, np.full((5, 5), 50.0))
+    grid = UxGrid(ds.uxgrid, z=ds.coords["zf"], mesh="flat")
+
+    with pytest.raises(ValueError, match="requires the time index ti"):
+        grid.search(np.array([10.0]), np.array([2e3]), np.array([2e3]))
diff --git a/tests/test_xgrid.py b/tests/test_xgrid.py
index 2d28d4aef9..d98e11362b 100644
--- a/tests/test_xgrid.py
+++ b/tests/test_xgrid.py
@@ -12,6 +12,7 @@
     LEFT_OUT_OF_BOUNDS,
     RIGHT_OUT_OF_BOUNDS,
     _search_1d_array,
+    _search_1d_columns,
 )
 from parcels._core.utils.time import timedelta_to_float
 from parcels._core.xgrid import (
@@ -244,6 +245,19 @@ def corner_to_cell_center_points(lat, lon):
     return lat_c, lon_c
 
 
+def _search_shared_column(array, x):
+    """Run ``_search_1d_columns`` with every query point searching the same column ``array``."""
+    x = np.atleast_1d(np.asarray(x, dtype=float))
+    return _search_1d_columns(np.broadcast_to(array, (x.size, array.size)), x)
+
+
+SEARCH_1D_FUNCTIONS = [
+    pytest.param(_search_1d_array, id="search_1d_array"),
+    pytest.param(_search_shared_column, id="search_1d_columns"),
+]
+
+
+@pytest.mark.parametrize("search_1d", SEARCH_1D_FUNCTIONS)
 @pytest.mark.parametrize(
     "array, x, expected_xi, expected_xsi",
     [
@@ -253,12 +267,13 @@ def corner_to_cell_center_points(lat, lon):
         (np.array([1, 2, 3, 4, 5]), 4.5, 3, 0.5),
     ],
 )
-def test_search_1d_array(array, x, expected_xi, expected_xsi):
-    xi, xsi = _search_1d_array(array, x)
+def test_search_1d_array(search_1d, array, x, expected_xi, expected_xsi):
+    xi, xsi = search_1d(array, x)
     np.testing.assert_array_equal(xi, expected_xi)
     np.testing.assert_allclose(xsi, expected_xsi)
 
 
+@pytest.mark.parametrize("search_1d", SEARCH_1D_FUNCTIONS)
 @pytest.mark.parametrize(
     "array, x, expected_xi",
     [
@@ -266,11 +281,12 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi):
         (np.array([1, 2, 3, 4, 5]), 6.5, RIGHT_OUT_OF_BOUNDS),
     ],
 )
-def test_search_1d_array_out_of_bounds(array, x, expected_xi):
-    xi, _xsi = _search_1d_array(array, x)
+def test_search_1d_array_out_of_bounds(search_1d, array, x, expected_xi):
+    xi, _xsi = search_1d(array, x)
     assert xi == expected_xi
 
 
+@pytest.mark.parametrize("search_1d", SEARCH_1D_FUNCTIONS)
 @pytest.mark.parametrize(
     "array, x, expected_xi",
     [
@@ -278,8 +294,8 @@ def test_search_1d_array_out_of_bounds(array, x, expected_xi):
         (np.array([1, 2, 3, 4, 5]), (6.5, 1), (RIGHT_OUT_OF_BOUNDS, 0)),
     ],
 )
-def test_search_1d_array_some_out_of_bounds(array, x, expected_xi):
-    xi, _ = _search_1d_array(array, x)
+def test_search_1d_array_some_out_of_bounds(search_1d, array, x, expected_xi):
+    xi, _ = search_1d(array, x)
     np.testing.assert_array_equal(xi, expected_xi)
 
 

From 57cf6ddbc5c1dde7be681302aead82ca641109ec Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 5 Oct 2026 21:00:50 +0000
Subject: [PATCH 02/22] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 ...orial_unstructured_sigma_coordinates.ipynb | 50 ++++++++++++++-----
 .../_datasets/unstructured/generated.py       | 15 ++----
 2 files changed, 42 insertions(+), 23 deletions(-)

diff --git a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
index ee32d8ede2..00b444f70e 100644
--- a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
+++ b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
@@ -42,10 +42,16 @@
    "source": [
     "nx = 21  # nodes along and across the channel\n",
     "n_layers = 10\n",
-    "x_nodes, _ = np.meshgrid(np.linspace(0.0, 40e3, nx), np.linspace(0.0, 10e3, nx), indexing=\"ij\")\n",
-    "bottom_depth = 10.0 + 90.0 * np.log10(1.0 + 9.0 * x_nodes / 40e3)  # (x, y) array, in meters\n",
-    "\n",
-    "ds_static = sigma_coordinate_lattice_dataset(nx, (0.0, 40e3), (0.0, 10e3), n_layers + 1, bottom_depth)\n",
+    "x_nodes, _ = np.meshgrid(\n",
+    "    np.linspace(0.0, 40e3, nx), np.linspace(0.0, 10e3, nx), indexing=\"ij\"\n",
+    ")\n",
+    "bottom_depth = 10.0 + 90.0 * np.log10(\n",
+    "    1.0 + 9.0 * x_nodes / 40e3\n",
+    ")  # (x, y) array, in meters\n",
+    "\n",
+    "ds_static = sigma_coordinate_lattice_dataset(\n",
+    "    nx, (0.0, 40e3), (0.0, 10e3), n_layers + 1, bottom_depth\n",
+    ")\n",
     "node_x, node_y = ds_static.uxgrid.node_lon.values, ds_static.uxgrid.node_lat.values\n",
     "face_nodes = ds_static.uxgrid.face_node_connectivity.values\n",
     "\n",
@@ -93,7 +99,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "centre_line = np.isclose(node_y, 5e3) \n",
+    "centre_line = np.isclose(node_y, 5e3)\n",
     "xs = node_x[centre_line] / 1e3\n",
     "\n",
     "fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)\n",
@@ -131,9 +137,15 @@
     "\n",
     "# for illustrative purposes we put a 2m surface wave that propogates with time\n",
     "wavelength, wave_period = 20e3, 6 * 3600.0\n",
-    "wave = 2 * np.sin(2.0 * np.pi * (x_nodes[np.newaxis] / wavelength - t[:, np.newaxis, np.newaxis] / wave_period))\n",
-    "\n",
-    "ds = sigma_coordinate_lattice_dataset(nx, (0.0, 40e3), (0.0, 10e3), n_layers + 1, bottom_depth, eta=wave)\n",
+    "wave = 2 * np.sin(\n",
+    "    2.0\n",
+    "    * np.pi\n",
+    "    * (x_nodes[np.newaxis] / wavelength - t[:, np.newaxis, np.newaxis] / wave_period)\n",
+    ")\n",
+    "\n",
+    "ds = sigma_coordinate_lattice_dataset(\n",
+    "    nx, (0.0, 40e3), (0.0, 10e3), n_layers + 1, bottom_depth, eta=wave\n",
+    ")\n",
     "time = ds[\"time\"].values\n",
     "z_moving = ds[\"zf\"].values  # (time, zf, n_node)\n",
     "eta = -z_moving[:, 0, :]  # the surface interface sits at depth -eta"
@@ -159,11 +171,13 @@
     "    ax.set_xlabel(\"x [km]\")\n",
     "plt.tight_layout()\n",
     "\n",
+    "\n",
     "def animate(snapshot):\n",
     "    for line, interface_depths in zip(z_lines, z_moving[snapshot][:, centre_line]):\n",
     "        line.set_ydata(interface_depths)\n",
     "    z_title.set_text(f\"Layer interfaces in z space at {time[snapshot]}\")\n",
     "\n",
+    "\n",
     "anim = FuncAnimation(fig, animate, frames=time.size, interval=500)\n",
     "plt.close()\n",
     "anim"
@@ -222,7 +236,12 @@
     "z0 = np.linspace(3.0, 22.0, n_particles)  # depths in m\n",
     "\n",
     "pset = parcels.ParticleSet(\n",
-    "    fieldset=fieldset, pclass=parcels.Particle, t=np.full(n_particles, time[0]), z=z0, y=y0, x=x0\n",
+    "    fieldset=fieldset,\n",
+    "    pclass=parcels.Particle,\n",
+    "    t=np.full(n_particles, time[0]),\n",
+    "    z=z0,\n",
+    "    y=y0,\n",
+    "    x=x0,\n",
     ")"
    ]
   },
@@ -272,7 +291,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "outputfile = parcels.ParticleFile(path=\"sigma_particles.parquet\", outputdt=np.timedelta64(30, \"m\"), mode=\"w\")\n",
+    "outputfile = parcels.ParticleFile(\n",
+    "    path=\"sigma_particles.parquet\", outputdt=np.timedelta64(30, \"m\"), mode=\"w\"\n",
+    ")\n",
     "\n",
     "pset.execute(\n",
     "    [parcels.kernels.AdvectionRK4_3D],\n",
@@ -299,7 +320,9 @@
     "psigma = np.empty_like(pz)\n",
     "for frame, snapshot in enumerate(snapshots):\n",
     "    eta_p = np.interp(px[frame], x_line, eta[snapshot, centre_line])\n",
-    "    psigma[frame] = (pz[frame] + eta_p) / (np.interp(px[frame], x_line, H[centre_line]) + eta_p)\n",
+    "    psigma[frame] = (pz[frame] + eta_p) / (\n",
+    "        np.interp(px[frame], x_line, H[centre_line]) + eta_p\n",
+    "    )\n",
     "\n",
     "colors = [f\"C{p}\" for p in range(n_particles)]\n",
     "\n",
@@ -320,13 +343,16 @@
     "    ax.set_xlabel(\"x [km]\")\n",
     "plt.tight_layout()\n",
     "\n",
+    "\n",
     "def animate(frame):\n",
     "    for p in range(n_particles):\n",
     "        sigma_trails[p].set_data(px[: frame + 1, p] / 1e3, psigma[: frame + 1, p])\n",
     "        z_trails[p].set_data(px[: frame + 1, p] / 1e3, pz[: frame + 1, p])\n",
     "    sigma_dots.set_offsets(np.c_[px[frame] / 1e3, psigma[frame]])\n",
     "    z_dots.set_offsets(np.c_[px[frame] / 1e3, pz[frame]])\n",
-    "    for line, interface_depths in zip(z_lines, z_moving[snapshots[frame]][:, centre_line]):\n",
+    "    for line, interface_depths in zip(\n",
+    "        z_lines, z_moving[snapshots[frame]][:, centre_line]\n",
+    "    ):\n",
     "        line.set_ydata(interface_depths)\n",
     "    z_title.set_text(f\"Particles in z space at {output_times[frame]}\")\n",
     "\n",
diff --git a/src/parcels/_datasets/unstructured/generated.py b/src/parcels/_datasets/unstructured/generated.py
index a31145de08..68e7c96313 100644
--- a/src/parcels/_datasets/unstructured/generated.py
+++ b/src/parcels/_datasets/unstructured/generated.py
@@ -377,10 +377,10 @@ def sigma_coordinate_lattice_dataset(nx, x_range, y_range, nz, bottom_depth, eta
 
     if eta is None:
         if np.ndim(bottom_depth) == 0:
-            zf = sigma*bottom_depth # (nz,)
-        elif np.ndim(bottom_depth)== 2:
-            zf = (sigma*(bottom_depth.T.ravel()[:, np.newaxis])).T # (nz, n_node)
-            zf = zf[np.newaxis, :, :] # (1, nz, n_node)
+            zf = sigma * bottom_depth  # (nz,)
+        elif np.ndim(bottom_depth) == 2:
+            zf = (sigma * (bottom_depth.T.ravel()[:, np.newaxis])).T  # (nz, n_node)
+            zf = zf[np.newaxis, :, :]  # (1, nz, n_node)
     else:
         total_depth = np.swapaxes(bottom_depth + eta, 1, 2).reshape(eta.shape[0], -1)
         eta_nodes = np.swapaxes(eta, 1, 2).reshape(eta.shape[0], -1)
@@ -401,10 +401,3 @@ def sigma_coordinate_lattice_dataset(nx, x_range, y_range, nz, bottom_depth, eta
     velocities = {name: (velocity_dims, np.zeros(velocity_shape), velocity_attrs) for name in ("U", "V", "W")}
 
     return ux.UxDataset(xr.Dataset(velocities, coords={"time": time, **vertical_coords}), uxgrid=uxgrid)
-
-
-
-
-
-
-

From 4f921d6b352e20bafc8b59aa553cb18b3b3b0973 Mon Sep 17 00:00:00 2001
From: Erik van Sebille <e.vansebille@uu.nl>
Date: Tue, 6 Oct 2026 08:25:06 +0200
Subject: [PATCH 03/22] Removing unused polars import

---
 .../examples/tutorial_unstructured_sigma_coordinates.ipynb       | 1 -
 1 file changed, 1 deletion(-)

diff --git a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
index 00b444f70e..bc7f6d6240 100644
--- a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
+++ b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
@@ -15,7 +15,6 @@
    "source": [
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
-    "import polars as pl\n",
     "from matplotlib.animation import FuncAnimation\n",
     "\n",
     "import parcels\n",

From a12142d341eb764aeab6426beb26aa64784f57c8 Mon Sep 17 00:00:00 2001
From: Wyatt Sieminski <wyatt@fluidnumerics.com>
Date: Tue, 6 Oct 2026 19:17:00 +0000
Subject: [PATCH 04/22] Updated markdown cells in new tutorial

---
 ...orial_unstructured_sigma_coordinates.ipynb | 120 +++++-------------
 1 file changed, 31 insertions(+), 89 deletions(-)

diff --git a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
index ee32d8ede2..c8c8818910 100644
--- a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
+++ b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
@@ -15,13 +15,11 @@
    "source": [
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
-    "import polars as pl\n",
     "from matplotlib.animation import FuncAnimation\n",
     "\n",
     "import parcels\n",
     "from parcels._datasets.unstructured.generated import sigma_coordinate_lattice_dataset\n",
     "\n",
-    "# for interactive display of animations\n",
     "plt.rcParams[\"animation.html\"] = \"jshtml\""
    ]
   },
@@ -29,9 +27,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Build the channel mesh\n",
+    "## Building the Dataset\n",
     "\n",
-    "We begin by constructing a 2D mesh that we will later extend into 3 dimensions."
+    "For this tutorial, we construct our dataset using the parcels built in {py:func}`parcels._datasets.unstructured.generated.sigma_coordinate_lattice_dataset`, which makes an unstructured horizontally gridded dataset that uses a linear vertical $\\sigma$ grid. 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 height 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:"
    ]
   },
   {
@@ -40,85 +38,35 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "nx = 21  # nodes along and across the channel\n",
-    "n_layers = 10\n",
-    "x_nodes, _ = np.meshgrid(np.linspace(0.0, 40e3, nx), np.linspace(0.0, 10e3, nx), indexing=\"ij\")\n",
-    "bottom_depth = 10.0 + 90.0 * np.log10(1.0 + 9.0 * x_nodes / 40e3)  # (x, y) array, in meters\n",
+    "nx = 40 # the number of grid cells in x and y\n",
+    "x_range, y_range = (0.0, 40e3), (0.0, 10e3) # we set the bounds of the domain to be a 40km x 10km channel\n",
+    "n_vertical_layers = 10\n",
     "\n",
-    "ds_static = sigma_coordinate_lattice_dataset(nx, (0.0, 40e3), (0.0, 10e3), n_layers + 1, bottom_depth)\n",
-    "node_x, node_y = ds_static.uxgrid.node_lon.values, ds_static.uxgrid.node_lat.values\n",
-    "face_nodes = ds_static.uxgrid.face_node_connectivity.values\n",
+    "x_vals = np.linspace(x_range[0], x_range[1], nx)\n",
+    "y_vals = np.linspace(y_range[0], y_range[1], nx)\n",
     "\n",
-    "print(f\"n_node={ds_static.uxgrid.n_node}, n_face={ds_static.uxgrid.n_face}\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "fig, ax = plt.subplots(figsize=(10, 3))\n",
-    "ax.triplot(node_x / 1e3, node_y / 1e3, face_nodes, color=\"k\", linewidth=0.5)\n",
-    "ax.set_aspect(\"equal\")\n",
-    "ax.set_xlabel(\"x [km]\")\n",
-    "ax.set_ylabel(\"y [km]\")\n",
-    "ax.set_title(\"Channel mesh\")\n",
-    "plt.show()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Build the vertical grid\n",
-    "\n",
-    "Next we give the channel a terrain-following (sigma) vertical coordinate. The bottom depth $H$ increases logarithmically along the channel, from 10 m at the shallow left end to 100 m at the deep right end, so the bottom slopes steeply near the shallow end and flattens towards the deep end, typical of a coastal shelf. Following Parcels' convention that depth is positive downward, we define $\\sigma = \\frac{z + \\eta}{H + \\eta}$, so $\\sigma = 0$ at the surface and $\\sigma = 1$ at the bottom. Initially, we assume $\\eta$ to be zero. We place 11 layer interfaces (10 layers) equally spaced in $\\sigma$, so every node has its own column of interface depths $z = \\sigma H$."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "sigma = np.linspace(0.0, 1.0, n_layers + 1)\n",
-    "z = ds_static[\"zf\"].values[0]  # interface depths at every node, shape (zf, n_node)\n",
-    "H = z[-1]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "centre_line = np.isclose(node_y, 5e3) \n",
-    "xs = node_x[centre_line] / 1e3\n",
+    "# we create a time array that goes from hour 0 to hour 12\n",
+    "t = np.arange(13) * 3600.0\n",
     "\n",
-    "fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)\n",
-    "ax_sigma.plot(xs, np.broadcast_to(sigma, (xs.size, sigma.size)), \"k\", linewidth=0.8)\n",
-    "ax_sigma.set_ylabel(r\"$\\sigma$\")\n",
-    "ax_sigma.set_title(r\"Layer interfaces in $\\sigma$ space\")\n",
+    "# we construct a depth profile that is 10m deep near shore and logarithmically increases to 100m depth offshore\n",
+    "bottom_depth = 10+90*np.log10(1 + 9*(x_vals/x_range[1]))\n",
+    "bottom_depth = np.broadcast_to(bottom_depth[:, np.newaxis], (nx, nx))\n",
     "\n",
-    "ax_z.plot(xs, z[:, centre_line].T, \"k\", linewidth=0.8)\n",
-    "ax_z.set_ylabel(\"depth [m]\")\n",
-    "ax_z.set_title(\"Layer interfaces in z space\")\n",
+    "# for illustrative purposes we create 2m surface wave that propogates with time\n",
+    "wavelength, wave_period = 20e3, 6 * 3600.0\n",
+    "amplitude = 2\n",
+    "wave = amplitude * np.sin(2.0 * np.pi * (np.broadcast_to(x_vals[:, np.newaxis], (nx, nx)) / wavelength - t[:, np.newaxis, np.newaxis] / wave_period))\n",
     "\n",
-    "for ax in (ax_sigma, ax_z):\n",
-    "    ax.invert_yaxis()\n",
-    "    ax.set_xlabel(\"x [km]\")\n",
-    "plt.tight_layout()\n",
-    "plt.show()"
+    "# construct the dataset itself\n",
+    "ds = sigma_coordinate_lattice_dataset(nx, x_range, y_range, n_vertical_layers + 1, bottom_depth, eta=wave)\n",
+    "ds"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Temporal Variation in $\\eta$\n",
-    "\n",
-    "Next, we add a fluctuating sea surface height field to the vertical grid. For illustrative purposes, we use a large amplitude sea surface height wave that's independent of the velocity field. In sigma space, the vertical grid remains exactly the same, since it accounts for $\\eta$ when calculating the fractional vertical coordinate. In z space however, the change in $\\eta$ causes fluctuations in the grid itself."
+    "We 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 near shore, where the domain is shallower, this demonstrates the strength of a terrain following vertical grid. We can also see that the vertical has time dependency in z space."
    ]
   },
   {
@@ -127,24 +75,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "t = np.arange(13) * 3600.0\n",
+    "time = ds[\"time\"].values.astype(\"datetime64[s]\")\n",
+    "z_moving = ds[\"zf\"].values  # interface depths, shape (time, zf, n_node)\n",
+    "sigma = np.linspace(0.0, 1.0, n_vertical_layers + 1)\n",
     "\n",
-    "# for illustrative purposes we put a 2m surface wave that propogates with time\n",
-    "wavelength, wave_period = 20e3, 6 * 3600.0\n",
-    "wave = 2 * np.sin(2.0 * np.pi * (x_nodes[np.newaxis] / wavelength - t[:, np.newaxis, np.newaxis] / wave_period))\n",
+    "node_x, node_y = ds.uxgrid.node_lon.values, ds.uxgrid.node_lat.values\n",
+    "centre_line = np.isclose(node_y, y_vals[nx // 2])  # the row of nodes closest to the middle of the channel\n",
+    "xs = node_x[centre_line] / 1e3\n",
     "\n",
-    "ds = sigma_coordinate_lattice_dataset(nx, (0.0, 40e3), (0.0, 10e3), n_layers + 1, bottom_depth, eta=wave)\n",
-    "time = ds[\"time\"].values\n",
-    "z_moving = ds[\"zf\"].values  # (time, zf, n_node)\n",
-    "eta = -z_moving[:, 0, :]  # the surface interface sits at depth -eta"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
     "fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)\n",
     "ax_sigma.plot(xs, np.broadcast_to(sigma, (xs.size, sigma.size)), \"k\", linewidth=0.8)\n",
     "ax_sigma.set_ylabel(r\"$\\sigma$\")\n",
@@ -159,11 +97,13 @@
     "    ax.set_xlabel(\"x [km]\")\n",
     "plt.tight_layout()\n",
     "\n",
+    "\n",
     "def animate(snapshot):\n",
     "    for line, interface_depths in zip(z_lines, z_moving[snapshot][:, centre_line]):\n",
     "        line.set_ydata(interface_depths)\n",
     "    z_title.set_text(f\"Layer interfaces in z space at {time[snapshot]}\")\n",
     "\n",
+    "\n",
     "anim = FuncAnimation(fig, animate, frames=time.size, interval=500)\n",
     "plt.close()\n",
     "anim"
@@ -232,6 +172,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "H = z_moving[0, -1]  # the bottom interface sits at the bottom depth, shape (n_node,)\n",
+    "eta = -z_moving[:, 0]  # the surface interface sits at depth -eta, shape (time, n_node)\n",
     "x_line = node_x[centre_line]\n",
     "H0 = np.interp(x0, x_line, H[centre_line])\n",
     "eta0 = np.interp(x0, x_line, eta[0, centre_line])\n",
@@ -294,7 +236,7 @@
     "px = df[\"x\"].to_numpy().reshape(n_frames, n_particles)\n",
     "pz = df[\"z\"].to_numpy().reshape(n_frames, n_particles)\n",
     "\n",
-    "# sigma coordinate of every particle at every output time, on the latest grid snapshot at or before that time\n",
+    "# calcualte the sigma coordinate of every particle at every output time\n",
     "snapshots = np.searchsorted(time, output_times.to_numpy(), side=\"right\") - 1\n",
     "psigma = np.empty_like(pz)\n",
     "for frame, snapshot in enumerate(snapshots):\n",

From 85aa9d41e1ed7b58c5f61ed70c99d80cfe84abe4 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 6 Oct 2026 19:23:44 +0000
Subject: [PATCH 05/22] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 ...orial_unstructured_sigma_coordinates.ipynb | 26 ++++++++++++++-----
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
index 4113683727..ea054b1314 100644
--- a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
+++ b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
@@ -38,8 +38,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "nx = 40 # the number of grid cells in x and y\n",
-    "x_range, y_range = (0.0, 40e3), (0.0, 10e3) # we set the bounds of the domain to be a 40km x 10km channel\n",
+    "nx = 40  # the number of grid cells in x and y\n",
+    "x_range, y_range = (\n",
+    "    (0.0, 40e3),\n",
+    "    (0.0, 10e3),\n",
+    ")  # we set the bounds of the domain to be a 40km x 10km channel\n",
     "n_vertical_layers = 10\n",
     "\n",
     "x_vals = np.linspace(x_range[0], x_range[1], nx)\n",
@@ -49,16 +52,25 @@
     "t = np.arange(13) * 3600.0\n",
     "\n",
     "# we construct a depth profile that is 10m deep near shore and logarithmically increases to 100m depth offshore\n",
-    "bottom_depth = 10+90*np.log10(1 + 9*(x_vals/x_range[1]))\n",
+    "bottom_depth = 10 + 90 * np.log10(1 + 9 * (x_vals / x_range[1]))\n",
     "bottom_depth = np.broadcast_to(bottom_depth[:, np.newaxis], (nx, nx))\n",
     "\n",
     "# for illustrative purposes we create 2m surface wave that propogates with time\n",
     "wavelength, wave_period = 20e3, 6 * 3600.0\n",
     "amplitude = 2\n",
-    "wave = amplitude * np.sin(2.0 * np.pi * (np.broadcast_to(x_vals[:, np.newaxis], (nx, nx)) / wavelength - t[:, np.newaxis, np.newaxis] / wave_period))\n",
+    "wave = amplitude * np.sin(\n",
+    "    2.0\n",
+    "    * np.pi\n",
+    "    * (\n",
+    "        np.broadcast_to(x_vals[:, np.newaxis], (nx, nx)) / wavelength\n",
+    "        - t[:, np.newaxis, np.newaxis] / wave_period\n",
+    "    )\n",
+    ")\n",
     "\n",
     "# construct the dataset itself\n",
-    "ds = sigma_coordinate_lattice_dataset(nx, x_range, y_range, n_vertical_layers + 1, bottom_depth, eta=wave)\n",
+    "ds = sigma_coordinate_lattice_dataset(\n",
+    "    nx, x_range, y_range, n_vertical_layers + 1, bottom_depth, eta=wave\n",
+    ")\n",
     "ds"
    ]
   },
@@ -80,7 +92,9 @@
     "sigma = np.linspace(0.0, 1.0, n_vertical_layers + 1)\n",
     "\n",
     "node_x, node_y = ds.uxgrid.node_lon.values, ds.uxgrid.node_lat.values\n",
-    "centre_line = np.isclose(node_y, y_vals[nx // 2])  # the row of nodes closest to the middle of the channel\n",
+    "centre_line = np.isclose(\n",
+    "    node_y, y_vals[nx // 2]\n",
+    ")  # the row of nodes closest to the middle of the channel\n",
     "xs = node_x[centre_line] / 1e3\n",
     "\n",
     "fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)\n",

From 21c2cf524127d4ee86922c8603e894db8b1b705a Mon Sep 17 00:00:00 2001
From: Wyatt Sieminski <wyatt@fluidnumerics.com>
Date: Thu, 8 Oct 2026 19:18:08 +0000
Subject: [PATCH 06/22] Temporal Linear interpolatation of the z coordinate
 during grid search

---
 ...orial_unstructured_sigma_coordinates.ipynb | 21 ++------------
 src/parcels/_core/basegrid.py                 |  4 ++-
 src/parcels/_core/field.py                    |  2 +-
 src/parcels/_core/uxgrid.py                   | 29 +++++++++++++++----
 src/parcels/_core/xgrid.py                    |  2 +-
 tests/test_uxarray_fieldset.py                | 10 ++++---
 6 files changed, 37 insertions(+), 31 deletions(-)

diff --git a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
index ea054b1314..f1a81aabe2 100644
--- a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
+++ b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
@@ -221,29 +221,14 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": [
-    "## Advect the particles\n",
-    "\n",
-    "We advect the particles for 12 hours using the 3D fourth-order Runge-Kutta kernel `parcels.kernels.AdvectionRK4_3D` with a 5-minute time step, and write their positions every 30 minutes. "
-   ]
+   "source": "## Advect the particles\n\nWe advect the particles for 12 hours using the 3D fourth-order Runge-Kutta kernel `parcels.kernels.AdvectionRK4_3D` with a 5-minute time step, and write their positions every hour. "
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "outputfile = parcels.ParticleFile(\n",
-    "    path=\"sigma_particles.parquet\", outputdt=np.timedelta64(30, \"m\"), mode=\"w\"\n",
-    ")\n",
-    "\n",
-    "pset.execute(\n",
-    "    [parcels.kernels.AdvectionRK4_3D],\n",
-    "    runtime=np.timedelta64(12, \"h\"),\n",
-    "    dt=np.timedelta64(5, \"m\"),\n",
-    "    output_file=outputfile,\n",
-    ")"
-   ]
+   "source": "outputfile = parcels.ParticleFile(\n    path=\"sigma_particles.parquet\", outputdt=np.timedelta64(1, \"h\"), mode=\"w\"\n)\n\npset.execute(\n    [parcels.kernels.AdvectionRK4_3D],\n    runtime=np.timedelta64(12, \"h\"),\n    dt=np.timedelta64(5, \"m\"),\n    output_file=outputfile,\n)"
   },
   {
    "cell_type": "code",
@@ -334,4 +319,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}
\ No newline at end of file
diff --git a/src/parcels/_core/basegrid.py b/src/parcels/_core/basegrid.py
index ac3ae9c664..a0300d595b 100644
--- a/src/parcels/_core/basegrid.py
+++ b/src/parcels/_core/basegrid.py
@@ -29,7 +29,7 @@ class BaseGrid(ABC):
     _mesh: FlatMesh | SphericalMesh
 
     @abstractmethod
-    def search(self, z: float, y: float, x: float, ei=None, ti=None) -> dict[str, tuple[int, float | np.ndarray]]:
+    def search(self, z: float, y: float, x: float, ei=None, ti=None, tau=None) -> dict[str, tuple[int, float | np.ndarray]]:
         """
         Perform a spatial (and optionally vertical) search to locate the grid element
         that contains a given point (x, y, z).
@@ -51,6 +51,8 @@ def search(self, z: float, y: float, x: float, ei=None, ti=None) -> dict[str, tu
             a global or local search strategy.
         ti : np.ndarray, optional
             Time index of each query point, as returned by ``_search_time_index``.
+        tau : np.ndarray, optional
+            Barycentric time coordinate of each query point, as returned by ``_search_time_index``.
         search2D : bool, default=False
             If True, perform only a 2D search (x, y), ignoring the vertical component z.
 
diff --git a/src/parcels/_core/field.py b/src/parcels/_core/field.py
index c8a470248d..c667295138 100644
--- a/src/parcels/_core/field.py
+++ b/src/parcels/_core/field.py
@@ -400,7 +400,7 @@ def _get_positions(field: Field, t, z, y, x, particles, _ei) -> tuple[dict, dict
     grid_positions = {}
     time_positions = _search_time_index(field, t)
     grid_positions.update(time_positions)
-    grid_positions.update(field.grid.search(z, y, x, ei=_ei, ti=time_positions["T"]["index"]))
+    grid_positions.update(field.grid.search(z, y, x, ei=_ei, ti=time_positions["T"]["index"], tau=time_positions["T"]["bcoord"]))
     _update_particles_ei(particles, grid_positions, field)
     _update_particle_states_position(particles, grid_positions)
     return particle_positions, grid_positions
diff --git a/src/parcels/_core/uxgrid.py b/src/parcels/_core/uxgrid.py
index 6c90b04d38..b123f8b13b 100644
--- a/src/parcels/_core/uxgrid.py
+++ b/src/parcels/_core/uxgrid.py
@@ -92,7 +92,7 @@ def deg2m(self) -> float:
             return self._mesh.deg2m
         return 1.0
 
-    def search(self, z, y, x, ei=None, ti=None, tol=1e-6):
+    def search(self, z, y, x, ei=None, ti=None, tau=None, tol=1e-6):
         """
         Search for the grid cell (face) and vertical layer that contains the given points.
 
@@ -111,8 +111,11 @@ def search(self, z, y, x, ei=None, ti=None, tol=1e-6):
             if the points are within the faces specified by these indices. For cells where the particles
             are not found, a nearest neighbor search will be performed. As a last resort, the spatial hash will be used.
         ti : np.ndarray, optional
-            Time index of each point, as returned by ``_search_time_index``. Required when z is 3D; selects the z
-            snapshot used for each particle's column.
+            Time index of each point, as returned by ``_search_time_index``. Required when z is 3D; selects the
+            earlier of the two z snapshots (``ti`` and ``ti + 1``) that are linearly interpolated with ``tau``.
+        tau : np.ndarray, optional
+            Barycentric time coordinate of each point, as returned by ``_search_time_index``. Required when z is 3D;
+            used for linear interpolation of the z coordinate for the construction of a particle's column.
         tol : float, optional
             Tolerance for barycentric coordinate checks. Default is 1e-6.
         """
@@ -142,13 +145,27 @@ def search(self, z, y, x, ei=None, ti=None, tol=1e-6):
 
         found = fi >= 0
         if self.z.ndim == 3:
-            if ti is None:
-                raise ValueError("Searching a UxGrid with a time-varying (3D) z requires the time index ti")
-            node_columns = self.z.isel(
+            if ti is None or tau is None:
+                raise ValueError("Searching a UxGrid with a time-varying (3D) z requires the time index ti and barycentric coordinate tau")
+
+            cols_ti = self.z.isel(
                 time=xr.DataArray(np.broadcast_to(ti, fi.shape)[found], dims="points"),
                 n_node=xr.DataArray(self.uxgrid.face_node_connectivity[fi[found], :].values, dims=("points", "nodes")),
                 ignore_grid=True,
             ).compute()
+
+            if self.z.shape[0] == 1:
+                node_columns = cols_ti
+            else:
+                cols_tnext = self.z.isel(
+                    time=xr.DataArray(np.broadcast_to(ti + 1, fi.shape)[found], dims="points"),
+                    n_node=xr.DataArray(self.uxgrid.face_node_connectivity[fi[found], :].values, dims=("points", "nodes")),
+                    ignore_grid=True,
+                ).compute()
+
+                tau_points = xr.DataArray(np.broadcast_to(tau, fi.shape)[found], dims="points")
+                node_columns = (1 - tau_points) * cols_ti + tau_points * cols_tnext
+
             bcoords = xr.DataArray(coords[found], dims=("points", "nodes"))
             # Particles outside the mesh are given a NaN z column for vertical searching
             columns = np.full((fi.size, self.z.sizes["zf"]), np.nan)
diff --git a/src/parcels/_core/xgrid.py b/src/parcels/_core/xgrid.py
index 5d4c368e5e..1e01650514 100644
--- a/src/parcels/_core/xgrid.py
+++ b/src/parcels/_core/xgrid.py
@@ -313,7 +313,7 @@ def _gtype(self):
             else:
                 return GridType.CurvilinearSGrid
 
-    def search(self, z, y, x, ei=None, ti=None):
+    def search(self, z, y, x, ei=None, ti=None, tau=None):
         ds = self._ds
 
         if "Z" in self.axes:
diff --git a/tests/test_uxarray_fieldset.py b/tests/test_uxarray_fieldset.py
index 7535ed4a26..266e07421c 100644
--- a/tests/test_uxarray_fieldset.py
+++ b/tests/test_uxarray_fieldset.py
@@ -243,8 +243,9 @@ def test_fieldset_from_sigma_coordinate_dataset(eta, expected_n_snapshots):
 def test_sigma_linear_field_is_exact_on_moving_sigma_grid(horizontal_dim, expected_interpolator):
     """A field equal to a + b * sigma on every interface must evaluate to exactly a + b * sigma at any particle.
 
-    Bottom depth and eta are linear in x and y, so barycentric interpolation of the node columns is exact and a
-    particle placed at sigma = s has a known depth.
+    Bottom depth and eta are linear in x and y, and each interface depth is affine in eta, so barycentric
+    interpolation of the node columns and linear interpolation of the columns between snapshots are both exact. A
+    particle placed at sigma = s therefore has a known depth.
     """
     nx, nz = 11, 8
     n_snapshots = 5
@@ -270,9 +271,10 @@ def eta(snapshot, x, y):
     x = rng.uniform(1.0, 10e3 - 1.0, n_particles)
     y = rng.uniform(1.0, 10e3 - 1.0, n_particles)
     t = np.concatenate([[0.0, 3600.0, 7200.0], rng.uniform(0.0, (n_snapshots - 1) * 3600.0, n_particles - 3)])
-    snapshot = np.maximum(np.ceil(t / 3600.0).astype(int) - 1, 0)
     sigma = rng.uniform(0.02, 0.98, n_particles)
-    eta_particles = eta(snapshot, x, y)
+    snapshot = (t // 3600.0).astype(int)
+    tau = t / 3600.0 - snapshot
+    eta_particles = (1 - tau) * eta(snapshot, x, y) + tau * eta(snapshot + 1, x, y)
     z = -eta_particles + sigma * (bottom_depth(x, y) + eta_particles)
 
     np.testing.assert_allclose(fieldset.F.eval(t, z, y, x), a + b * sigma, rtol=1e-6)

From c9bc9135889d652adbba731634ec567a8332d37b Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 8 Oct 2026 19:18:23 +0000
Subject: [PATCH 07/22] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../tutorial_unstructured_sigma_coordinates.ipynb         | 2 +-
 src/parcels/_core/basegrid.py                             | 4 +++-
 src/parcels/_core/field.py                                | 4 +++-
 src/parcels/_core/uxgrid.py                               | 8 ++++++--
 4 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
index f1a81aabe2..86f7093fb3 100644
--- a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
+++ b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
@@ -319,4 +319,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
\ No newline at end of file
+}
diff --git a/src/parcels/_core/basegrid.py b/src/parcels/_core/basegrid.py
index a0300d595b..5998f4a605 100644
--- a/src/parcels/_core/basegrid.py
+++ b/src/parcels/_core/basegrid.py
@@ -29,7 +29,9 @@ class BaseGrid(ABC):
     _mesh: FlatMesh | SphericalMesh
 
     @abstractmethod
-    def search(self, z: float, y: float, x: float, ei=None, ti=None, tau=None) -> dict[str, tuple[int, float | np.ndarray]]:
+    def search(
+        self, z: float, y: float, x: float, ei=None, ti=None, tau=None
+    ) -> dict[str, tuple[int, float | np.ndarray]]:
         """
         Perform a spatial (and optionally vertical) search to locate the grid element
         that contains a given point (x, y, z).
diff --git a/src/parcels/_core/field.py b/src/parcels/_core/field.py
index c667295138..b629ea6243 100644
--- a/src/parcels/_core/field.py
+++ b/src/parcels/_core/field.py
@@ -400,7 +400,9 @@ def _get_positions(field: Field, t, z, y, x, particles, _ei) -> tuple[dict, dict
     grid_positions = {}
     time_positions = _search_time_index(field, t)
     grid_positions.update(time_positions)
-    grid_positions.update(field.grid.search(z, y, x, ei=_ei, ti=time_positions["T"]["index"], tau=time_positions["T"]["bcoord"]))
+    grid_positions.update(
+        field.grid.search(z, y, x, ei=_ei, ti=time_positions["T"]["index"], tau=time_positions["T"]["bcoord"])
+    )
     _update_particles_ei(particles, grid_positions, field)
     _update_particle_states_position(particles, grid_positions)
     return particle_positions, grid_positions
diff --git a/src/parcels/_core/uxgrid.py b/src/parcels/_core/uxgrid.py
index b123f8b13b..c1d50343c6 100644
--- a/src/parcels/_core/uxgrid.py
+++ b/src/parcels/_core/uxgrid.py
@@ -146,7 +146,9 @@ def search(self, z, y, x, ei=None, ti=None, tau=None, tol=1e-6):
         found = fi >= 0
         if self.z.ndim == 3:
             if ti is None or tau is None:
-                raise ValueError("Searching a UxGrid with a time-varying (3D) z requires the time index ti and barycentric coordinate tau")
+                raise ValueError(
+                    "Searching a UxGrid with a time-varying (3D) z requires the time index ti and barycentric coordinate tau"
+                )
 
             cols_ti = self.z.isel(
                 time=xr.DataArray(np.broadcast_to(ti, fi.shape)[found], dims="points"),
@@ -159,7 +161,9 @@ def search(self, z, y, x, ei=None, ti=None, tau=None, tol=1e-6):
             else:
                 cols_tnext = self.z.isel(
                     time=xr.DataArray(np.broadcast_to(ti + 1, fi.shape)[found], dims="points"),
-                    n_node=xr.DataArray(self.uxgrid.face_node_connectivity[fi[found], :].values, dims=("points", "nodes")),
+                    n_node=xr.DataArray(
+                        self.uxgrid.face_node_connectivity[fi[found], :].values, dims=("points", "nodes")
+                    ),
                     ignore_grid=True,
                 ).compute()
 

From dea15dc7a152098b43bf919df0ed351ea0612b29 Mon Sep 17 00:00:00 2001
From: Wyatt Sieminski <wyatt@fluidnumerics.com>
Date: Thu, 8 Oct 2026 20:01:04 +0000
Subject: [PATCH 08/22] Fixed ci failures

---
 ...orial_unstructured_sigma_coordinates.ipynb | 23 +++++++++++++++----
 .../_datasets/unstructured/generated.py       |  6 ++---
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
index 86f7093fb3..cefab39708 100644
--- a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
+++ b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
@@ -221,14 +221,29 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": "## Advect the particles\n\nWe advect the particles for 12 hours using the 3D fourth-order Runge-Kutta kernel `parcels.kernels.AdvectionRK4_3D` with a 5-minute time step, and write their positions every hour. "
+   "source": [
+    "## Advect the particles\n",
+    "\n",
+    "We advect the particles for 12 hours using the 3D fourth-order Runge-Kutta kernel `parcels.kernels.AdvectionRK4_3D` with a 5-minute time step, and write their positions every hour. "
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "outputfile = parcels.ParticleFile(\n    path=\"sigma_particles.parquet\", outputdt=np.timedelta64(1, \"h\"), mode=\"w\"\n)\n\npset.execute(\n    [parcels.kernels.AdvectionRK4_3D],\n    runtime=np.timedelta64(12, \"h\"),\n    dt=np.timedelta64(5, \"m\"),\n    output_file=outputfile,\n)"
+   "source": [
+    "outputfile = parcels.ParticleFile(\n",
+    "    path=\"sigma_particles.parquet\", outputdt=np.timedelta64(1, \"h\"), mode=\"w\"\n",
+    ")\n",
+    "\n",
+    "pset.execute(\n",
+    "    [parcels.kernels.AdvectionRK4_3D],\n",
+    "    runtime=np.timedelta64(12, \"h\"),\n",
+    "    dt=np.timedelta64(5, \"m\"),\n",
+    "    output_file=outputfile,\n",
+    ")"
+   ]
   },
   {
    "cell_type": "code",
@@ -300,9 +315,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python (parcels-pixi)",
+   "display_name": "Parcels:docs (3.14.6)",
    "language": "python",
-   "name": "parcels-pixi"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
diff --git a/src/parcels/_datasets/unstructured/generated.py b/src/parcels/_datasets/unstructured/generated.py
index 68e7c96313..1b57bb7390 100644
--- a/src/parcels/_datasets/unstructured/generated.py
+++ b/src/parcels/_datasets/unstructured/generated.py
@@ -169,8 +169,8 @@ def _build_triangulated_lattice_grid(nx, x_range, y_range):
     """Build a flat UxGrid by splitting each cell of a regular nx-by-nx node lattice over x_range by y_range (in meters) into two triangles."""
     X, Y = np.meshgrid(np.linspace(x_range[0], x_range[1], nx), np.linspace(y_range[0], y_range[1], nx))
     node_x, node_y = X.ravel(), Y.ravel()
-    i, j = np.meshgrid(np.arange(nx - 1), np.arange(nx - 1))
-    i, j = i.ravel(), j.ravel()
+    i_grid, j_grid = np.meshgrid(np.arange(nx - 1), np.arange(nx - 1))
+    i, j = i_grid.ravel(), j_grid.ravel()
     sw = (j * nx) + i
     se = sw + 1
     nw = sw + nx
@@ -391,7 +391,7 @@ def sigma_coordinate_lattice_dataset(nx, x_range, y_range, nz, bottom_depth, eta
     zc = 0.5 * (zf[:-1] + zf[1:]) if zf.ndim == 1 else 0.5 * (zf[:, :-1] + zf[:, 1:])
 
     if zf.ndim == 1:
-        vertical_coords = {"zf": ("zf", zf), "zc": ("zc", zc)}
+        vertical_coords: dict[str, tuple] = {"zf": ("zf", zf), "zc": ("zc", zc)}
     else:
         vertical_coords = {"zf": (("time", "zf", "n_node"), zf), "zc": (("time", "zc", "n_node"), zc)}
 

From debe984fc57816b90ff84a622dd12714ea6dcf3b Mon Sep 17 00:00:00 2001
From: Wyatt Sieminski <wyatt@fluidnumerics.com>
Date: Thu, 8 Oct 2026 20:35:34 +0000
Subject: [PATCH 09/22] Added experimental/memory warning during 3D z
 coordinate fieldset construction

---
 src/parcels/_core/uxgrid.py    | 12 ++++++++++++
 tests/test_uxadvection.py      |  1 +
 tests/test_uxarray_fieldset.py |  3 +++
 tests/test_uxgrid.py           |  1 +
 4 files changed, 17 insertions(+)

diff --git a/src/parcels/_core/uxgrid.py b/src/parcels/_core/uxgrid.py
index c1d50343c6..23415832b3 100644
--- a/src/parcels/_core/uxgrid.py
+++ b/src/parcels/_core/uxgrid.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import warnings
 from typing import TYPE_CHECKING, Literal
 
 import numpy as np
@@ -9,6 +10,7 @@
 from parcels._core.basegrid import BaseGrid
 from parcels._core.index_search import GRID_SEARCH_ERROR, _search_1d_columns, uxgrid_point_in_cell
 from parcels._core.mesh import SphericalMesh, get_mesh
+from parcels._core.warnings import FieldSetWarning
 
 if TYPE_CHECKING:
     import uxarray as ux
@@ -50,6 +52,16 @@ def __init__(
             raise ValueError(f"z must be a 1D or 3D array of vertical coordinates, got {z.ndim}D")
         if z.ndim == 3 and z.dims != ("time", "zf", "n_node"):
             raise ValueError(f"A 3D z must have dims ('time', 'zf', 'n_node'), got {z.dims}")
+        if z.ndim == 3:
+            warnings.warn(
+                "Time-varying (3D) z coordinates are experimental and may cause significant memory overhead that "
+                f"leads to OOM errors. This z coordinate has sizes {dict(z.sizes)} ({z.nbytes / 1e9:.3g} 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.",
+                FieldSetWarning,
+                stacklevel=4,
+            )
         self.z = z
         self._mesh = get_mesh(mesh)
         self._spatialhash = None
diff --git a/tests/test_uxadvection.py b/tests/test_uxadvection.py
index 85ee920768..e9206e75e9 100644
--- a/tests/test_uxadvection.py
+++ b/tests/test_uxadvection.py
@@ -30,6 +30,7 @@ def test_ux_constant_flow_face_centered_2D(integrator, tmp_parquet):
     np.testing.assert_allclose(df["x"].iloc[-1], expected_lon, atol=1e-5)
 
 
+@pytest.mark.filterwarnings("ignore:Time-varying \\(3D\\) z coordinates are experimental")
 def test_ux_advection_on_moving_sigma_grid_keeps_depth():
     """With uniform u and w = 0, particles keep their depth while the sigma levels move past them."""
     nx, nz = 5, 6
diff --git a/tests/test_uxarray_fieldset.py b/tests/test_uxarray_fieldset.py
index 266e07421c..077c4037f7 100644
--- a/tests/test_uxarray_fieldset.py
+++ b/tests/test_uxarray_fieldset.py
@@ -214,6 +214,7 @@ def test_nestedgrids_triangulation_spherical_search():
 RISING_ETA = np.broadcast_to(np.linspace(-1.0, 1.0, 3)[:, np.newaxis, np.newaxis], (3, 5, 5))
 
 
+@pytest.mark.filterwarnings("ignore:Time-varying \\(3D\\) z coordinates are experimental")
 @pytest.mark.parametrize(
     "eta, expected_n_snapshots",
     [
@@ -233,6 +234,7 @@ def test_fieldset_from_sigma_coordinate_dataset(eta, expected_n_snapshots):
         assert isinstance(field.interp_method, UxLinearNodeLinearZF)
 
 
+@pytest.mark.filterwarnings("ignore:Time-varying \\(3D\\) z coordinates are experimental")
 @pytest.mark.parametrize(
     "horizontal_dim, expected_interpolator",
     [
@@ -280,6 +282,7 @@ def eta(snapshot, x, y):
     np.testing.assert_allclose(fieldset.F.eval(t, z, y, x), a + b * sigma, rtol=1e-6)
 
 
+@pytest.mark.filterwarnings("ignore:Time-varying \\(3D\\) z coordinates are experimental")
 def test_unstructured_model_data_rejects_z_with_different_time_coordinate():
     ds = sigma_coordinate_lattice_dataset(5, (0.0, 4e3), (0.0, 4e3), 4, SLOPING_BOTTOM_DEPTH, RISING_ETA)
     ds_one_day_later = sigma_coordinate_lattice_dataset(
diff --git a/tests/test_uxgrid.py b/tests/test_uxgrid.py
index 3d9800734b..8b4dd8ee0f 100644
--- a/tests/test_uxgrid.py
+++ b/tests/test_uxgrid.py
@@ -36,6 +36,7 @@ def test_xgrid_get_axis_dim(uxds):
     assert grid.get_axis_dim("Z") == 2
 
 
+@pytest.mark.filterwarnings("ignore:Time-varying \\(3D\\) z coordinates are experimental")
 def test_uxgrid_search_3d_z_requires_ti():
     ds = sigma_coordinate_lattice_dataset(5, (0.0, 4e3), (0.0, 4e3), 4, np.full((5, 5), 50.0))
     grid = UxGrid(ds.uxgrid, z=ds.coords["zf"], mesh="flat")

From 09fee291e6f6c57207a119be015e8c41cc272bbb Mon Sep 17 00:00:00 2001
From: Wyatt Sieminski <wyatt@fluidnumerics.com>
Date: Fri, 9 Oct 2026 15:32:50 +0000
Subject: [PATCH 10/22] time varying z coordinate support for XGrids

---
 src/parcels/_core/particleset.py              |   8 +-
 src/parcels/_core/xgrid.py                    | 100 ++++++++++++++----
 src/parcels/_datasets/structured/generated.py |  22 ++++
 tests/test_advection.py                       |  23 ++++
 tests/test_fieldset.py                        |  45 ++++++++
 tests/test_uxarray_fieldset.py                |  49 ---------
 tests/test_uxgrid.py                          |  11 --
 7 files changed, 173 insertions(+), 85 deletions(-)

diff --git a/src/parcels/_core/particleset.py b/src/parcels/_core/particleset.py
index d98de7f02f..83e3df44c8 100644
--- a/src/parcels/_core/particleset.py
+++ b/src/parcels/_core/particleset.py
@@ -21,6 +21,7 @@
 )
 from parcels._core.uxgrid import UxGrid
 from parcels._core.warnings import ParticleSetWarning
+from parcels._core.xgrid import XGrid
 from parcels._logger import logger
 
 __all__ = ["ParticleSet"]
@@ -83,9 +84,12 @@ def __init__(
         if z is None:
             minz = None
             for field in self.fieldset.fields.values():
-                if isinstance(field.grid, UxGrid) and field.grid.z.ndim == 3:
+                has_time_varying_z = (isinstance(field.grid, UxGrid) and field.grid.z.ndim == 3) or (
+                    isinstance(field.grid, XGrid) and "Z" in field.grid.axes and field.grid._ds["depth"].ndim == 4
+                )
+                if has_time_varying_z:
                     raise ValueError(
-                        f"Field {field.name!r} has a time-varying (3D) vertical grid, so there is no default "
+                        f"Field {field.name!r} has a time-varying vertical grid, so there is no default "
                         "particle depth. Pass the particle depths explicitly with `z`."
                     )
                 for depth in field.grid.depth:
diff --git a/src/parcels/_core/xgrid.py b/src/parcels/_core/xgrid.py
index 1e01650514..747e9f08a1 100644
--- a/src/parcels/_core/xgrid.py
+++ b/src/parcels/_core/xgrid.py
@@ -1,3 +1,4 @@
+import warnings
 from collections.abc import Hashable, Sequence
 from functools import cached_property
 from typing import Literal, cast
@@ -10,9 +11,11 @@
 import parcels._sgrid as sgrid
 import parcels._typing as ptyping
 from parcels._core.basegrid import BaseGrid
-from parcels._core.index_search import _search_1d_array, _search_indices_curvilinear_2d
+from parcels._core.index_search import _search_1d_array, _search_1d_columns, _search_indices_curvilinear_2d
 from parcels._core.mesh import SphericalMesh, get_mesh
+from parcels._core.warnings import FieldSetWarning
 from parcels._sgrid.accessor import _get_dim_to_axis_mapping, get_dim_position
+from parcels.interpolators._xinterpolators import _get_offsets_dictionary
 
 _FIELD_DATA_ORDERING: Sequence[ptyping.XgcmAxisDirection] = "TZYX"
 _XGRID_AXES_ORDERING: Sequence[ptyping.XgridAxis] = "ZYX"
@@ -135,7 +138,25 @@ def __init__(self, model_data: xr.Dataset, mesh: Literal["flat", "spherical"] |
             assert_valid_lat_lon(ds["lat"], ds["lon"], self.sgrid_metadata)
 
         if "Z" in axes:
-            assert_valid_depth(ds["depth"])
+            depth = ds["depth"]
+            if depth.ndim not in (1, 4):
+                raise ValueError(f"depth must be a 1D or 4D array of vertical coordinates, got {depth.ndim}D")
+            if depth.ndim == 4:
+                assert self.sgrid_metadata.vertical_dimensions is not None
+                vertical_dim = self.sgrid_metadata.vertical_dimensions[0].node
+                fnp_x, fnp_y = self.sgrid_metadata.face_dimensions
+                warnings.warn(
+                    "Time-varying (4D) z coordinates are experimental and may cause significant memory overhead that "
+                    f"leads to OOM errors. This z coordinate has sizes {dict(depth.sizes)} ({depth.nbytes / 1e9:.3g} GB). "
+                    f"Assumptions: z has dims ('time', {vertical_dim!r}, {fnp_y.face!r}, {fnp_x.face!r}), i.e. it is "
+                    f"defined at the layer interfaces on the cell centres, and is strictly increasing along "
+                    f"{vertical_dim!r}; 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.",
+                    FieldSetWarning,
+                    stacklevel=5,
+                )
+            else:
+                assert_valid_depth(depth)
 
         self._ds = ds
 
@@ -316,12 +337,17 @@ def _gtype(self):
     def search(self, z, y, x, ei=None, ti=None, tau=None):
         ds = self._ds
 
-        if "Z" in self.axes:
-            zi, zeta = _search_1d_array(ds.depth.values, z)
-        else:
-            zi, zeta = np.zeros(z.shape, dtype=int), np.zeros(z.shape, dtype=float)
+        if "X" not in self.axes or ds.lon.ndim == 1:
+            if "Y" in self.axes:
+                yi, eta = _search_1d_array(ds.lat.values, y)
+            else:
+                yi, eta = np.zeros(y.shape, dtype=int), np.zeros(y.shape, dtype=float)
 
-        if "X" in self.axes and "Y" in self.axes and ds.lon.ndim == 2:
+            if "X" in self.axes:
+                xi, xsi = _search_1d_array(ds.lon.values, x)
+            else:
+                xi, xsi = np.zeros(x.shape, dtype=int), np.zeros(x.shape, dtype=float)
+        elif ds.lon.ndim == 2:
             yi, xi = None, None
             if ei is not None:
                 axis_indices = self.unravel_index(ei)
@@ -329,25 +355,15 @@ def search(self, z, y, x, ei=None, ti=None, tau=None):
                 yi = axis_indices.get("Y")
 
             yi, eta, xi, xsi = _search_indices_curvilinear_2d(self, y, x, yi, xi)
-
-            return {
-                "Z": {"index": zi, "bcoord": zeta},
-                "Y": {"index": yi, "bcoord": eta},
-                "X": {"index": xi, "bcoord": xsi},
-            }
-
-        if "X" in self.axes and ds.lon.ndim > 2:
-            raise NotImplementedError("Searching in >2D lon/lat arrays is not implemented yet.")
-
-        if "Y" in self.axes:
-            yi, eta = _search_1d_array(ds.lat.values, y)
         else:
-            yi, eta = np.zeros(y.shape, dtype=int), np.zeros(y.shape, dtype=float)
+            raise NotImplementedError("Searching in >2D lon/lat arrays is not implemented yet.")
 
-        if "X" in self.axes:
-            xi, xsi = _search_1d_array(ds.lon.values, x)
+        if "Z" in self.axes and ds.depth.ndim == 4:
+            zi, zeta = self._search_time_varying_depth(z, yi, xi, ti, tau)
+        elif "Z" in self.axes:
+            zi, zeta = _search_1d_array(ds.depth.values, z)
         else:
-            xi, xsi = np.zeros(x.shape, dtype=int), np.zeros(x.shape, dtype=float)
+            zi, zeta = np.zeros(z.shape, dtype=int), np.zeros(z.shape, dtype=float)
 
         return {
             "Z": {"index": zi, "bcoord": zeta},
@@ -355,6 +371,44 @@ def search(self, z, y, x, ei=None, ti=None, tau=None):
             "X": {"index": xi, "bcoord": xsi},
         }
 
+    def _search_time_varying_depth(self, z, yi, xi, ti, tau):
+        """Vertical search of each particle in the z column at the centre of its cell using temporal linear
+        interpolation between the snapshots ti and ti + 1.
+        """
+        if ti is None or tau is None:
+            raise ValueError(
+                "Searching an XGrid with a time-varying (4D) depth requires the time index ti and barycentric coordinate tau"
+            )
+
+        found = (yi >= 0) & (xi >= 0)
+        offsets = _get_offsets_dictionary(self)
+        yi_centre = yi[found] + offsets["Y"]
+        xi_centre = xi[found] + offsets["X"]
+
+        depth = self._ds["depth"]
+        fnp_x, fnp_y = self.sgrid_metadata.face_dimensions
+        cell_centres = {fnp_y.face: xr.DataArray(yi_centre, dims="points"), fnp_x.face: xr.DataArray(xi_centre, dims="points")}
+
+        ti_found = np.broadcast_to(ti, yi.shape)[found]
+        columns_ti = depth.isel({"time": xr.DataArray(ti_found, dims="points"), **cell_centres})
+
+        if depth.sizes["time"] == 1:
+            cell_columns = columns_ti
+        else:
+            columns_tnext = depth.isel({"time": xr.DataArray(ti_found + 1, dims="points"), **cell_centres})
+            tau_found = xr.DataArray(np.broadcast_to(tau, yi.shape)[found], dims="points")
+            cell_columns = (1 - tau_found) * columns_ti + tau_found * columns_tnext
+
+        cell_columns = cell_columns.transpose("points", ...)
+
+        # Particles outside the grid are given a NaN z column for vertical searching
+        columns = np.full((yi.size, cell_columns.shape[1]), np.nan)
+        columns[found] = cell_columns.values
+        zi, zeta = _search_1d_columns(columns, z)
+
+        # Particles outside the grid are given a 0 vertical index and a NaN vertical barycentric coordinate
+        return np.where(found, zi, 0), np.where(found, zeta, np.nan)
+
     @cached_property
     def _fpoint_info(self) -> dict[ptyping.XgridAxis, sgrid.Padding]:
         """Returns a mapping of the spatial axes in the Grid to their Padding values (node positions)."""
diff --git a/src/parcels/_datasets/structured/generated.py b/src/parcels/_datasets/structured/generated.py
index 05ad96dd6e..0913bf857b 100644
--- a/src/parcels/_datasets/structured/generated.py
+++ b/src/parcels/_datasets/structured/generated.py
@@ -39,6 +39,28 @@ def simple_UV_dataset(dims=(360, 2, 30, 4), maxdepth=1, mesh="spherical"):
     )
 
 
+def simple_UV_sigma_dataset(dims, bottom_depth, eta=None, mesh="flat", start_time="2000-01-01"):
+    """simple_UV_dataset with zero U, V and W and a time-varying (4D) depth on the cell centres: dims[1] sigma interfaces from the surface (depth -eta) to bottom_depth ((YC, XC) array, meters); eta is None or (time, YC, XC), hourly from start_time."""
+    n_times, nz, ny, nx = dims
+    if np.shape(bottom_depth) != (ny, nx):
+        raise ValueError(f"bottom_depth must be a (YC, XC) array of shape {(ny, nx)}, got shape {np.shape(bottom_depth)}")
+    if eta is None:
+        eta = np.zeros((n_times, ny, nx))
+    elif np.shape(eta) != (n_times, ny, nx):
+        raise ValueError(f"eta must be a (time, YC, XC) array of shape {(n_times, ny, nx)}, got shape {np.shape(eta)}")
+
+    sigma = np.linspace(0.0, 1.0, nz)
+    total_depth = bottom_depth[np.newaxis, :, :] + eta
+    zf = sigma[np.newaxis, :, np.newaxis, np.newaxis] * total_depth[:, np.newaxis] - eta[:, np.newaxis]  # (time, depth, YC, XC)
+
+    ds = simple_UV_dataset(dims=dims, mesh=mesh)
+    ds["W"] = (["time", "depth", "YG", "XG"], np.zeros(dims))
+    return ds.assign_coords(
+        time=(["time"], xr.date_range(start_time, periods=n_times, freq="1h"), {"axis": "T"}),
+        depth=(["time", "depth", "YC", "XC"], zf, {"axis": "Z"}),
+    )
+
+
 def radial_rotation_dataset(xdim=200, ydim=200):  # Define 2D flat, square fieldset for testing purposes.
     lon = np.linspace(0, 60, xdim, dtype=np.float32)
     lat = np.linspace(0, 60, ydim, dtype=np.float32)
diff --git a/tests/test_advection.py b/tests/test_advection.py
index 451e5e485c..c027946d3a 100644
--- a/tests/test_advection.py
+++ b/tests/test_advection.py
@@ -24,6 +24,7 @@
     peninsula_dataset,
     radial_rotation_dataset,
     simple_UV_dataset,
+    simple_UV_sigma_dataset,
     stommel_gyre_dataset,
 )
 from parcels._datasets.structured.generic import datasets_sgrid
@@ -508,3 +509,25 @@ def test_mitgcm():
         1952691.93845841,
     ]
     np.testing.assert_allclose(pset.y, lat_v3, atol=1)
+
+
+@pytest.mark.filterwarnings("ignore:Time-varying \\(4D\\) z coordinates are experimental")
+def test_advection_on_moving_sigma_grid_keeps_depth():
+    """With uniform u and w = 0, particles keep their depth while the sigma levels move past them."""
+    n_snapshots, nz, ny, nx = 3, 6, 5, 5
+    lon = simple_UV_dataset(dims=(n_snapshots, nz, ny, nx), mesh="flat")["lon"].values
+    x_centres = lon - 0.5 * (lon[1] - lon[0])  # LOW padding: cell centre i lies between nodes i - 1 and i
+    bottom_depth = np.broadcast_to(20.0 + 5e-6 * x_centres, (ny, nx))
+    eta = np.stack([np.broadcast_to(np.sin(snapshot + 1e-6 * x_centres), (ny, nx)) for snapshot in range(n_snapshots)])
+    ds = simple_UV_sigma_dataset((n_snapshots, nz, ny, nx), bottom_depth, eta)
+    u0 = 0.5
+    ds["U"].values[:] = u0
+    fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat")
+
+    x0, z0 = 0.0, np.array([2.0, 8.0, 15.0])
+    pset = ParticleSet(fieldset, x=np.full(z0.size, x0), y=np.zeros(z0.size), z=z0)
+    pset.execute(AdvectionRK4_3D, runtime=np.timedelta64(2, "h"), dt=np.timedelta64(300, "s"), verbose_progress=False)
+
+    np.testing.assert_array_equal(pset.z, z0)
+    np.testing.assert_allclose(pset.x, x0 + u0 * 7200.0)
+    assert np.all(pset.state < StatusCode.Error)
diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py
index 765cefc303..6f9639e7db 100644
--- a/tests/test_fieldset.py
+++ b/tests/test_fieldset.py
@@ -15,8 +15,10 @@
 from parcels._core.fieldset import FieldSet, IncompatibleMeshesException, _datetime_to_msg
 from parcels._core.mesh import SphericalMesh
 from parcels._core.model import _default_vector_field_components
+from parcels._datasets.structured.generated import simple_UV_sigma_dataset
 from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels._datasets.structured.generic import datasets_sgrid
+from parcels._datasets.unstructured.generated import sigma_coordinate_lattice_dataset
 from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
 
 ds = datasets_structured["ds_2d_left"]
@@ -560,3 +562,46 @@ def test_fieldset_describe_backends(tmp_path):
     assert actual == expected
 
     # TODO: Add test for the ChunkedArray backend (can also refactor this test at the same time)
+
+
+@pytest.mark.filterwarnings("ignore:Time-varying \\(4D\\) z coordinates are experimental")
+@pytest.mark.parametrize(
+    "n_snapshots",
+    [
+        pytest.param(1, id="single_snapshot_4d_z"),
+        pytest.param(3, id="time_varying_4d_z"),
+    ],
+)
+def test_fieldset_from_simple_UV_sigma_dataset(n_snapshots):
+    nz, ny, nx = 4, 6, 5
+    eta = np.broadcast_to(np.linspace(-1.0, 1.0, n_snapshots)[:, np.newaxis, np.newaxis], (n_snapshots, ny, nx))
+    ds = simple_UV_sigma_dataset((n_snapshots, nz, ny, nx), np.full((ny, nx), 50.0), eta)
+    fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat")
+
+    depth = fieldset.U.grid._ds["depth"]
+    assert depth.dims == ("time", "depth", "YC", "XC")
+    assert depth.sizes["time"] == n_snapshots
+    assert fieldset.U.grid.get_axis_dim("Z") == nz
+
+
+@pytest.mark.filterwarnings("ignore:Time-varying \\([34]D\\) z coordinates are experimental")
+@pytest.mark.parametrize(
+    "fset_convention, ds",
+    [
+        pytest.param(
+            FieldSet.from_sgrid_conventions,
+            simple_UV_sigma_dataset((1, 4, 6, 5), np.full((6, 5), 50.0)),
+            id="structured",
+        ),
+        pytest.param(
+            FieldSet.from_ugrid_conventions,
+            sigma_coordinate_lattice_dataset(5, (0.0, 4e3), (0.0, 4e3), 4, np.full((5, 5), 50.0)),
+            id="unstructured",
+        ),
+    ],
+)
+def test_search_time_varying_z_requires_ti(fset_convention, ds):
+    grid = fset_convention(ds, mesh="flat").U.grid
+
+    with pytest.raises(ValueError, match="requires the time index ti"):
+        grid.search(np.array([10.0]), np.array([2e3]), np.array([2e3]))
diff --git a/tests/test_uxarray_fieldset.py b/tests/test_uxarray_fieldset.py
index 077c4037f7..ba9ecd74e1 100644
--- a/tests/test_uxarray_fieldset.py
+++ b/tests/test_uxarray_fieldset.py
@@ -16,7 +16,6 @@
 from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
 from parcels.interpolators import (
     UxConstantFaceConstantZC,
-    UxConstantFaceLinearZF,
     UxLinearNodeLinearZF,
 )
 from tests.utils import create_uxgrid_from_triangulation
@@ -234,54 +233,6 @@ def test_fieldset_from_sigma_coordinate_dataset(eta, expected_n_snapshots):
         assert isinstance(field.interp_method, UxLinearNodeLinearZF)
 
 
-@pytest.mark.filterwarnings("ignore:Time-varying \\(3D\\) z coordinates are experimental")
-@pytest.mark.parametrize(
-    "horizontal_dim, expected_interpolator",
-    [
-        pytest.param("n_node", UxLinearNodeLinearZF, id="node_registered"),
-        pytest.param("n_face", UxConstantFaceLinearZF, id="face_registered"),
-    ],
-)
-def test_sigma_linear_field_is_exact_on_moving_sigma_grid(horizontal_dim, expected_interpolator):
-    """A field equal to a + b * sigma on every interface must evaluate to exactly a + b * sigma at any particle.
-
-    Bottom depth and eta are linear in x and y, and each interface depth is affine in eta, so barycentric
-    interpolation of the node columns and linear interpolation of the columns between snapshots are both exact. A
-    particle placed at sigma = s therefore has a known depth.
-    """
-    nx, nz = 11, 8
-    n_snapshots = 5
-    a, b = 0.3, 1.7
-
-    def bottom_depth(x, y):
-        return 20.0 + 4e-3 * x + 2e-3 * y
-
-    def eta(snapshot, x, y):
-        return 0.3 * snapshot + (-1.0) ** snapshot * 1e-4 * x + 5e-5 * y
-
-    x_nodes, y_nodes = np.meshgrid(np.linspace(0.0, 10e3, nx), np.linspace(0.0, 10e3, nx), indexing="ij")
-    eta_nodes = np.stack([eta(snapshot, x_nodes, y_nodes) for snapshot in range(n_snapshots)])
-    ds = sigma_coordinate_lattice_dataset(nx, (0.0, 10e3), (0.0, 10e3), nz, bottom_depth(x_nodes, y_nodes), eta_nodes)
-    field_on_interfaces = a + b * np.linspace(0.0, 1.0, nz)
-    field_shape = (n_snapshots, nz, getattr(ds.uxgrid, horizontal_dim))
-    ds["F"] = (("time", "zf", horizontal_dim), np.broadcast_to(field_on_interfaces[None, :, None], field_shape).copy())
-    fieldset = FieldSet.from_ugrid_conventions(ds, mesh="flat")
-    assert isinstance(fieldset.F.interp_method, expected_interpolator)
-
-    rng = np.random.default_rng(0)
-    n_particles = 100
-    x = rng.uniform(1.0, 10e3 - 1.0, n_particles)
-    y = rng.uniform(1.0, 10e3 - 1.0, n_particles)
-    t = np.concatenate([[0.0, 3600.0, 7200.0], rng.uniform(0.0, (n_snapshots - 1) * 3600.0, n_particles - 3)])
-    sigma = rng.uniform(0.02, 0.98, n_particles)
-    snapshot = (t // 3600.0).astype(int)
-    tau = t / 3600.0 - snapshot
-    eta_particles = (1 - tau) * eta(snapshot, x, y) + tau * eta(snapshot + 1, x, y)
-    z = -eta_particles + sigma * (bottom_depth(x, y) + eta_particles)
-
-    np.testing.assert_allclose(fieldset.F.eval(t, z, y, x), a + b * sigma, rtol=1e-6)
-
-
 @pytest.mark.filterwarnings("ignore:Time-varying \\(3D\\) z coordinates are experimental")
 def test_unstructured_model_data_rejects_z_with_different_time_coordinate():
     ds = sigma_coordinate_lattice_dataset(5, (0.0, 4e3), (0.0, 4e3), 4, SLOPING_BOTTOM_DEPTH, RISING_ETA)
diff --git a/tests/test_uxgrid.py b/tests/test_uxgrid.py
index 8b4dd8ee0f..63df6a3cae 100644
--- a/tests/test_uxgrid.py
+++ b/tests/test_uxgrid.py
@@ -1,8 +1,6 @@
-import numpy as np
 import pytest
 
 from parcels import UxGrid
-from parcels._datasets.unstructured.generated import sigma_coordinate_lattice_dataset
 from parcels._datasets.unstructured.generic import datasets as uxdatasets
 
 GENERIC_Z_COORDS = ["nz", "zf", "depth_2"]
@@ -34,12 +32,3 @@ def test_xgrid_get_axis_dim(uxds):
 
     assert grid.get_axis_dim("FACE") == 721
     assert grid.get_axis_dim("Z") == 2
-
-
-@pytest.mark.filterwarnings("ignore:Time-varying \\(3D\\) z coordinates are experimental")
-def test_uxgrid_search_3d_z_requires_ti():
-    ds = sigma_coordinate_lattice_dataset(5, (0.0, 4e3), (0.0, 4e3), 4, np.full((5, 5), 50.0))
-    grid = UxGrid(ds.uxgrid, z=ds.coords["zf"], mesh="flat")
-
-    with pytest.raises(ValueError, match="requires the time index ti"):
-        grid.search(np.array([10.0]), np.array([2e3]), np.array([2e3]))

From 0701f1cd5cad34bcb7df0fed05bd75c114bdfd04 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 9 Oct 2026 15:33:05 +0000
Subject: [PATCH 11/22] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 src/parcels/_core/xgrid.py                    | 5 ++++-
 src/parcels/_datasets/structured/generated.py | 8 ++++++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/parcels/_core/xgrid.py b/src/parcels/_core/xgrid.py
index 747e9f08a1..50aa5bd565 100644
--- a/src/parcels/_core/xgrid.py
+++ b/src/parcels/_core/xgrid.py
@@ -387,7 +387,10 @@ def _search_time_varying_depth(self, z, yi, xi, ti, tau):
 
         depth = self._ds["depth"]
         fnp_x, fnp_y = self.sgrid_metadata.face_dimensions
-        cell_centres = {fnp_y.face: xr.DataArray(yi_centre, dims="points"), fnp_x.face: xr.DataArray(xi_centre, dims="points")}
+        cell_centres = {
+            fnp_y.face: xr.DataArray(yi_centre, dims="points"),
+            fnp_x.face: xr.DataArray(xi_centre, dims="points"),
+        }
 
         ti_found = np.broadcast_to(ti, yi.shape)[found]
         columns_ti = depth.isel({"time": xr.DataArray(ti_found, dims="points"), **cell_centres})
diff --git a/src/parcels/_datasets/structured/generated.py b/src/parcels/_datasets/structured/generated.py
index 0913bf857b..7beff48511 100644
--- a/src/parcels/_datasets/structured/generated.py
+++ b/src/parcels/_datasets/structured/generated.py
@@ -43,7 +43,9 @@ def simple_UV_sigma_dataset(dims, bottom_depth, eta=None, mesh="flat", start_tim
     """simple_UV_dataset with zero U, V and W and a time-varying (4D) depth on the cell centres: dims[1] sigma interfaces from the surface (depth -eta) to bottom_depth ((YC, XC) array, meters); eta is None or (time, YC, XC), hourly from start_time."""
     n_times, nz, ny, nx = dims
     if np.shape(bottom_depth) != (ny, nx):
-        raise ValueError(f"bottom_depth must be a (YC, XC) array of shape {(ny, nx)}, got shape {np.shape(bottom_depth)}")
+        raise ValueError(
+            f"bottom_depth must be a (YC, XC) array of shape {(ny, nx)}, got shape {np.shape(bottom_depth)}"
+        )
     if eta is None:
         eta = np.zeros((n_times, ny, nx))
     elif np.shape(eta) != (n_times, ny, nx):
@@ -51,7 +53,9 @@ def simple_UV_sigma_dataset(dims, bottom_depth, eta=None, mesh="flat", start_tim
 
     sigma = np.linspace(0.0, 1.0, nz)
     total_depth = bottom_depth[np.newaxis, :, :] + eta
-    zf = sigma[np.newaxis, :, np.newaxis, np.newaxis] * total_depth[:, np.newaxis] - eta[:, np.newaxis]  # (time, depth, YC, XC)
+    zf = (
+        sigma[np.newaxis, :, np.newaxis, np.newaxis] * total_depth[:, np.newaxis] - eta[:, np.newaxis]
+    )  # (time, depth, YC, XC)
 
     ds = simple_UV_dataset(dims=dims, mesh=mesh)
     ds["W"] = (["time", "depth", "YG", "XG"], np.zeros(dims))

From 0ab30228ce9f1b81a67e23cd2f16952c2e175093 Mon Sep 17 00:00:00 2001
From: Wyatt Sieminski <wyatt@fluidnumerics.com>
Date: Fri, 9 Oct 2026 18:42:25 +0000
Subject: [PATCH 12/22] Changed tutorial to include structured grids

---
 .../examples/tutorial_sigma_coordinates.ipynb | 482 ++++++++++++++++++
 ...orial_unstructured_sigma_coordinates.ipynb | 337 ------------
 docs/user_guide/index.md                      |   2 +-
 3 files changed, 483 insertions(+), 338 deletions(-)
 create mode 100644 docs/user_guide/examples/tutorial_sigma_coordinates.ipynb
 delete mode 100644 docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb

diff --git a/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb b/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb
new file mode 100644
index 0000000000..79f41ff4c6
--- /dev/null
+++ b/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb
@@ -0,0 +1,482 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 🖥️ Sigma Coordinate Vertical Grids"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "from matplotlib.animation import FuncAnimation\n",
+    "\n",
+    "import parcels\n",
+    "from parcels._datasets.structured.generated import (\n",
+    "    simple_UV_dataset,\n",
+    "    simple_UV_sigma_dataset,\n",
+    ")\n",
+    "from parcels._datasets.unstructured.generated import sigma_coordinate_lattice_dataset\n",
+    "\n",
+    "plt.rcParams[\"animation.html\"] = \"jshtml\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Building the Dataset\n",
+    "\n",
+    "For this tutorial, we construct our dataset using the parcels built in {py:func}`parcels._datasets.structured.generated.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."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "n_days = 10\n",
+    "n_snapshots = 24 * n_days + 1  # hourly snapshots from hour 0 to the end of day 10\n",
+    "n_vertical_layers = 10\n",
+    "ny, nx = 5, 81  # the number of grid nodes in y and x, so that the cells are 25km wide in x\n",
+    "dims = (n_snapshots, n_vertical_layers + 1, ny, nx)\n",
+    "\n",
+    "lon = simple_UV_dataset(dims=dims, mesh=\"flat\")[\"lon\"].values\n",
+    "dx = lon[1] - lon[0]\n",
+    "x_centres = lon - 0.5 * dx\n",
+    "\n",
+    "# we construct a depth profile that is 100m deep in the west and logarithmically increases to 1000m depth in the east\n",
+    "domain_length = lon[-1] - lon[0]\n",
+    "bottom_depth = 100 + 900 * np.log10(1 + 9 * (x_centres - lon[0]) / domain_length)\n",
+    "bottom_depth = np.broadcast_to(bottom_depth, (ny, nx))\n",
+    "\n",
+    "# for illustrative purposes we create a 50m surface wave with a 1000km wavelength that propagates east with time\n",
+    "t = np.arange(n_snapshots) * 3600.0\n",
+    "wavelength, wave_period = 1000e3, 2 * 86400.0\n",
+    "amplitude = 50\n",
+    "wave = amplitude * np.sin(\n",
+    "    2.0 * np.pi * (x_centres[np.newaxis, :] / wavelength - t[:, np.newaxis] / wave_period)\n",
+    ")\n",
+    "wave = np.broadcast_to(wave[:, np.newaxis, :], (n_snapshots, ny, nx))\n",
+    "\n",
+    "# construct the dataset itself\n",
+    "ds = simple_UV_sigma_dataset(dims, bottom_depth, eta=wave)\n",
+    "ds"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We 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."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "time = ds[\"time\"].values.astype(\"datetime64[s]\")\n",
+    "z_moving = ds[\"depth\"].values  # interface depths, shape (time, depth, YC, XC)\n",
+    "sigma = np.linspace(0.0, 1.0, n_vertical_layers + 1)\n",
+    "\n",
+    "centre_row = ny // 2  # the row of cells in the middle of the domain\n",
+    "xs = x_centres / 1e3\n",
+    "\n",
+    "fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)\n",
+    "ax_sigma.plot(xs, np.broadcast_to(sigma, (xs.size, sigma.size)), \"k\", linewidth=0.8)\n",
+    "ax_sigma.set_ylabel(r\"$\\sigma$\")\n",
+    "ax_sigma.set_title(r\"Layer interfaces in $\\sigma$ space\")\n",
+    "\n",
+    "z_lines = ax_z.plot(\n",
+    "    xs, z_moving[0][:, centre_row].T, \"k\", linewidth=0.8, drawstyle=\"steps-mid\"\n",
+    ")\n",
+    "ax_z.set_ylabel(\"depth [m]\")\n",
+    "z_title = ax_z.set_title(f\"Layer interfaces in z space at {time[0]}\")\n",
+    "\n",
+    "for ax in (ax_sigma, ax_z):\n",
+    "    ax.invert_yaxis()\n",
+    "    ax.set_xlabel(\"x [km]\")\n",
+    "plt.tight_layout()\n",
+    "\n",
+    "\n",
+    "def animate(snapshot):\n",
+    "    for line, interface_depths in zip(z_lines, z_moving[snapshot][:, centre_row]):\n",
+    "        line.set_ydata(interface_depths)\n",
+    "    z_title.set_text(f\"Layer interfaces in z space at {time[snapshot]}\")\n",
+    "\n",
+    "\n",
+    "anim = FuncAnimation(fig, animate, frames=range(0, time.size, 6), interval=300)\n",
+    "plt.close()\n",
+    "anim"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Building the `FieldSet`\n",
+    "\n",
+    "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.\n",
+    "\n",
+    "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."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "u0 = 1.0\n",
+    "ds[\"U\"].values[:] = u0\n",
+    "ds"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fieldset = parcels.FieldSet.from_sgrid_conventions(ds, mesh=\"flat\")\n",
+    "fieldset.describe()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Release the particles\n",
+    "\n",
+    "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."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "n_particles = 10\n",
+    "x0 = np.full(n_particles, -890e3)\n",
+    "y0 = np.zeros(n_particles)\n",
+    "z0 = np.linspace(60.0, 240.0, n_particles)  # depths in m\n",
+    "\n",
+    "pset = parcels.ParticleSet(\n",
+    "    fieldset=fieldset,\n",
+    "    pclass=parcels.Particle,\n",
+    "    t=np.full(n_particles, time[0]),\n",
+    "    z=z0,\n",
+    "    y=y0,\n",
+    "    x=x0,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "H = z_moving[0, -1, centre_row]\n",
+    "eta = -z_moving[:, 0, centre_row] \n",
+    "cell = np.searchsorted(lon, x0)  # index of the cell centre whose z column Parcels uses for each particle\n",
+    "sigma0 = (z0 + eta[0, cell]) / (H[cell] + eta[0, cell])  # initial sigma coordinate of each particle\n",
+    "\n",
+    "colors = [f\"C{p}\" for p in range(n_particles)]\n",
+    "\n",
+    "fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)\n",
+    "ax_sigma.plot(xs, np.broadcast_to(sigma, (xs.size, sigma.size)), \"k\", linewidth=0.8)\n",
+    "ax_sigma.scatter(x0 / 1e3, sigma0, color=colors, zorder=3)\n",
+    "ax_sigma.set_ylabel(r\"$\\sigma$\")\n",
+    "ax_sigma.set_title(r\"Initial particle positions in $\\sigma$ space\")\n",
+    "\n",
+    "ax_z.plot(xs, z_moving[0][:, centre_row].T, \"k\", linewidth=0.8, drawstyle=\"steps-mid\")\n",
+    "ax_z.scatter(x0 / 1e3, z0, color=colors, zorder=3)\n",
+    "ax_z.set_ylabel(\"depth [m]\")\n",
+    "ax_z.set_title(f\"Initial particle positions in z space at {time[0]}\")\n",
+    "\n",
+    "for ax in (ax_sigma, ax_z):\n",
+    "    ax.invert_yaxis()\n",
+    "    ax.set_xlabel(\"x [km]\")\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Advect the particles\n",
+    "\n",
+    "We advect the particles for 10 days using the 3D fourth-order Runge-Kutta kernel {py:func}`parcels.kernels.AdvectionRK4_3D` with a 20-minute time step, and write their positions every 6 hours. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "outputfile = parcels.ParticleFile(\n",
+    "    path=\"sigma_particles.parquet\", outputdt=np.timedelta64(6, \"h\"), mode=\"w\"\n",
+    ")\n",
+    "\n",
+    "pset.execute(\n",
+    "    [parcels.kernels.AdvectionRK4_3D],\n",
+    "    runtime=np.timedelta64(n_days, \"D\"),\n",
+    "    dt=np.timedelta64(20, \"m\"),\n",
+    "    output_file=outputfile,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df = parcels.read_particlefile(\"sigma_particles.parquet\").sort([\"t\", \"particle_id\"])\n",
+    "output_times = df[\"t\"].unique().sort()\n",
+    "n_frames = len(output_times)\n",
+    "px = df[\"x\"].to_numpy().reshape(n_frames, n_particles)\n",
+    "pz = df[\"z\"].to_numpy().reshape(n_frames, n_particles)\n",
+    "\n",
+    "# calculate the sigma coordinate of every particle at every output time, using the z column of its cell\n",
+    "snapshots = np.searchsorted(time, output_times.to_numpy(), side=\"right\") - 1\n",
+    "cells = np.searchsorted(lon, px)\n",
+    "eta_p = eta[snapshots[:, np.newaxis], cells]\n",
+    "psigma = (pz + eta_p) / (H[cells] + eta_p)\n",
+    "\n",
+    "colors = [f\"C{p}\" for p in range(n_particles)]\n",
+    "\n",
+    "fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)\n",
+    "ax_sigma.plot(xs, np.broadcast_to(sigma, (xs.size, sigma.size)), \"k\", linewidth=0.8)\n",
+    "z_lines = ax_z.plot(\n",
+    "    xs, z_moving[0][:, centre_row].T, \"k\", linewidth=0.8, drawstyle=\"steps-mid\"\n",
+    ")\n",
+    "sigma_trails = [ax_sigma.plot([], [], color=c, linewidth=1.5)[0] for c in colors]\n",
+    "z_trails = [ax_z.plot([], [], color=c, linewidth=1.5)[0] for c in colors]\n",
+    "sigma_dots = ax_sigma.scatter(px[0] / 1e3, psigma[0], c=colors, zorder=3)\n",
+    "z_dots = ax_z.scatter(px[0] / 1e3, pz[0], c=colors, zorder=3)\n",
+    "\n",
+    "ax_sigma.set_ylabel(r\"$\\sigma$\")\n",
+    "ax_sigma.set_title(r\"Particles in $\\sigma$ space\")\n",
+    "ax_z.set_ylabel(\"depth [m]\")\n",
+    "z_title = ax_z.set_title(\"\")\n",
+    "for ax in (ax_sigma, ax_z):\n",
+    "    ax.invert_yaxis()\n",
+    "    ax.set_xlabel(\"x [km]\")\n",
+    "plt.tight_layout()\n",
+    "\n",
+    "\n",
+    "def animate(frame):\n",
+    "    for p in range(n_particles):\n",
+    "        sigma_trails[p].set_data(px[: frame + 1, p] / 1e3, psigma[: frame + 1, p])\n",
+    "        z_trails[p].set_data(px[: frame + 1, p] / 1e3, pz[: frame + 1, p])\n",
+    "    sigma_dots.set_offsets(np.c_[px[frame] / 1e3, psigma[frame]])\n",
+    "    z_dots.set_offsets(np.c_[px[frame] / 1e3, pz[frame]])\n",
+    "    for line, interface_depths in zip(z_lines, z_moving[snapshots[frame]][:, centre_row]):\n",
+    "        line.set_ydata(interface_depths)\n",
+    "    z_title.set_text(f\"Particles in z space at {output_times[frame]}\")\n",
+    "\n",
+    "\n",
+    "animate(0)\n",
+    "anim = FuncAnimation(fig, animate, frames=n_frames, interval=300)\n",
+    "plt.close()\n",
+    "anim"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "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."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Sigma coordinates on unstructured grids\n",
+    "\n",
+    "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.\n",
+    "\n",
+    "Below we build the same basin on a triangulated lattice with {py:func}`parcels._datasets.unstructured.generated.sigma_coordinate_lattice_dataset`, evaluating the bathymetry and sea surface height at the lattice nodes, and compare the two z coordinates."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "n_lattice_nodes = 21  # the number of lattice nodes in x and y, so that the nodes are 100km apart\n",
+    "x_nodes = np.linspace(-1000e3, 1000e3, n_lattice_nodes)\n",
+    "\n",
+    "# the same bathymetry and surface wave as above, evaluated at the lattice nodes, with shapes (x, y) and (time, x, y)\n",
+    "bottom_depth_nodes = 100 + 900 * np.log10(1 + 9 * (x_nodes - x_nodes[0]) / domain_length)\n",
+    "bottom_depth_nodes = np.broadcast_to(\n",
+    "    bottom_depth_nodes[:, np.newaxis], (n_lattice_nodes, n_lattice_nodes)\n",
+    ")\n",
+    "wave_nodes = amplitude * np.sin(\n",
+    "    2.0 * np.pi * (x_nodes[np.newaxis, :] / wavelength - t[:, np.newaxis] / wave_period)\n",
+    ")\n",
+    "wave_nodes = np.broadcast_to(\n",
+    "    wave_nodes[:, :, np.newaxis], (n_snapshots, n_lattice_nodes, n_lattice_nodes)\n",
+    ")\n",
+    "\n",
+    "ux_ds = sigma_coordinate_lattice_dataset(\n",
+    "    n_lattice_nodes,\n",
+    "    (-1000e3, 1000e3),\n",
+    "    (-1000e3, 1000e3),\n",
+    "    n_vertical_layers + 1,\n",
+    "    bottom_depth_nodes,\n",
+    "    eta=wave_nodes,\n",
+    ")\n",
+    "print(\"structured z:  \", ds[\"depth\"].dims, ds[\"depth\"].shape)\n",
+    "print(\"unstructured z:\", ux_ds[\"zf\"].dims, ux_ds[\"zf\"].shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ux_ds[\"U\"].values[:] = u0\n",
+    "ux_fieldset = parcels.FieldSet.from_ugrid_conventions(ux_ds, mesh=\"flat\")\n",
+    "ux_fieldset.describe()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "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."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ux_pset = parcels.ParticleSet(\n",
+    "    fieldset=ux_fieldset,\n",
+    "    pclass=parcels.Particle,\n",
+    "    t=np.full(n_particles, time[0]),\n",
+    "    z=z0,\n",
+    "    y=y0,\n",
+    "    x=x0,\n",
+    ")\n",
+    "ux_outputfile = parcels.ParticleFile(\n",
+    "    path=\"ux_sigma_particles.parquet\", outputdt=np.timedelta64(6, \"h\"), mode=\"w\"\n",
+    ")\n",
+    "\n",
+    "ux_pset.execute(\n",
+    "    [parcels.kernels.AdvectionRK4_3D],\n",
+    "    runtime=np.timedelta64(n_days, \"D\"),\n",
+    "    dt=np.timedelta64(20, \"m\"),\n",
+    "    output_file=ux_outputfile,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ux_df = parcels.read_particlefile(\"ux_sigma_particles.parquet\").sort([\"t\", \"particle_id\"])\n",
+    "ux_px = ux_df[\"x\"].to_numpy().reshape(n_frames, n_particles)\n",
+    "ux_pz = ux_df[\"z\"].to_numpy().reshape(n_frames, n_particles)\n",
+    "\n",
+    "# the particles move along the row of lattice nodes at y = 0, i.e. along triangle edges, so the z column that\n",
+    "# Parcels uses for each particle is the linear interpolation in x of the z columns at the nodes of that row\n",
+    "row = np.isclose(ux_ds.uxgrid.node_lat.values, 0.0)\n",
+    "z_nodes = ux_ds[\"zf\"].values[:, :, row]  # interface depths along y = 0, shape (time, zf, x)\n",
+    "H_nodes, eta_nodes = z_nodes[0, -1], -z_nodes[:, 0]\n",
+    "ux_psigma = np.empty_like(ux_pz)\n",
+    "for frame, snapshot in enumerate(snapshots):\n",
+    "    ux_eta_p = np.interp(ux_px[frame], x_nodes, eta_nodes[snapshot])\n",
+    "    ux_psigma[frame] = (ux_pz[frame] + ux_eta_p) / (\n",
+    "        np.interp(ux_px[frame], x_nodes, H_nodes) + ux_eta_p\n",
+    "    )\n",
+    "\n",
+    "fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)\n",
+    "ax_sigma.plot(\n",
+    "    x_nodes / 1e3, np.broadcast_to(sigma, (x_nodes.size, sigma.size)), \"k\", linewidth=0.8\n",
+    ")\n",
+    "z_lines = ax_z.plot(x_nodes / 1e3, z_nodes[0].T, \"k\", linewidth=0.8)\n",
+    "sigma_trails = [ax_sigma.plot([], [], color=c, linewidth=1.5)[0] for c in colors]\n",
+    "z_trails = [ax_z.plot([], [], color=c, linewidth=1.5)[0] for c in colors]\n",
+    "sigma_dots = ax_sigma.scatter(ux_px[0] / 1e3, ux_psigma[0], c=colors, zorder=3)\n",
+    "z_dots = ax_z.scatter(ux_px[0] / 1e3, ux_pz[0], c=colors, zorder=3)\n",
+    "\n",
+    "ax_sigma.set_ylabel(r\"$\\sigma$\")\n",
+    "ax_sigma.set_title(r\"Particles in $\\sigma$ space on the unstructured grid\")\n",
+    "ax_z.set_ylabel(\"depth [m]\")\n",
+    "z_title = ax_z.set_title(\"\")\n",
+    "for ax in (ax_sigma, ax_z):\n",
+    "    ax.invert_yaxis()\n",
+    "    ax.set_xlabel(\"x [km]\")\n",
+    "plt.tight_layout()\n",
+    "\n",
+    "\n",
+    "def animate(frame):\n",
+    "    for p in range(n_particles):\n",
+    "        sigma_trails[p].set_data(ux_px[: frame + 1, p] / 1e3, ux_psigma[: frame + 1, p])\n",
+    "        z_trails[p].set_data(ux_px[: frame + 1, p] / 1e3, ux_pz[: frame + 1, p])\n",
+    "    sigma_dots.set_offsets(np.c_[ux_px[frame] / 1e3, ux_psigma[frame]])\n",
+    "    z_dots.set_offsets(np.c_[ux_px[frame] / 1e3, ux_pz[frame]])\n",
+    "    for line, interface_depths in zip(z_lines, z_nodes[snapshots[frame]]):\n",
+    "        line.set_ydata(interface_depths)\n",
+    "    z_title.set_text(f\"Particles in z space at {output_times[frame]}\")\n",
+    "\n",
+    "\n",
+    "animate(0)\n",
+    "anim = FuncAnimation(fig, animate, frames=n_frames, interval=300)\n",
+    "plt.close()\n",
+    "anim"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "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."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.14.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}
diff --git a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb b/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
deleted file mode 100644
index cefab39708..0000000000
--- a/docs/user_guide/examples/tutorial_unstructured_sigma_coordinates.ipynb
+++ /dev/null
@@ -1,337 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# 🖥️ Sigma Coordinate Vertical Grids on Unstructured Grids"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import matplotlib.pyplot as plt\n",
-    "import numpy as np\n",
-    "from matplotlib.animation import FuncAnimation\n",
-    "\n",
-    "import parcels\n",
-    "from parcels._datasets.unstructured.generated import sigma_coordinate_lattice_dataset\n",
-    "\n",
-    "plt.rcParams[\"animation.html\"] = \"jshtml\""
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Building the Dataset\n",
-    "\n",
-    "For this tutorial, we construct our dataset using the parcels built in {py:func}`parcels._datasets.unstructured.generated.sigma_coordinate_lattice_dataset`, which makes an unstructured horizontally gridded dataset that uses a linear vertical $\\sigma$ grid. 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 height 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:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "nx = 40  # the number of grid cells in x and y\n",
-    "x_range, y_range = (\n",
-    "    (0.0, 40e3),\n",
-    "    (0.0, 10e3),\n",
-    ")  # we set the bounds of the domain to be a 40km x 10km channel\n",
-    "n_vertical_layers = 10\n",
-    "\n",
-    "x_vals = np.linspace(x_range[0], x_range[1], nx)\n",
-    "y_vals = np.linspace(y_range[0], y_range[1], nx)\n",
-    "\n",
-    "# we create a time array that goes from hour 0 to hour 12\n",
-    "t = np.arange(13) * 3600.0\n",
-    "\n",
-    "# we construct a depth profile that is 10m deep near shore and logarithmically increases to 100m depth offshore\n",
-    "bottom_depth = 10 + 90 * np.log10(1 + 9 * (x_vals / x_range[1]))\n",
-    "bottom_depth = np.broadcast_to(bottom_depth[:, np.newaxis], (nx, nx))\n",
-    "\n",
-    "# for illustrative purposes we create 2m surface wave that propogates with time\n",
-    "wavelength, wave_period = 20e3, 6 * 3600.0\n",
-    "amplitude = 2\n",
-    "wave = amplitude * np.sin(\n",
-    "    2.0\n",
-    "    * np.pi\n",
-    "    * (\n",
-    "        np.broadcast_to(x_vals[:, np.newaxis], (nx, nx)) / wavelength\n",
-    "        - t[:, np.newaxis, np.newaxis] / wave_period\n",
-    "    )\n",
-    ")\n",
-    "\n",
-    "# construct the dataset itself\n",
-    "ds = sigma_coordinate_lattice_dataset(\n",
-    "    nx, x_range, y_range, n_vertical_layers + 1, bottom_depth, eta=wave\n",
-    ")\n",
-    "ds"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "We 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 near shore, where the domain is shallower, this demonstrates the strength of a terrain following vertical grid. We can also see that the vertical has time dependency in z space."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "time = ds[\"time\"].values.astype(\"datetime64[s]\")\n",
-    "z_moving = ds[\"zf\"].values  # interface depths, shape (time, zf, n_node)\n",
-    "sigma = np.linspace(0.0, 1.0, n_vertical_layers + 1)\n",
-    "\n",
-    "node_x, node_y = ds.uxgrid.node_lon.values, ds.uxgrid.node_lat.values\n",
-    "centre_line = np.isclose(\n",
-    "    node_y, y_vals[nx // 2]\n",
-    ")  # the row of nodes closest to the middle of the channel\n",
-    "xs = node_x[centre_line] / 1e3\n",
-    "\n",
-    "fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)\n",
-    "ax_sigma.plot(xs, np.broadcast_to(sigma, (xs.size, sigma.size)), \"k\", linewidth=0.8)\n",
-    "ax_sigma.set_ylabel(r\"$\\sigma$\")\n",
-    "ax_sigma.set_title(r\"Layer interfaces in $\\sigma$ space\")\n",
-    "\n",
-    "z_lines = ax_z.plot(xs, z_moving[0][:, centre_line].T, \"k\", linewidth=0.8)\n",
-    "ax_z.set_ylabel(\"depth [m]\")\n",
-    "z_title = ax_z.set_title(f\"Layer interfaces in z space at {time[0]}\")\n",
-    "\n",
-    "for ax in (ax_sigma, ax_z):\n",
-    "    ax.invert_yaxis()\n",
-    "    ax.set_xlabel(\"x [km]\")\n",
-    "plt.tight_layout()\n",
-    "\n",
-    "\n",
-    "def animate(snapshot):\n",
-    "    for line, interface_depths in zip(z_lines, z_moving[snapshot][:, centre_line]):\n",
-    "        line.set_ydata(interface_depths)\n",
-    "    z_title.set_text(f\"Layer interfaces in z space at {time[snapshot]}\")\n",
-    "\n",
-    "\n",
-    "anim = FuncAnimation(fig, animate, frames=time.size, interval=500)\n",
-    "plt.close()\n",
-    "anim"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Building the dataset and `FieldSet`\n",
-    "\n",
-    "We prescribe a uniform along channel velocity of $u_0 = 0.5$ m/s, with $v = w = 0$. All three components sit on the mesh nodes at the layer interfaces (`zf`), so Parcels interpolates them barycentrically in the horizontal and linearly in the vertical.\n",
-    "\n",
-    "To advect the particles on sigma grid, for which the z grid varies in time, we use a 3D `zf` coordinate with dimensions of `(time, zf, n_node)`. Parcels then uses it to find each particle's own column at the current time, before searching that column for the particle's layer. `parcels.FieldSet.from_ugrid_conventions()` also expects a `zc` dimension for the layer centres, so we add the layer-centre depths as well. We pass `mesh=\"flat\"` because the coordinates are Cartesian meters."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "u0 = 0.5\n",
-    "ds[\"U\"].values[:] = u0\n",
-    "ds"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "fieldset = parcels.FieldSet.from_ugrid_conventions(ds, mesh=\"flat\")\n",
-    "fieldset.describe()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Release the particles\n",
-    "\n",
-    "We release 10 particles at the first snapshot near the shallow end of the channel, at $x = 2$ km on the centre line, at depths from 3 m to 22 m. The water column there is about 22.5 m deep, and the shallowest particle starts below the deepest wave trough (2 m), so every particle will remain in the water column throughout the simulation."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "n_particles = 10\n",
-    "x0 = np.full(n_particles, 2e3)\n",
-    "y0 = np.full(n_particles, 5e3)\n",
-    "z0 = np.linspace(3.0, 22.0, n_particles)  # depths in m\n",
-    "\n",
-    "pset = parcels.ParticleSet(\n",
-    "    fieldset=fieldset,\n",
-    "    pclass=parcels.Particle,\n",
-    "    t=np.full(n_particles, time[0]),\n",
-    "    z=z0,\n",
-    "    y=y0,\n",
-    "    x=x0,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "H = z_moving[0, -1]  # the bottom interface sits at the bottom depth, shape (n_node,)\n",
-    "eta = -z_moving[:, 0]  # the surface interface sits at depth -eta, shape (time, n_node)\n",
-    "x_line = node_x[centre_line]\n",
-    "H0 = np.interp(x0, x_line, H[centre_line])\n",
-    "eta0 = np.interp(x0, x_line, eta[0, centre_line])\n",
-    "sigma0 = (z0 + eta0) / (H0 + eta0)  # initial sigma coordinate of each particle\n",
-    "\n",
-    "colors = [f\"C{p}\" for p in range(n_particles)]\n",
-    "\n",
-    "fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)\n",
-    "ax_sigma.plot(xs, np.broadcast_to(sigma, (xs.size, sigma.size)), \"k\", linewidth=0.8)\n",
-    "ax_sigma.scatter(x0 / 1e3, sigma0, color=colors, zorder=3)\n",
-    "ax_sigma.set_ylabel(r\"$\\sigma$\")\n",
-    "ax_sigma.set_title(r\"Initial particle positions in $\\sigma$ space\")\n",
-    "\n",
-    "ax_z.plot(xs, z_moving[0][:, centre_line].T, \"k\", linewidth=0.8)\n",
-    "ax_z.scatter(x0 / 1e3, z0, color=colors, zorder=3)\n",
-    "ax_z.set_ylabel(\"depth [m]\")\n",
-    "ax_z.set_title(f\"Initial particle positions in z space at {time[0]}\")\n",
-    "\n",
-    "for ax in (ax_sigma, ax_z):\n",
-    "    ax.invert_yaxis()\n",
-    "    ax.set_xlabel(\"x [km]\")\n",
-    "plt.tight_layout()\n",
-    "plt.show()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Advect the particles\n",
-    "\n",
-    "We advect the particles for 12 hours using the 3D fourth-order Runge-Kutta kernel `parcels.kernels.AdvectionRK4_3D` with a 5-minute time step, and write their positions every hour. "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "outputfile = parcels.ParticleFile(\n",
-    "    path=\"sigma_particles.parquet\", outputdt=np.timedelta64(1, \"h\"), mode=\"w\"\n",
-    ")\n",
-    "\n",
-    "pset.execute(\n",
-    "    [parcels.kernels.AdvectionRK4_3D],\n",
-    "    runtime=np.timedelta64(12, \"h\"),\n",
-    "    dt=np.timedelta64(5, \"m\"),\n",
-    "    output_file=outputfile,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "df = parcels.read_particlefile(\"sigma_particles.parquet\").sort([\"t\", \"particle_id\"])\n",
-    "output_times = df[\"t\"].unique().sort()\n",
-    "n_frames = len(output_times)\n",
-    "px = df[\"x\"].to_numpy().reshape(n_frames, n_particles)\n",
-    "pz = df[\"z\"].to_numpy().reshape(n_frames, n_particles)\n",
-    "\n",
-    "# calcualte the sigma coordinate of every particle at every output time\n",
-    "snapshots = np.searchsorted(time, output_times.to_numpy(), side=\"right\") - 1\n",
-    "psigma = np.empty_like(pz)\n",
-    "for frame, snapshot in enumerate(snapshots):\n",
-    "    eta_p = np.interp(px[frame], x_line, eta[snapshot, centre_line])\n",
-    "    psigma[frame] = (pz[frame] + eta_p) / (\n",
-    "        np.interp(px[frame], x_line, H[centre_line]) + eta_p\n",
-    "    )\n",
-    "\n",
-    "colors = [f\"C{p}\" for p in range(n_particles)]\n",
-    "\n",
-    "fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)\n",
-    "ax_sigma.plot(xs, np.broadcast_to(sigma, (xs.size, sigma.size)), \"k\", linewidth=0.8)\n",
-    "z_lines = ax_z.plot(xs, z_moving[0][:, centre_line].T, \"k\", linewidth=0.8)\n",
-    "sigma_trails = [ax_sigma.plot([], [], color=c, linewidth=1.5)[0] for c in colors]\n",
-    "z_trails = [ax_z.plot([], [], color=c, linewidth=1.5)[0] for c in colors]\n",
-    "sigma_dots = ax_sigma.scatter(px[0] / 1e3, psigma[0], c=colors, zorder=3)\n",
-    "z_dots = ax_z.scatter(px[0] / 1e3, pz[0], c=colors, zorder=3)\n",
-    "\n",
-    "ax_sigma.set_ylabel(r\"$\\sigma$\")\n",
-    "ax_sigma.set_title(r\"Particles in $\\sigma$ space\")\n",
-    "ax_z.set_ylabel(\"depth [m]\")\n",
-    "z_title = ax_z.set_title(\"\")\n",
-    "for ax in (ax_sigma, ax_z):\n",
-    "    ax.invert_yaxis()\n",
-    "    ax.set_xlabel(\"x [km]\")\n",
-    "plt.tight_layout()\n",
-    "\n",
-    "\n",
-    "def animate(frame):\n",
-    "    for p in range(n_particles):\n",
-    "        sigma_trails[p].set_data(px[: frame + 1, p] / 1e3, psigma[: frame + 1, p])\n",
-    "        z_trails[p].set_data(px[: frame + 1, p] / 1e3, pz[: frame + 1, p])\n",
-    "    sigma_dots.set_offsets(np.c_[px[frame] / 1e3, psigma[frame]])\n",
-    "    z_dots.set_offsets(np.c_[px[frame] / 1e3, pz[frame]])\n",
-    "    for line, interface_depths in zip(\n",
-    "        z_lines, z_moving[snapshots[frame]][:, centre_line]\n",
-    "    ):\n",
-    "        line.set_ydata(interface_depths)\n",
-    "    z_title.set_text(f\"Particles in z space at {output_times[frame]}\")\n",
-    "\n",
-    "\n",
-    "animate(0)\n",
-    "anim = FuncAnimation(fig, animate, frames=n_frames, interval=300)\n",
-    "plt.close()\n",
-    "anim"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "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."
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Parcels:docs (3.14.6)",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.14.6"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 4
-}
diff --git a/docs/user_guide/index.md b/docs/user_guide/index.md
index 17107a5e4a..7e1d38c448 100644
--- a/docs/user_guide/index.md
+++ b/docs/user_guide/index.md
@@ -42,7 +42,6 @@ examples/tutorial_mitgcm.ipynb
 examples/tutorial_delft3d.ipynb
 examples/tutorial_fesom.ipynb
 examples/tutorial_schism.ipynb
-examples/tutorial_unstructured_sigma_coordinates.ipynb
 ```
 
 ## Work with FieldSets
@@ -52,6 +51,7 @@ examples/tutorial_unstructured_sigma_coordinates.ipynb
 :name: work-with-fieldsets
 :titlesonly:
 examples/explanation_grids.ipynb
+examples/tutorial_sigma_coordinates.ipynb
 examples/tutorial_velocityconversion.ipynb
 examples/tutorial_nestedgrids.ipynb
 examples/tutorial_manipulating_field_data.ipynb

From 27b819231f73366b75caa8e1ed51b112f9249df9 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 9 Oct 2026 18:43:31 +0000
Subject: [PATCH 13/22] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../examples/tutorial_sigma_coordinates.ipynb | 44 ++++++++++++++-----
 1 file changed, 33 insertions(+), 11 deletions(-)

diff --git a/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb b/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb
index 79f41ff4c6..4db340e2e7 100644
--- a/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb
+++ b/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb
@@ -45,7 +45,10 @@
     "n_days = 10\n",
     "n_snapshots = 24 * n_days + 1  # hourly snapshots from hour 0 to the end of day 10\n",
     "n_vertical_layers = 10\n",
-    "ny, nx = 5, 81  # the number of grid nodes in y and x, so that the cells are 25km wide in x\n",
+    "ny, nx = (\n",
+    "    5,\n",
+    "    81,\n",
+    ")  # the number of grid nodes in y and x, so that the cells are 25km wide in x\n",
     "dims = (n_snapshots, n_vertical_layers + 1, ny, nx)\n",
     "\n",
     "lon = simple_UV_dataset(dims=dims, mesh=\"flat\")[\"lon\"].values\n",
@@ -62,7 +65,9 @@
     "wavelength, wave_period = 1000e3, 2 * 86400.0\n",
     "amplitude = 50\n",
     "wave = amplitude * np.sin(\n",
-    "    2.0 * np.pi * (x_centres[np.newaxis, :] / wavelength - t[:, np.newaxis] / wave_period)\n",
+    "    2.0\n",
+    "    * np.pi\n",
+    "    * (x_centres[np.newaxis, :] / wavelength - t[:, np.newaxis] / wave_period)\n",
     ")\n",
     "wave = np.broadcast_to(wave[:, np.newaxis, :], (n_snapshots, ny, nx))\n",
     "\n",
@@ -188,9 +193,13 @@
    "outputs": [],
    "source": [
     "H = z_moving[0, -1, centre_row]\n",
-    "eta = -z_moving[:, 0, centre_row] \n",
-    "cell = np.searchsorted(lon, x0)  # index of the cell centre whose z column Parcels uses for each particle\n",
-    "sigma0 = (z0 + eta[0, cell]) / (H[cell] + eta[0, cell])  # initial sigma coordinate of each particle\n",
+    "eta = -z_moving[:, 0, centre_row]\n",
+    "cell = np.searchsorted(\n",
+    "    lon, x0\n",
+    ")  # index of the cell centre whose z column Parcels uses for each particle\n",
+    "sigma0 = (z0 + eta[0, cell]) / (\n",
+    "    H[cell] + eta[0, cell]\n",
+    ")  # initial sigma coordinate of each particle\n",
     "\n",
     "colors = [f\"C{p}\" for p in range(n_particles)]\n",
     "\n",
@@ -285,7 +294,9 @@
     "        z_trails[p].set_data(px[: frame + 1, p] / 1e3, pz[: frame + 1, p])\n",
     "    sigma_dots.set_offsets(np.c_[px[frame] / 1e3, psigma[frame]])\n",
     "    z_dots.set_offsets(np.c_[px[frame] / 1e3, pz[frame]])\n",
-    "    for line, interface_depths in zip(z_lines, z_moving[snapshots[frame]][:, centre_row]):\n",
+    "    for line, interface_depths in zip(\n",
+    "        z_lines, z_moving[snapshots[frame]][:, centre_row]\n",
+    "    ):\n",
     "        line.set_ydata(interface_depths)\n",
     "    z_title.set_text(f\"Particles in z space at {output_times[frame]}\")\n",
     "\n",
@@ -320,11 +331,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "n_lattice_nodes = 21  # the number of lattice nodes in x and y, so that the nodes are 100km apart\n",
+    "n_lattice_nodes = (\n",
+    "    21  # the number of lattice nodes in x and y, so that the nodes are 100km apart\n",
+    ")\n",
     "x_nodes = np.linspace(-1000e3, 1000e3, n_lattice_nodes)\n",
     "\n",
     "# the same bathymetry and surface wave as above, evaluated at the lattice nodes, with shapes (x, y) and (time, x, y)\n",
-    "bottom_depth_nodes = 100 + 900 * np.log10(1 + 9 * (x_nodes - x_nodes[0]) / domain_length)\n",
+    "bottom_depth_nodes = 100 + 900 * np.log10(\n",
+    "    1 + 9 * (x_nodes - x_nodes[0]) / domain_length\n",
+    ")\n",
     "bottom_depth_nodes = np.broadcast_to(\n",
     "    bottom_depth_nodes[:, np.newaxis], (n_lattice_nodes, n_lattice_nodes)\n",
     ")\n",
@@ -397,14 +412,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ux_df = parcels.read_particlefile(\"ux_sigma_particles.parquet\").sort([\"t\", \"particle_id\"])\n",
+    "ux_df = parcels.read_particlefile(\"ux_sigma_particles.parquet\").sort(\n",
+    "    [\"t\", \"particle_id\"]\n",
+    ")\n",
     "ux_px = ux_df[\"x\"].to_numpy().reshape(n_frames, n_particles)\n",
     "ux_pz = ux_df[\"z\"].to_numpy().reshape(n_frames, n_particles)\n",
     "\n",
     "# the particles move along the row of lattice nodes at y = 0, i.e. along triangle edges, so the z column that\n",
     "# Parcels uses for each particle is the linear interpolation in x of the z columns at the nodes of that row\n",
     "row = np.isclose(ux_ds.uxgrid.node_lat.values, 0.0)\n",
-    "z_nodes = ux_ds[\"zf\"].values[:, :, row]  # interface depths along y = 0, shape (time, zf, x)\n",
+    "z_nodes = ux_ds[\"zf\"].values[\n",
+    "    :, :, row\n",
+    "]  # interface depths along y = 0, shape (time, zf, x)\n",
     "H_nodes, eta_nodes = z_nodes[0, -1], -z_nodes[:, 0]\n",
     "ux_psigma = np.empty_like(ux_pz)\n",
     "for frame, snapshot in enumerate(snapshots):\n",
@@ -415,7 +434,10 @@
     "\n",
     "fig, (ax_sigma, ax_z) = plt.subplots(1, 2, figsize=(12, 4), sharex=True)\n",
     "ax_sigma.plot(\n",
-    "    x_nodes / 1e3, np.broadcast_to(sigma, (x_nodes.size, sigma.size)), \"k\", linewidth=0.8\n",
+    "    x_nodes / 1e3,\n",
+    "    np.broadcast_to(sigma, (x_nodes.size, sigma.size)),\n",
+    "    \"k\",\n",
+    "    linewidth=0.8,\n",
     ")\n",
     "z_lines = ax_z.plot(x_nodes / 1e3, z_nodes[0].T, \"k\", linewidth=0.8)\n",
     "sigma_trails = [ax_sigma.plot([], [], color=c, linewidth=1.5)[0] for c in colors]\n",

From 096033dd918ad4268d65107950884a2cdcc1562b Mon Sep 17 00:00:00 2001
From: Wyatt Sieminski <wyatt@fluidnumerics.com>
Date: Fri, 9 Oct 2026 19:05:47 +0000
Subject: [PATCH 14/22] Enforce correct dimensions for Z on XGrid construction

---
 src/parcels/_core/xgrid.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/parcels/_core/xgrid.py b/src/parcels/_core/xgrid.py
index 50aa5bd565..45e354a3b5 100644
--- a/src/parcels/_core/xgrid.py
+++ b/src/parcels/_core/xgrid.py
@@ -145,10 +145,13 @@ def __init__(self, model_data: xr.Dataset, mesh: Literal["flat", "spherical"] |
                 assert self.sgrid_metadata.vertical_dimensions is not None
                 vertical_dim = self.sgrid_metadata.vertical_dimensions[0].node
                 fnp_x, fnp_y = self.sgrid_metadata.face_dimensions
+                expected_dims = ("time", vertical_dim, fnp_y.face, fnp_x.face)
+                if depth.dims != expected_dims:
+                    raise ValueError(f"A 4D depth must have dims {expected_dims}, got {depth.dims}")
                 warnings.warn(
                     "Time-varying (4D) z coordinates are experimental and may cause significant memory overhead that "
                     f"leads to OOM errors. This z coordinate has sizes {dict(depth.sizes)} ({depth.nbytes / 1e9:.3g} GB). "
-                    f"Assumptions: z has dims ('time', {vertical_dim!r}, {fnp_y.face!r}, {fnp_x.face!r}), i.e. it is "
+                    f"Assumptions: z has dims {expected_dims}, i.e. it is "
                     f"defined at the layer interfaces on the cell centres, and is strictly increasing along "
                     f"{vertical_dim!r}; 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.",

From de187632d8c6ec80f821f774b1070d480cac45b0 Mon Sep 17 00:00:00 2001
From: Wyatt Sieminski <wyatt@fluidnumerics.com>
Date: Tue, 13 Oct 2026 13:48:34 +0000
Subject: [PATCH 15/22] Switched tutorial kernel to 2D RK2 and bumped dt to
 30min

---
 .../examples/tutorial_sigma_coordinates.ipynb          | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb b/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb
index 4db340e2e7..c3434e0443 100644
--- a/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb
+++ b/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb
@@ -227,7 +227,7 @@
    "source": [
     "## Advect the particles\n",
     "\n",
-    "We advect the particles for 10 days using the 3D fourth-order Runge-Kutta kernel {py:func}`parcels.kernels.AdvectionRK4_3D` with a 20-minute time step, and write their positions every 6 hours. "
+    "We advect the particles for 10 days using the 2D second-order Runge-Kutta kernel {py:func}`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."
    ]
   },
   {
@@ -241,9 +241,9 @@
     ")\n",
     "\n",
     "pset.execute(\n",
-    "    [parcels.kernels.AdvectionRK4_3D],\n",
+    "    [parcels.kernels.AdvectionRK2],\n",
     "    runtime=np.timedelta64(n_days, \"D\"),\n",
-    "    dt=np.timedelta64(20, \"m\"),\n",
+    "    dt=np.timedelta64(30, \"m\"),\n",
     "    output_file=outputfile,\n",
     ")"
    ]
@@ -399,9 +399,9 @@
     ")\n",
     "\n",
     "ux_pset.execute(\n",
-    "    [parcels.kernels.AdvectionRK4_3D],\n",
+    "    [parcels.kernels.AdvectionRK2],\n",
     "    runtime=np.timedelta64(n_days, \"D\"),\n",
-    "    dt=np.timedelta64(20, \"m\"),\n",
+    "    dt=np.timedelta64(30, \"m\"),\n",
     "    output_file=ux_outputfile,\n",
     ")"
    ]

From f6c79f9645e8be64d7498413d15b1116380034c1 Mon Sep 17 00:00:00 2001
From: Wyatt Sieminski <wyatt@fluidnumerics.com>
Date: Tue, 13 Oct 2026 11:53:08 -0400
Subject: [PATCH 16/22] change wording of particle vertical position in
 ParticleSet error

Co-authored-by: Erik van Sebille <e.vansebille@uu.nl>
---
 src/parcels/_core/particleset.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/parcels/_core/particleset.py b/src/parcels/_core/particleset.py
index 83e3df44c8..72a4cc32a8 100644
--- a/src/parcels/_core/particleset.py
+++ b/src/parcels/_core/particleset.py
@@ -90,7 +90,7 @@ def __init__(
                 if has_time_varying_z:
                     raise ValueError(
                         f"Field {field.name!r} has a time-varying vertical grid, so there is no default "
-                        "particle depth. Pass the particle depths explicitly with `z`."
+                        "particle.z. Pass the particle vertical start locations explicitly with `z`."
                     )
                 for depth in field.grid.depth:
                     if minz is None or np.abs(depth) < np.abs(minz):

From 35464a353881d4c7282432755ff22637675e67db Mon Sep 17 00:00:00 2001
From: Wyatt Sieminski <wyatt@fluidnumerics.com>
Date: Tue, 13 Oct 2026 11:55:06 -0400
Subject: [PATCH 17/22] simplify simple_UV_sigma_dataset call signature

Co-authored-by: Erik van Sebille <e.vansebille@uu.nl>
---
 src/parcels/_datasets/structured/generated.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/parcels/_datasets/structured/generated.py b/src/parcels/_datasets/structured/generated.py
index 7beff48511..27bb26aa81 100644
--- a/src/parcels/_datasets/structured/generated.py
+++ b/src/parcels/_datasets/structured/generated.py
@@ -39,9 +39,9 @@ def simple_UV_dataset(dims=(360, 2, 30, 4), maxdepth=1, mesh="spherical"):
     )
 
 
-def simple_UV_sigma_dataset(dims, bottom_depth, eta=None, mesh="flat", start_time="2000-01-01"):
+def simple_UV_sigma_dataset(n_times, nz, ny, nx, bottom_depth, eta=None, mesh="flat", start_time="2000-01-01"):
     """simple_UV_dataset with zero U, V and W and a time-varying (4D) depth on the cell centres: dims[1] sigma interfaces from the surface (depth -eta) to bottom_depth ((YC, XC) array, meters); eta is None or (time, YC, XC), hourly from start_time."""
-    n_times, nz, ny, nx = dims
+
     if np.shape(bottom_depth) != (ny, nx):
         raise ValueError(
             f"bottom_depth must be a (YC, XC) array of shape {(ny, nx)}, got shape {np.shape(bottom_depth)}"

From 2bac6e2a6e141313ffb6bd91e63207b5bd76c76b Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 13 Oct 2026 15:55:21 +0000
Subject: [PATCH 18/22] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 src/parcels/_datasets/structured/generated.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/parcels/_datasets/structured/generated.py b/src/parcels/_datasets/structured/generated.py
index 27bb26aa81..7c62166560 100644
--- a/src/parcels/_datasets/structured/generated.py
+++ b/src/parcels/_datasets/structured/generated.py
@@ -41,7 +41,6 @@ def simple_UV_dataset(dims=(360, 2, 30, 4), maxdepth=1, mesh="spherical"):
 
 def simple_UV_sigma_dataset(n_times, nz, ny, nx, bottom_depth, eta=None, mesh="flat", start_time="2000-01-01"):
     """simple_UV_dataset with zero U, V and W and a time-varying (4D) depth on the cell centres: dims[1] sigma interfaces from the surface (depth -eta) to bottom_depth ((YC, XC) array, meters); eta is None or (time, YC, XC), hourly from start_time."""
-
     if np.shape(bottom_depth) != (ny, nx):
         raise ValueError(
             f"bottom_depth must be a (YC, XC) array of shape {(ny, nx)}, got shape {np.shape(bottom_depth)}"

From 40ed761225930d52c8764aeba9ccc54b93076962 Mon Sep 17 00:00:00 2001
From: Wyatt Sieminski <wyatt@fluidnumerics.com>
Date: Tue, 13 Oct 2026 16:32:09 +0000
Subject: [PATCH 19/22] Changed simple xgrid dataset generators to use explicit
 dims rather than dims tuple. Removed ti/tau warning from grid.search
 functions.

---
 .../examples/tutorial_diffusion.ipynb         | 11 +-----
 .../examples/tutorial_interaction.ipynb       | 12 +-----
 .../examples/tutorial_interpolation.ipynb     | 14 +------
 .../examples/tutorial_nestedgrids.ipynb       | 15 +-------
 .../examples/tutorial_sigma_coordinates.ipynb | 37 +------------------
 .../examples/tutorial_statuscodes.md          |  2 +-
 .../tutorial_velocityconversion.ipynb         | 16 +-------
 src/parcels/_core/uxgrid.py                   | 11 ++----
 src/parcels/_core/xgrid.py                    |  5 ---
 src/parcels/_datasets/structured/generated.py | 33 ++++++++++-------
 tests/test_advection.py                       |  6 +--
 tests/test_backends.py                        |  2 +-
 tests/test_diffusion.py                       |  4 +-
 tests/test_fieldset.py                        | 26 +------------
 tests/test_kernel.py                          |  2 +-
 tests/test_particleset.py                     |  4 +-
 tests/test_particleset_execute.py             |  2 +-
 tests/utils.py                                |  4 +-
 18 files changed, 48 insertions(+), 158 deletions(-)

diff --git a/docs/user_guide/examples/tutorial_diffusion.ipynb b/docs/user_guide/examples/tutorial_diffusion.ipynb
index 8f52efcf54..279766619e 100644
--- a/docs/user_guide/examples/tutorial_diffusion.ipynb
+++ b/docs/user_guide/examples/tutorial_diffusion.ipynb
@@ -187,14 +187,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from parcels._datasets.structured.generated import simple_UV_dataset\n",
-    "\n",
-    "ds = simple_UV_dataset(dims=(1, 1, Ny, 1), mesh=\"flat\")\n",
-    "ds[\"lat\"][:] = np.linspace(-0.01, 1.01, Ny)\n",
-    "ds[\"Kh_meridional\"] = ([\"YG\", \"XG\"], Kh_meridional[:, None])\n",
-    "ds"
-   ]
+   "source": "from parcels._datasets.structured.generated import simple_UV_dataset\n\nds = simple_UV_dataset(n_times=1, nz=1, ny=Ny, nx=1, mesh=\"flat\")\nds[\"lat\"][:] = np.linspace(-0.01, 1.01, Ny)\nds[\"Kh_meridional\"] = ([\"YG\", \"XG\"], Kh_meridional[:, None])\nds"
   },
   {
    "cell_type": "code",
@@ -638,4 +631,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}
\ No newline at end of file
diff --git a/docs/user_guide/examples/tutorial_interaction.ipynb b/docs/user_guide/examples/tutorial_interaction.ipynb
index 713ee6d862..7bb721136a 100644
--- a/docs/user_guide/examples/tutorial_interaction.ipynb
+++ b/docs/user_guide/examples/tutorial_interaction.ipynb
@@ -96,15 +96,7 @@
    "id": "5",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "def DiffusionFieldSet():\n",
-    "    \"\"\"Define a fieldset with only diffusion\"\"\"\n",
-    "    ds = simple_UV_dataset(dims=(1, 1, 1, 1), mesh=\"flat\")\n",
-    "    fieldset = parcels.FieldSet.from_sgrid_conventions(ds, mesh=\"flat\")\n",
-    "    fieldset.add_constant_field(\"Kh_zonal\", 0.0005)\n",
-    "    fieldset.add_constant_field(\"Kh_meridional\", 0.0005)\n",
-    "    return fieldset"
-   ]
+   "source": "def DiffusionFieldSet():\n    \"\"\"Define a fieldset with only diffusion\"\"\"\n    ds = simple_UV_dataset(n_times=1, nz=1, ny=1, nx=1, mesh=\"flat\")\n    fieldset = parcels.FieldSet.from_sgrid_conventions(ds, mesh=\"flat\")\n    fieldset.add_constant_field(\"Kh_zonal\", 0.0005)\n    fieldset.add_constant_field(\"Kh_meridional\", 0.0005)\n    return fieldset"
   },
   {
    "cell_type": "code",
@@ -380,4 +372,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}
\ No newline at end of file
diff --git a/docs/user_guide/examples/tutorial_interpolation.ipynb b/docs/user_guide/examples/tutorial_interpolation.ipynb
index a5ac9088a9..0ef76edee0 100644
--- a/docs/user_guide/examples/tutorial_interpolation.ipynb
+++ b/docs/user_guide/examples/tutorial_interpolation.ipynb
@@ -43,17 +43,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from parcels._datasets.structured.generated import simple_UV_dataset\n",
-    "\n",
-    "ds = simple_UV_dataset(dims=(1, 1, 5, 4), mesh=\"flat\")\n",
-    "ds[\"lat\"][:] = np.linspace(0.0, 1.0, len(ds.YG))\n",
-    "ds[\"lon\"][:] = np.linspace(0.0, 1.0, len(ds.XG))\n",
-    "dx, dy = 1.0 / len(ds.XG), 1.0 / len(ds.YG)\n",
-    "ds[\"P\"] = ds[\"U\"] + np.random.rand(1, 1, 5, 4) + 0.1\n",
-    "ds[\"P\"][:, :, 1, 1] = 0\n",
-    "ds"
-   ]
+   "source": "from parcels._datasets.structured.generated import simple_UV_dataset\n\nds = simple_UV_dataset(n_times=1, nz=1, ny=5, nx=4, mesh=\"flat\")\nds[\"lat\"][:] = np.linspace(0.0, 1.0, len(ds.YG))\nds[\"lon\"][:] = np.linspace(0.0, 1.0, len(ds.XG))\ndx, dy = 1.0 / len(ds.XG), 1.0 / len(ds.YG)\nds[\"P\"] = ds[\"U\"] + np.random.rand(1, 1, 5, 4) + 0.1\nds[\"P\"][:, :, 1, 1] = 0\nds"
   },
   {
    "cell_type": "markdown",
@@ -499,4 +489,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}
\ No newline at end of file
diff --git a/docs/user_guide/examples/tutorial_nestedgrids.ipynb b/docs/user_guide/examples/tutorial_nestedgrids.ipynb
index 838a712392..37201eec14 100644
--- a/docs/user_guide/examples/tutorial_nestedgrids.ipynb
+++ b/docs/user_guide/examples/tutorial_nestedgrids.ipynb
@@ -58,18 +58,7 @@
    "id": "4",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "ds_in = []\n",
-    "grid_polygons = []\n",
-    "\n",
-    "\n",
-    "def setup_nested_ds(polygon):\n",
-    "    xdim, ydim = 20, 20\n",
-    "    ds = simple_UV_dataset(dims=(1, 1, ydim, xdim), mesh=\"spherical\")\n",
-    "    ds[\"lon\"][:] = np.linspace(polygon[:, 0].min(), polygon[:, 0].max(), xdim)\n",
-    "    ds[\"lat\"][:] = np.linspace(polygon[:, 1].min(), polygon[:, 1].max(), ydim)\n",
-    "    return ds"
-   ]
+   "source": "ds_in = []\ngrid_polygons = []\n\n\ndef setup_nested_ds(polygon):\n    xdim, ydim = 20, 20\n    ds = simple_UV_dataset(n_times=1, nz=1, ny=ydim, nx=xdim, mesh=\"spherical\")\n    ds[\"lon\"][:] = np.linspace(polygon[:, 0].min(), polygon[:, 0].max(), xdim)\n    ds[\"lat\"][:] = np.linspace(polygon[:, 1].min(), polygon[:, 1].max(), ydim)\n    return ds"
   },
   {
    "cell_type": "markdown",
@@ -559,4 +548,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}
\ No newline at end of file
diff --git a/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb b/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb
index c3434e0443..0adace8c85 100644
--- a/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb
+++ b/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb
@@ -41,40 +41,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "n_days = 10\n",
-    "n_snapshots = 24 * n_days + 1  # hourly snapshots from hour 0 to the end of day 10\n",
-    "n_vertical_layers = 10\n",
-    "ny, nx = (\n",
-    "    5,\n",
-    "    81,\n",
-    ")  # the number of grid nodes in y and x, so that the cells are 25km wide in x\n",
-    "dims = (n_snapshots, n_vertical_layers + 1, ny, nx)\n",
-    "\n",
-    "lon = simple_UV_dataset(dims=dims, mesh=\"flat\")[\"lon\"].values\n",
-    "dx = lon[1] - lon[0]\n",
-    "x_centres = lon - 0.5 * dx\n",
-    "\n",
-    "# we construct a depth profile that is 100m deep in the west and logarithmically increases to 1000m depth in the east\n",
-    "domain_length = lon[-1] - lon[0]\n",
-    "bottom_depth = 100 + 900 * np.log10(1 + 9 * (x_centres - lon[0]) / domain_length)\n",
-    "bottom_depth = np.broadcast_to(bottom_depth, (ny, nx))\n",
-    "\n",
-    "# for illustrative purposes we create a 50m surface wave with a 1000km wavelength that propagates east with time\n",
-    "t = np.arange(n_snapshots) * 3600.0\n",
-    "wavelength, wave_period = 1000e3, 2 * 86400.0\n",
-    "amplitude = 50\n",
-    "wave = amplitude * np.sin(\n",
-    "    2.0\n",
-    "    * np.pi\n",
-    "    * (x_centres[np.newaxis, :] / wavelength - t[:, np.newaxis] / wave_period)\n",
-    ")\n",
-    "wave = np.broadcast_to(wave[:, np.newaxis, :], (n_snapshots, ny, nx))\n",
-    "\n",
-    "# construct the dataset itself\n",
-    "ds = simple_UV_sigma_dataset(dims, bottom_depth, eta=wave)\n",
-    "ds"
-   ]
+   "source": "n_days = 10\nn_snapshots = 24 * n_days + 1  # hourly snapshots from hour 0 to the end of day 10\nn_vertical_layers = 10\nny, nx = (\n    5,\n    81,\n)  # the number of grid nodes in y and x, so that the cells are 25km wide in x\n\nlon = simple_UV_dataset(\n    n_times=n_snapshots, nz=n_vertical_layers + 1, ny=ny, nx=nx, mesh=\"flat\"\n)[\"lon\"].values\ndx = lon[1] - lon[0]\nx_centres = lon - 0.5 * dx\n\n# we construct a depth profile that is 100m deep in the west and logarithmically increases to 1000m depth in the east\ndomain_length = lon[-1] - lon[0]\nbottom_depth = 100 + 900 * np.log10(1 + 9 * (x_centres - lon[0]) / domain_length)\nbottom_depth = np.broadcast_to(bottom_depth, (ny, nx))\n\n# for illustrative purposes we create a 50m surface wave with a 1000km wavelength that propagates east with time\nt = np.arange(n_snapshots) * 3600.0\nwavelength, wave_period = 1000e3, 2 * 86400.0\namplitude = 50\nwave = amplitude * np.sin(\n    2.0\n    * np.pi\n    * (x_centres[np.newaxis, :] / wavelength - t[:, np.newaxis] / wave_period)\n)\nwave = np.broadcast_to(wave[:, np.newaxis, :], (n_snapshots, ny, nx))\n\n# construct the dataset itself\nds = simple_UV_sigma_dataset(\n    n_times=n_snapshots,\n    nz=n_vertical_layers + 1,\n    ny=ny,\n    nx=nx,\n    bottom_depth=bottom_depth,\n    eta=wave,\n)\nds"
   },
   {
    "cell_type": "markdown",
@@ -501,4 +468,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}
\ No newline at end of file
diff --git a/docs/user_guide/examples/tutorial_statuscodes.md b/docs/user_guide/examples/tutorial_statuscodes.md
index 7888b79d70..38c018bc4a 100644
--- a/docs/user_guide/examples/tutorial_statuscodes.md
+++ b/docs/user_guide/examples/tutorial_statuscodes.md
@@ -63,7 +63,7 @@ Let's add the `KeepInOcean` Kernel to an particle simulation where particles mov
 import numpy as np
 from parcels._datasets.structured.generated import simple_UV_dataset
 
-ds = simple_UV_dataset(dims=(1, 2, 5, 4), mesh="flat")
+ds = simple_UV_dataset(n_times=1, nz=2, ny=5, nx=4, mesh="flat")
 
 dx, dy = 1.0 / len(ds.XG), 1.0 / len(ds.YG)
 
diff --git a/docs/user_guide/examples/tutorial_velocityconversion.ipynb b/docs/user_guide/examples/tutorial_velocityconversion.ipynb
index ee5b32c805..015e6c678c 100644
--- a/docs/user_guide/examples/tutorial_velocityconversion.ipynb
+++ b/docs/user_guide/examples/tutorial_velocityconversion.ipynb
@@ -46,19 +46,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from parcels._datasets.structured.generated import simple_UV_dataset\n",
-    "\n",
-    "nlat = 10\n",
-    "nlon = 18\n",
-    "ds = simple_UV_dataset(dims=(1, 1, nlat, nlon), mesh=\"spherical\")\n",
-    "ds[\"temperature\"] = ds[\"U\"] + 20  # add temperature field of 20 deg\n",
-    "ds[\"U\"].data[:] = 1.0  # set U to 1 m/s\n",
-    "ds[\"V\"].data[:] = 1.0  # set V to 1 m/s\n",
-    "display(ds)\n",
-    "\n",
-    "fieldset = parcels.FieldSet.from_sgrid_conventions(ds, mesh=\"spherical\")"
-   ]
+   "source": "from parcels._datasets.structured.generated import simple_UV_dataset\n\nnlat = 10\nnlon = 18\nds = simple_UV_dataset(n_times=1, nz=1, ny=nlat, nx=nlon, mesh=\"spherical\")\nds[\"temperature\"] = ds[\"U\"] + 20  # add temperature field of 20 deg\nds[\"U\"].data[:] = 1.0  # set U to 1 m/s\nds[\"V\"].data[:] = 1.0  # set V to 1 m/s\ndisplay(ds)\n\nfieldset = parcels.FieldSet.from_sgrid_conventions(ds, mesh=\"spherical\")"
   },
   {
    "attachments": {},
@@ -210,4 +198,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file
diff --git a/src/parcels/_core/uxgrid.py b/src/parcels/_core/uxgrid.py
index 23415832b3..b40f10bfeb 100644
--- a/src/parcels/_core/uxgrid.py
+++ b/src/parcels/_core/uxgrid.py
@@ -8,7 +8,7 @@
 from dask import is_dask_collection
 
 from parcels._core.basegrid import BaseGrid
-from parcels._core.index_search import GRID_SEARCH_ERROR, _search_1d_columns, uxgrid_point_in_cell
+from parcels._core.index_search import GRID_SEARCH_ERROR, _search_1d_array, _search_1d_columns, uxgrid_point_in_cell
 from parcels._core.mesh import SphericalMesh, get_mesh
 from parcels._core.warnings import FieldSetWarning
 
@@ -157,11 +157,6 @@ def search(self, z, y, x, ei=None, ti=None, tau=None, tol=1e-6):
 
         found = fi >= 0
         if self.z.ndim == 3:
-            if ti is None or tau is None:
-                raise ValueError(
-                    "Searching a UxGrid with a time-varying (3D) z requires the time index ti and barycentric coordinate tau"
-                )
-
             cols_ti = self.z.isel(
                 time=xr.DataArray(np.broadcast_to(ti, fi.shape)[found], dims="points"),
                 n_node=xr.DataArray(self.uxgrid.face_node_connectivity[fi[found], :].values, dims=("points", "nodes")),
@@ -186,10 +181,10 @@ def search(self, z, y, x, ei=None, ti=None, tau=None, tol=1e-6):
             # Particles outside the mesh are given a NaN z column for vertical searching
             columns = np.full((fi.size, self.z.sizes["zf"]), np.nan)
             columns[found] = xr.dot(node_columns, bcoords, dim="nodes").transpose("points", "zf").values
+            zi, zeta = _search_1d_columns(columns, z)
         else:
-            columns = np.broadcast_to(self.z.values, (z.size, self.z.size))
+            zi, zeta = _search_1d_array(self.z.values, z)
 
-        zi, zeta = _search_1d_columns(columns, z)
         # Particles outside the mesh are given a 0 vertical index and a NaN vertical barycentric coordinate
         zi = np.where(found, zi, 0)
         zeta = np.where(found, zeta, np.nan)
diff --git a/src/parcels/_core/xgrid.py b/src/parcels/_core/xgrid.py
index 45e354a3b5..da9859821b 100644
--- a/src/parcels/_core/xgrid.py
+++ b/src/parcels/_core/xgrid.py
@@ -378,11 +378,6 @@ def _search_time_varying_depth(self, z, yi, xi, ti, tau):
         """Vertical search of each particle in the z column at the centre of its cell using temporal linear
         interpolation between the snapshots ti and ti + 1.
         """
-        if ti is None or tau is None:
-            raise ValueError(
-                "Searching an XGrid with a time-varying (4D) depth requires the time index ti and barycentric coordinate tau"
-            )
-
         found = (yi >= 0) & (xi >= 0)
         offsets = _get_offsets_dictionary(self)
         yi_centre = yi[found] + offsets["Y"]
diff --git a/src/parcels/_datasets/structured/generated.py b/src/parcels/_datasets/structured/generated.py
index 7c62166560..f3c21f52a1 100644
--- a/src/parcels/_datasets/structured/generated.py
+++ b/src/parcels/_datasets/structured/generated.py
@@ -7,21 +7,22 @@
 from parcels._core.utils.time import timedelta_to_float
 
 
-def simple_UV_dataset(dims=(360, 2, 30, 4), maxdepth=1, mesh="spherical"):
+def simple_UV_dataset(n_times=360, nz=2, ny=30, nx=4, maxdepth=1, mesh="spherical"):
     max_lon = 180.0 if mesh == "spherical" else 1e6
     max_lat = 90.0 if mesh == "spherical" else 1e6
+    shape = (n_times, nz, ny, nx)
 
     return xr.Dataset(
-        {"U": (["time", "depth", "YG", "XG"], np.zeros(dims)), "V": (["time", "depth", "YG", "XG"], np.zeros(dims))},
+        {"U": (["time", "depth", "YG", "XG"], np.zeros(shape)), "V": (["time", "depth", "YG", "XG"], np.zeros(shape))},
         coords={
-            "time": (["time"], xr.date_range("2000", "2001", dims[0]), {"axis": "T"}),
-            "depth": (["depth"], np.linspace(0, maxdepth, dims[1]), {"axis": "Z"}),
-            "YC": (["YC"], np.arange(dims[2]) + 0.5, {"axis": "Y"}),
-            "YG": (["YG"], np.arange(dims[2]), {"axis": "Y", "c_grid_axis_shift": -0.5}),
-            "XC": (["XC"], np.arange(dims[3]) + 0.5, {"axis": "X"}),
-            "XG": (["XG"], np.arange(dims[3]), {"axis": "X", "c_grid_axis_shift": -0.5}),
-            "lat": (["YG"], np.linspace(-max_lat, max_lat, dims[2]), {"axis": "Y", "c_grid_axis_shift": -0.5}),
-            "lon": (["XG"], np.linspace(-max_lon, max_lon, dims[3]), {"axis": "X", "c_grid_axis_shift": -0.5}),
+            "time": (["time"], xr.date_range("2000", "2001", n_times), {"axis": "T"}),
+            "depth": (["depth"], np.linspace(0, maxdepth, nz), {"axis": "Z"}),
+            "YC": (["YC"], np.arange(ny) + 0.5, {"axis": "Y"}),
+            "YG": (["YG"], np.arange(ny), {"axis": "Y", "c_grid_axis_shift": -0.5}),
+            "XC": (["XC"], np.arange(nx) + 0.5, {"axis": "X"}),
+            "XG": (["XG"], np.arange(nx), {"axis": "X", "c_grid_axis_shift": -0.5}),
+            "lat": (["YG"], np.linspace(-max_lat, max_lat, ny), {"axis": "Y", "c_grid_axis_shift": -0.5}),
+            "lon": (["XG"], np.linspace(-max_lon, max_lon, nx), {"axis": "X", "c_grid_axis_shift": -0.5}),
         },
     ).pipe(
         sgrid._attach_sgrid_metadata,
@@ -39,8 +40,12 @@ def simple_UV_dataset(dims=(360, 2, 30, 4), maxdepth=1, mesh="spherical"):
     )
 
 
-def simple_UV_sigma_dataset(n_times, nz, ny, nx, bottom_depth, eta=None, mesh="flat", start_time="2000-01-01"):
-    """simple_UV_dataset with zero U, V and W and a time-varying (4D) depth on the cell centres: dims[1] sigma interfaces from the surface (depth -eta) to bottom_depth ((YC, XC) array, meters); eta is None or (time, YC, XC), hourly from start_time."""
+def simple_UV_sigma_dataset(
+    n_times=360, nz=2, ny=30, nx=4, bottom_depth=None, eta=None, mesh="flat", start_time="2000-01-01"
+):
+    """simple_UV_dataset with zero U, V and W and a time-varying (4D) depth on the cell centres: nz sigma interfaces from the surface (depth -eta) to bottom_depth ((YC, XC) array, meters; None for a flat bottom at 1 m); eta is None or (time, YC, XC), hourly from start_time."""
+    if bottom_depth is None:
+        bottom_depth = np.ones((ny, nx))
     if np.shape(bottom_depth) != (ny, nx):
         raise ValueError(
             f"bottom_depth must be a (YC, XC) array of shape {(ny, nx)}, got shape {np.shape(bottom_depth)}"
@@ -56,8 +61,8 @@ def simple_UV_sigma_dataset(n_times, nz, ny, nx, bottom_depth, eta=None, mesh="f
         sigma[np.newaxis, :, np.newaxis, np.newaxis] * total_depth[:, np.newaxis] - eta[:, np.newaxis]
     )  # (time, depth, YC, XC)
 
-    ds = simple_UV_dataset(dims=dims, mesh=mesh)
-    ds["W"] = (["time", "depth", "YG", "XG"], np.zeros(dims))
+    ds = simple_UV_dataset(n_times=n_times, nz=nz, ny=ny, nx=nx, mesh=mesh)
+    ds["W"] = (["time", "depth", "YG", "XG"], np.zeros((n_times, nz, ny, nx)))
     return ds.assign_coords(
         time=(["time"], xr.date_range(start_time, periods=n_times, freq="1h"), {"axis": "T"}),
         depth=(["time", "depth", "YC", "XC"], zf, {"axis": "Z"}),
diff --git a/tests/test_advection.py b/tests/test_advection.py
index c027946d3a..332801ac64 100644
--- a/tests/test_advection.py
+++ b/tests/test_advection.py
@@ -88,7 +88,7 @@ def periodicBC(particles, fieldset):
 
 @ignore_kernel_warnings
 def test_advection_zonal_periodic():
-    ds = simple_UV_dataset(dims=(2, 2, 2, 2), mesh="flat")
+    ds = simple_UV_dataset(n_times=2, nz=2, ny=2, nx=2, mesh="flat")
     ds["U"].data[:] = 0.1
     ds["lon"].data = np.array([0, 2])
     ds["lat"].data = np.array([0, 2])
@@ -515,11 +515,11 @@ def test_mitgcm():
 def test_advection_on_moving_sigma_grid_keeps_depth():
     """With uniform u and w = 0, particles keep their depth while the sigma levels move past them."""
     n_snapshots, nz, ny, nx = 3, 6, 5, 5
-    lon = simple_UV_dataset(dims=(n_snapshots, nz, ny, nx), mesh="flat")["lon"].values
+    lon = simple_UV_dataset(n_times=n_snapshots, nz=nz, ny=ny, nx=nx, mesh="flat")["lon"].values
     x_centres = lon - 0.5 * (lon[1] - lon[0])  # LOW padding: cell centre i lies between nodes i - 1 and i
     bottom_depth = np.broadcast_to(20.0 + 5e-6 * x_centres, (ny, nx))
     eta = np.stack([np.broadcast_to(np.sin(snapshot + 1e-6 * x_centres), (ny, nx)) for snapshot in range(n_snapshots)])
-    ds = simple_UV_sigma_dataset((n_snapshots, nz, ny, nx), bottom_depth, eta)
+    ds = simple_UV_sigma_dataset(n_snapshots, nz, ny, nx, bottom_depth, eta)
     u0 = 0.5
     ds["U"].values[:] = u0
     fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat")
diff --git a/tests/test_backends.py b/tests/test_backends.py
index 3019b5a78c..f5085bb434 100644
--- a/tests/test_backends.py
+++ b/tests/test_backends.py
@@ -130,7 +130,7 @@ def test_field_without_depth_identical_across_backends(backend: BackendT):
     """Fields without a depth dimension get a size-1 ``mockZ`` dimension, which must
     also be indexable by every backend (gh-2897).
     """
-    ds = simple_UV_dataset(dims=(10, 1, 20, 20), mesh="flat").isel(depth=0, drop=True)
+    ds = simple_UV_dataset(n_times=10, nz=1, ny=20, nx=20, mesh="flat").isel(depth=0, drop=True)
     ds["U"].data[:] = 1.0 + ds["XG"].values / 20  # spatially varying, so the corners matter
     ds["V"].data[:] = 0.5
 
diff --git a/tests/test_diffusion.py b/tests/test_diffusion.py
index d199a7ae0f..0d26031a13 100644
--- a/tests/test_diffusion.py
+++ b/tests/test_diffusion.py
@@ -22,7 +22,7 @@ def test_fieldKh_Brownian(mesh):
     kh_meridional = 50
     mesh_conversion = 1 / 1852.0 / 60 if mesh == "spherical" else 1
 
-    ds = simple_UV_dataset(dims=(2, 1, 2, 2), mesh=mesh)
+    ds = simple_UV_dataset(n_times=2, nz=1, ny=2, nx=2, mesh=mesh)
     ds["lon"].data = np.array([-1e6, 1e6])
     ds["lat"].data = np.array([-1e6, 1e6])
     fieldset = FieldSet.from_sgrid_conventions(ds, mesh=mesh)
@@ -53,7 +53,7 @@ def test_fieldKh_SpatiallyVaryingDiffusion(mesh, kernel):
     ydim, xdim = 100, 200
 
     mesh_conversion = 1 / 1852.0 / 60 if mesh == "spherical" else 1
-    ds = simple_UV_dataset(dims=(2, 1, ydim, xdim), mesh=mesh)
+    ds = simple_UV_dataset(n_times=2, nz=1, ny=ydim, nx=xdim, mesh=mesh)
     ds["lon"].data = np.linspace(-1e6, 1e6, xdim)
     ds["lat"].data = np.linspace(-1e6, 1e6, ydim)
 
diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py
index 6f9639e7db..4cc1bd658b 100644
--- a/tests/test_fieldset.py
+++ b/tests/test_fieldset.py
@@ -18,7 +18,6 @@
 from parcels._datasets.structured.generated import simple_UV_sigma_dataset
 from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels._datasets.structured.generic import datasets_sgrid
-from parcels._datasets.unstructured.generated import sigma_coordinate_lattice_dataset
 from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
 
 ds = datasets_structured["ds_2d_left"]
@@ -575,33 +574,10 @@ def test_fieldset_describe_backends(tmp_path):
 def test_fieldset_from_simple_UV_sigma_dataset(n_snapshots):
     nz, ny, nx = 4, 6, 5
     eta = np.broadcast_to(np.linspace(-1.0, 1.0, n_snapshots)[:, np.newaxis, np.newaxis], (n_snapshots, ny, nx))
-    ds = simple_UV_sigma_dataset((n_snapshots, nz, ny, nx), np.full((ny, nx), 50.0), eta)
+    ds = simple_UV_sigma_dataset(n_snapshots, nz, ny, nx, np.full((ny, nx), 50.0), eta)
     fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat")
 
     depth = fieldset.U.grid._ds["depth"]
     assert depth.dims == ("time", "depth", "YC", "XC")
     assert depth.sizes["time"] == n_snapshots
     assert fieldset.U.grid.get_axis_dim("Z") == nz
-
-
-@pytest.mark.filterwarnings("ignore:Time-varying \\([34]D\\) z coordinates are experimental")
-@pytest.mark.parametrize(
-    "fset_convention, ds",
-    [
-        pytest.param(
-            FieldSet.from_sgrid_conventions,
-            simple_UV_sigma_dataset((1, 4, 6, 5), np.full((6, 5), 50.0)),
-            id="structured",
-        ),
-        pytest.param(
-            FieldSet.from_ugrid_conventions,
-            sigma_coordinate_lattice_dataset(5, (0.0, 4e3), (0.0, 4e3), 4, np.full((5, 5), 50.0)),
-            id="unstructured",
-        ),
-    ],
-)
-def test_search_time_varying_z_requires_ti(fset_convention, ds):
-    grid = fset_convention(ds, mesh="flat").U.grid
-
-    with pytest.raises(ValueError, match="requires the time index ti"):
-        grid.search(np.array([10.0]), np.array([2e3]), np.array([2e3]))
diff --git a/tests/test_kernel.py b/tests/test_kernel.py
index 292deaffe2..1b1cb30012 100644
--- a/tests/test_kernel.py
+++ b/tests/test_kernel.py
@@ -130,7 +130,7 @@ def test_rk45_kernel_warnings(fieldset):
 @ignore_kernel_warnings
 @pytest.mark.parametrize("kernel_type", ["update_lon", "update_dlon"])
 def test_execution_order(kernel_type):
-    ds = simple_UV_dataset(dims=(1, 1, 2, 2), mesh="flat")
+    ds = simple_UV_dataset(n_times=1, nz=1, ny=2, nx=2, mesh="flat")
     ds["U"].data[:, :] = [[0, 1], [2, 3]]
     ds["lon"].data = [0, 2]
     fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat")
diff --git a/tests/test_particleset.py b/tests/test_particleset.py
index 0cfc35e155..d7db107af8 100644
--- a/tests/test_particleset.py
+++ b/tests/test_particleset.py
@@ -193,7 +193,7 @@ def test_pset_iterator(fieldset):
     ],
 )
 def test_pset_default_z_is_in_domain(depths):
-    ds = simple_UV_dataset(dims=(1, len(depths), 10, 10), mesh="flat")
+    ds = simple_UV_dataset(n_times=1, nz=len(depths), ny=10, nx=10, mesh="flat")
     ds = ds.assign_coords(depth=depths)
     fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat")
 
@@ -211,7 +211,7 @@ def test_pset_default_z_is_in_domain(depths):
     ],
 )
 def test_pset_default_z_closest_to_zero(depths):
-    ds = simple_UV_dataset(dims=(1, len(depths), 10, 10), mesh="flat")
+    ds = simple_UV_dataset(n_times=1, nz=len(depths), ny=10, nx=10, mesh="flat")
     ds = ds.assign_coords(depth=depths)
     fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat")
 
diff --git a/tests/test_particleset_execute.py b/tests/test_particleset_execute.py
index dad4601cc2..305b955077 100644
--- a/tests/test_particleset_execute.py
+++ b/tests/test_particleset_execute.py
@@ -56,7 +56,7 @@ def time_varying_zonal_flow_fieldset() -> FieldSet:
     clock and nothing else.
     """
     nt = 25
-    ds = simple_UV_dataset(dims=(nt, 2, 6, 6), mesh="flat")
+    ds = simple_UV_dataset(n_times=nt, nz=2, ny=6, nx=6, mesh="flat")
     times = np.array([np.timedelta64(3 * i, "h") for i in range(nt)])
     ds["time"] = ("time", times, {"axis": "T"})
     u = np.cos(2 * np.pi * (times / np.timedelta64(1, "s")) / 86400.0)  # 1 m/s, 1-day period
diff --git a/tests/utils.py b/tests/utils.py
index a649d36d49..f30f34f39a 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -77,7 +77,7 @@ def create_fieldset_global(xdim=200, ydim=100):
 def create_fieldset_zeros_conversion(mesh="spherical", xdim=200, ydim=100) -> FieldSet:
     """Zero velocity field with lat and lon determined by a conversion factor."""
     mesh_conversion = 1 / 1852.0 / 60 if mesh == "spherical" else 1
-    ds = simple_UV_dataset(dims=(2, 1, ydim, xdim), mesh=mesh)
+    ds = simple_UV_dataset(n_times=2, nz=1, ny=ydim, nx=xdim, mesh=mesh)
     ds["lon"].data = np.linspace(-1e6 * mesh_conversion, 1e6 * mesh_conversion, xdim)
     ds["lat"].data = np.linspace(-1e6 * mesh_conversion, 1e6 * mesh_conversion, ydim)
     return FieldSet.from_sgrid_conventions(ds, mesh=mesh)
@@ -229,7 +229,7 @@ def create_lonlat_grid(grid_type, face_deg, centre=(0.0, 0.0), n=8, mesh="spheri
     if grid_type == "uxgrid":
         grid = create_uxgrid_from_triangulation(nodes[:, 0], nodes[:, 1], faces, mesh=mesh)
     else:
-        ds = simple_UV_dataset(dims=(2, 2, n + 1, n + 1), mesh=mesh).assign_coords(
+        ds = simple_UV_dataset(n_times=2, nz=2, ny=n + 1, nx=n + 1, mesh=mesh).assign_coords(
             lon=(("YG", "XG"), grid_lon), lat=(("YG", "XG"), grid_lat)
         )
         grid = FieldSet.from_sgrid_conventions(ds, mesh=mesh).U.grid

From 3f474367c4a96801c7d1d371244c8c0a36fdb8e3 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 13 Oct 2026 16:32:29 +0000
Subject: [PATCH 20/22] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 docs/user_guide/examples/tutorial_diffusion.ipynb          | 2 +-
 docs/user_guide/examples/tutorial_interaction.ipynb        | 2 +-
 docs/user_guide/examples/tutorial_interpolation.ipynb      | 2 +-
 docs/user_guide/examples/tutorial_nestedgrids.ipynb        | 2 +-
 docs/user_guide/examples/tutorial_sigma_coordinates.ipynb  | 2 +-
 docs/user_guide/examples/tutorial_velocityconversion.ipynb | 2 +-
 6 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/docs/user_guide/examples/tutorial_diffusion.ipynb b/docs/user_guide/examples/tutorial_diffusion.ipynb
index 279766619e..7e2eb210b6 100644
--- a/docs/user_guide/examples/tutorial_diffusion.ipynb
+++ b/docs/user_guide/examples/tutorial_diffusion.ipynb
@@ -631,4 +631,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
\ No newline at end of file
+}
diff --git a/docs/user_guide/examples/tutorial_interaction.ipynb b/docs/user_guide/examples/tutorial_interaction.ipynb
index 7bb721136a..d952185221 100644
--- a/docs/user_guide/examples/tutorial_interaction.ipynb
+++ b/docs/user_guide/examples/tutorial_interaction.ipynb
@@ -372,4 +372,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
\ No newline at end of file
+}
diff --git a/docs/user_guide/examples/tutorial_interpolation.ipynb b/docs/user_guide/examples/tutorial_interpolation.ipynb
index 0ef76edee0..de0f88eceb 100644
--- a/docs/user_guide/examples/tutorial_interpolation.ipynb
+++ b/docs/user_guide/examples/tutorial_interpolation.ipynb
@@ -489,4 +489,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
\ No newline at end of file
+}
diff --git a/docs/user_guide/examples/tutorial_nestedgrids.ipynb b/docs/user_guide/examples/tutorial_nestedgrids.ipynb
index 37201eec14..38933ba216 100644
--- a/docs/user_guide/examples/tutorial_nestedgrids.ipynb
+++ b/docs/user_guide/examples/tutorial_nestedgrids.ipynb
@@ -548,4 +548,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
\ No newline at end of file
+}
diff --git a/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb b/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb
index 0adace8c85..1c615d3507 100644
--- a/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb
+++ b/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb
@@ -468,4 +468,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
\ No newline at end of file
+}
diff --git a/docs/user_guide/examples/tutorial_velocityconversion.ipynb b/docs/user_guide/examples/tutorial_velocityconversion.ipynb
index 015e6c678c..7be7552c6f 100644
--- a/docs/user_guide/examples/tutorial_velocityconversion.ipynb
+++ b/docs/user_guide/examples/tutorial_velocityconversion.ipynb
@@ -198,4 +198,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
\ No newline at end of file
+}

From b2b44c0597c66d21fa7964a80f719543de6b8492 Mon Sep 17 00:00:00 2001
From: Wyatt Sieminski <wyatt@fluidnumerics.com>
Date: Tue, 13 Oct 2026 17:05:06 +0000
Subject: [PATCH 21/22] Moved uxgrid time varying vertical grid search to a
 helper function

---
 src/parcels/_core/uxgrid.py | 59 +++++++++++++++++++------------------
 1 file changed, 31 insertions(+), 28 deletions(-)

diff --git a/src/parcels/_core/uxgrid.py b/src/parcels/_core/uxgrid.py
index b40f10bfeb..d6db846dc5 100644
--- a/src/parcels/_core/uxgrid.py
+++ b/src/parcels/_core/uxgrid.py
@@ -155,38 +155,41 @@ def search(self, z, y, x, ei=None, ti=None, tau=None, tol=1e-6):
             coords[zero_indices, :] = coords_q
             fi[zero_indices] = face_ids_q
 
-        found = fi >= 0
         if self.z.ndim == 3:
-            cols_ti = self.z.isel(
-                time=xr.DataArray(np.broadcast_to(ti, fi.shape)[found], dims="points"),
+            zi, zeta = self._search_time_varying_z(z, fi, coords, ti, tau)
+        else:
+            zi, zeta = _search_1d_array(self.z.values, z)
+
+        return {"Z": {"index": zi, "bcoord": zeta}, "FACE": {"index": fi, "bcoord": coords}}
+
+    def _search_time_varying_z(self, z, fi, coords, ti, tau):
+        """Vertical search of each particle in the z column interpolated barycentrically from the nodes of its face,
+        using temporal linear interpolation between the snapshots ti and ti + 1.
+        """
+        found = fi >= 0
+        cols_ti = self.z.isel(
+            time=xr.DataArray(np.broadcast_to(ti, fi.shape)[found], dims="points"),
+            n_node=xr.DataArray(self.uxgrid.face_node_connectivity[fi[found], :].values, dims=("points", "nodes")),
+            ignore_grid=True,
+        ).compute()
+
+        if self.z.shape[0] == 1:
+            node_columns = cols_ti
+        else:
+            cols_tnext = self.z.isel(
+                time=xr.DataArray(np.broadcast_to(ti + 1, fi.shape)[found], dims="points"),
                 n_node=xr.DataArray(self.uxgrid.face_node_connectivity[fi[found], :].values, dims=("points", "nodes")),
                 ignore_grid=True,
             ).compute()
 
-            if self.z.shape[0] == 1:
-                node_columns = cols_ti
-            else:
-                cols_tnext = self.z.isel(
-                    time=xr.DataArray(np.broadcast_to(ti + 1, fi.shape)[found], dims="points"),
-                    n_node=xr.DataArray(
-                        self.uxgrid.face_node_connectivity[fi[found], :].values, dims=("points", "nodes")
-                    ),
-                    ignore_grid=True,
-                ).compute()
-
-                tau_points = xr.DataArray(np.broadcast_to(tau, fi.shape)[found], dims="points")
-                node_columns = (1 - tau_points) * cols_ti + tau_points * cols_tnext
-
-            bcoords = xr.DataArray(coords[found], dims=("points", "nodes"))
-            # Particles outside the mesh are given a NaN z column for vertical searching
-            columns = np.full((fi.size, self.z.sizes["zf"]), np.nan)
-            columns[found] = xr.dot(node_columns, bcoords, dim="nodes").transpose("points", "zf").values
-            zi, zeta = _search_1d_columns(columns, z)
-        else:
-            zi, zeta = _search_1d_array(self.z.values, z)
+            tau_points = xr.DataArray(np.broadcast_to(tau, fi.shape)[found], dims="points")
+            node_columns = (1 - tau_points) * cols_ti + tau_points * cols_tnext
 
-        # Particles outside the mesh are given a 0 vertical index and a NaN vertical barycentric coordinate
-        zi = np.where(found, zi, 0)
-        zeta = np.where(found, zeta, np.nan)
+        bcoords = xr.DataArray(coords[found], dims=("points", "nodes"))
+        # Particles outside the mesh are given a NaN z column for vertical searching
+        columns = np.full((fi.size, self.z.sizes["zf"]), np.nan)
+        columns[found] = xr.dot(node_columns, bcoords, dim="nodes").transpose("points", "zf").values
+        zi, zeta = _search_1d_columns(columns, z)
 
-        return {"Z": {"index": zi, "bcoord": zeta}, "FACE": {"index": fi, "bcoord": coords}}
+        # Particles outside the mesh are given a 0 vertical index and a NaN vertical barycentric coordinate
+        return np.where(found, zi, 0), np.where(found, zeta, np.nan)

From 30b4dcecd68f77fbd8a6f5369da9ab115670e8bb Mon Sep 17 00:00:00 2001
From: Wyatt Sieminski <wyatt@fluidnumerics.com>
Date: Tue, 13 Oct 2026 18:34:44 +0000
Subject: [PATCH 22/22] Removed cell_columns transpose from grid searches

---
 src/parcels/_core/uxgrid.py | 2 +-
 src/parcels/_core/xgrid.py  | 2 --
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/parcels/_core/uxgrid.py b/src/parcels/_core/uxgrid.py
index d6db846dc5..0d8fb4c321 100644
--- a/src/parcels/_core/uxgrid.py
+++ b/src/parcels/_core/uxgrid.py
@@ -188,7 +188,7 @@ def _search_time_varying_z(self, z, fi, coords, ti, tau):
         bcoords = xr.DataArray(coords[found], dims=("points", "nodes"))
         # Particles outside the mesh are given a NaN z column for vertical searching
         columns = np.full((fi.size, self.z.sizes["zf"]), np.nan)
-        columns[found] = xr.dot(node_columns, bcoords, dim="nodes").transpose("points", "zf").values
+        columns[found] = xr.dot(node_columns, bcoords, dim="nodes").values
         zi, zeta = _search_1d_columns(columns, z)
 
         # Particles outside the mesh are given a 0 vertical index and a NaN vertical barycentric coordinate
diff --git a/src/parcels/_core/xgrid.py b/src/parcels/_core/xgrid.py
index da9859821b..037ff78ed6 100644
--- a/src/parcels/_core/xgrid.py
+++ b/src/parcels/_core/xgrid.py
@@ -400,8 +400,6 @@ def _search_time_varying_depth(self, z, yi, xi, ti, tau):
             tau_found = xr.DataArray(np.broadcast_to(tau, yi.shape)[found], dims="points")
             cell_columns = (1 - tau_found) * columns_ti + tau_found * columns_tnext
 
-        cell_columns = cell_columns.transpose("points", ...)
-
         # Particles outside the grid are given a NaN z column for vertical searching
         columns = np.full((yi.size, cell_columns.shape[1]), np.nan)
         columns[found] = cell_columns.values