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..4db340e2e7 --- /dev/null +++ b/docs/user_guide/examples/tutorial_sigma_coordinates.ipynb @@ -0,0 +1,504 @@ +{ + "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 = (\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" + ] + }, + { + "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(\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", + "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(\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", + "\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 = (\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(\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", + "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(\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[\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", + " 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,\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", + "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/index.md b/docs/user_guide/index.md index 533020cc1f..7e1d38c448 100644 --- a/docs/user_guide/index.md +++ b/docs/user_guide/index.md @@ -51,6 +51,7 @@ examples/tutorial_schism.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 diff --git a/src/parcels/_core/basegrid.py b/src/parcels/_core/basegrid.py index 7b7c7e9c49..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) -> 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). @@ -49,6 +51,10 @@ 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``. + 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 ded5993edd..b629ea6243 100644 --- a/src/parcels/_core/field.py +++ b/src/parcels/_core/field.py @@ -398,8 +398,11 @@ 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"], 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/index_search.py b/src/parcels/_core/index_search.py index 5631c1fc1e..58aafb9845 100644 --- a/src/parcels/_core/index_search.py +++ b/src/parcels/_core/index_search.py @@ -64,6 +64,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 b845311e92..83e3df44c8 100644 --- a/src/parcels/_core/particleset.py +++ b/src/parcels/_core/particleset.py @@ -19,7 +19,9 @@ float_to_datelike, timedelta_to_float, ) +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"] @@ -82,6 +84,14 @@ def __init__( if z is None: minz = None for field in self.fieldset.fields.values(): + 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 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 d609b9660f..23415832b3 100644 --- a/src/parcels/_core/uxgrid.py +++ b/src/parcels/_core/uxgrid.py @@ -1,13 +1,16 @@ from __future__ import annotations +import warnings 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 +from parcels._core.warnings import FieldSetWarning if TYPE_CHECKING: import uxarray as ux @@ -32,9 +35,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 +48,20 @@ 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}") + 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 @@ -76,6 +91,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 +104,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, tau=None, tol=1e-6): """ Search for the grid cell (face) and vertical layer that contains the given points. @@ -105,6 +122,12 @@ 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 + 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. """ @@ -112,8 +135,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 +155,43 @@ 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 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) + 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 b23bd4f24d..45e354a3b5 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,28 @@ 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 + 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 {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.", + FieldSetWarning, + stacklevel=5, + ) + else: + assert_valid_depth(depth) self._ds = ds @@ -313,15 +337,20 @@ 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, 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 +358,15 @@ def search(self, z, y, x, ei=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 +374,47 @@ def search(self, z, y, x, ei=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..7beff48511 100644 --- a/src/parcels/_datasets/structured/generated.py +++ b/src/parcels/_datasets/structured/generated.py @@ -39,6 +39,32 @@ 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/src/parcels/_datasets/unstructured/generated.py b/src/parcels/_datasets/unstructured/generated.py index 35d435ec47..1b57bb7390 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_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 + 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,40 @@ 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: dict[str, tuple] = {"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_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_uxadvection.py b/tests/test_uxadvection.py index 8ffd84b1a8..e9206e75e9 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,24 @@ 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) + + +@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 + 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..ba9ecd74e1 100644 --- a/tests/test_uxarray_fieldset.py +++ b/tests/test_uxarray_fieldset.py @@ -8,8 +8,11 @@ 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, @@ -204,3 +207,39 @@ 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.filterwarnings("ignore:Time-varying \\(3D\\) z coordinates are experimental") +@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.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( + 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_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)