diff --git a/docs/user_guide/examples/tutorial_time_varying_z_coordinates.ipynb b/docs/user_guide/examples/tutorial_time_varying_z_coordinates.ipynb new file mode 100644 index 000000000..d230a4a3d --- /dev/null +++ b/docs/user_guide/examples/tutorial_time_varying_z_coordinates.ipynb @@ -0,0 +1,658 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 🖥️ Datasets with Time Varying Z Coordinates" + ] + }, + { + "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 `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, and how to modify the interpolation method used on the multi-dimensional z coordinate." + ] + }, + { + "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", + "\n", + "lon = simple_UV_dataset(\n", + " n_times=n_snapshots, nz=n_vertical_layers + 1, ny=ny, nx=nx, mesh_str=\"flat\"\n", + ")[\"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(\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", + ")\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. The z coordinate is defined at the cell centres, and by default Parcels interpolates time-varying z bilinearly in space between the surrounding cell centres, so along a row of cells the layer interfaces are straight lines between the cell centres." + ] + }, + { + "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(xs, z_moving[0][:, centre_row].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(\n", + " z_lines, z_moving[snapshot][:, centre_row], strict=True\n", + " ):\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, Y, X)`. For each particle, Parcels builds a z column by interpolating the columns at the surrounding cell centres bilinearly in space and linearly in time between the two surrounding snapshots, and then searches that column for the particle's layer." + ] + }, + { + "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)\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", + "# initial sigma coordinate of each particle, using its z column interpolated linearly in x between the cell centres\n", + "eta0 = np.interp(x0, x_centres, eta[0])\n", + "sigma0 = (z0 + eta0) / (np.interp(x0, x_centres, H) + eta0)\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)\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 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." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-output" + ] + }, + "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.AdvectionRK2],\n", + " runtime=np.timedelta64(n_days, \"D\"),\n", + " dt=np.timedelta64(30, \"m\"),\n", + " output_file=outputfile,\n", + " verbose_progress=False,\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 its z column interpolated linearly in x\n", + "# between the cell centres\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_centres, eta[snapshot])\n", + " psigma[frame] = (pz[frame] + eta_p) / (np.interp(px[frame], x_centres, H) + 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_row].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_row], strict=True\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/n_face)`, rather than a 4D `depth` coordinate with dimensions `(time, depth, Y, X)`. \n", + "\n", + "Below we build the same basin on a triangulated lattice with `sigma_coordinate_lattice_dataset`, evaluating the bathymetry and sea surface height at the lattice nodes, and compare the two z coordinates." + ] + }, + { + "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)\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": { + "tags": [ + "hide-output" + ] + }, + "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.AdvectionRK2],\n", + " runtime=np.timedelta64(n_days, \"D\"),\n", + " dt=np.timedelta64(30, \"m\"),\n", + " output_file=ux_outputfile,\n", + " verbose_progress=False,\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." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Choosing how z is interpolated\n\nWith a 4d z coordinate, in order to find a particle z value at some time, we must interpolate the current z grid at that time onto the particle position. This reduces the 4d z coordinate to a 1d z coordiante where each entry is an interface depth at the particle's exact position in space. For more details on interpolation in general, see [tutorial_interpolation](tutorial_interpolation.ipynb). There are currently a number of built in interpolators for vertical coordinates:\n\n- {py:obj}`parcels.interpolators.XVerticalCoordinateLinear` (structured grids, z on the cell centres or corners): bilinear interpolation between the surrounding columns, so z is continuous in space.\n- {py:obj}`parcels.interpolators.XVerticalCoordinateNearest` (structured grids, z on the cell centres or corners): piecewise constant interpolation based on whatever cell is closest to the particle.\n- {py:obj}`parcels.interpolators.XVerticalCoordinateConstant` (structured grids, z on the cell centres): piecewise constant interpolation based on whatever cell the particle is in. \n- {py:obj}`parcels.interpolators.UxVerticalCoordinateLinearNode` (unstructured grids, z on the nodes): barycentric interpolation from the three corners of the particle's triangle.\n- {py:obj}`parcels.interpolators.UxVerticalCoordinateConstantFace` (unstructured grids, z on the faces): piecewise constant interpolation based on whatever cell the particle is in. \n\nAll of them interpolate z linearly in time between the two surrounding snapshots.\n\nThe effect of selecting different vertical coordinate interpolators can be seen by investigating what happens to particles near the surface:" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# the same basin as above on a coarser grid with 200 km wide cells, so that the difference between the interpolators is\n", + "# easy to see\n", + "coarse_nx = 11\n", + "coarse_lon = np.linspace(lon[0], lon[-1], coarse_nx)\n", + "coarse_x_centres = coarse_lon - 0.5 * (coarse_lon[1] - coarse_lon[0])\n", + "# the first cell centre lies west of the domain, where the depth profile is held at its western value of 100 m\n", + "coarse_bottom_depth = 100 + 900 * np.log10(\n", + " 1 + 9 * np.clip(coarse_x_centres - lon[0], 0.0, None) / domain_length\n", + ")\n", + "coarse_wave = amplitude * np.sin(\n", + " 2.0\n", + " * np.pi\n", + " * (coarse_x_centres[np.newaxis, :] / wavelength - t[:, np.newaxis] / wave_period)\n", + ")\n", + "coarse_ds = simple_UV_sigma_dataset(\n", + " n_times=n_snapshots,\n", + " nz=n_vertical_layers + 1,\n", + " ny=ny,\n", + " nx=coarse_nx,\n", + " bottom_depth=np.broadcast_to(coarse_bottom_depth, (ny, coarse_nx)),\n", + " eta=np.broadcast_to(coarse_wave[:, np.newaxis, :], (n_snapshots, ny, coarse_nx)),\n", + ")\n", + "coarse_ds[\"U\"].values[:] = u0\n", + "coarse_fieldset = parcels.FieldSet.from_sgrid_conventions(coarse_ds)\n", + "\n", + "\n", + "def DeleteParticle(particles, fieldset):\n", + " any_error = particles.state >= 50 # e.g. particles that are above the sea surface\n", + " particles[any_error].state = parcels.StatusCode.Delete\n", + "\n", + "\n", + "# release particles near the surface, 15 km west of the cell edge at x = -400 km, where the sea surface of the next\n", + "# cell is lower than in the cell they start in, and advect them for one day with each interpolator\n", + "surface_z0 = np.linspace(22.0, 31.0, 4) # depths in m\n", + "z_interp_methods = {\n", + " \"XVerticalCoordinateLinear\": parcels.interpolators.XVerticalCoordinateLinear(),\n", + " \"XVerticalCoordinateConstant\": parcels.interpolators.XVerticalCoordinateConstant(),\n", + "}\n", + "for name, z_interp_method in z_interp_methods.items():\n", + " coarse_fieldset.U.grid.z_interp_method = z_interp_method\n", + " surface_pset = parcels.ParticleSet(\n", + " fieldset=coarse_fieldset,\n", + " pclass=parcels.Particle,\n", + " t=np.full(surface_z0.size, time[0]),\n", + " z=surface_z0,\n", + " y=np.zeros(surface_z0.size),\n", + " x=np.full(surface_z0.size, -415e3),\n", + " )\n", + " surface_outputfile = parcels.ParticleFile(\n", + " path=f\"{name}_particles.parquet\", outputdt=np.timedelta64(1, \"h\"), mode=\"w\"\n", + " )\n", + " surface_pset.execute(\n", + " [parcels.kernels.AdvectionRK2, DeleteParticle],\n", + " runtime=np.timedelta64(1, \"D\"),\n", + " dt=np.timedelta64(30, \"m\"),\n", + " output_file=surface_outputfile,\n", + " verbose_progress=False,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "n_surface_frames = 25 # hourly outputs from hour 0 to hour 24\n", + "surface_px = {}\n", + "for name in z_interp_methods:\n", + " surface_df = parcels.read_particlefile(f\"{name}_particles.parquet\")\n", + " frames = (surface_df[\"t\"].to_numpy() - time[0]) // np.timedelta64(1, \"h\")\n", + " surface_px[name] = np.full((n_surface_frames, surface_z0.size), np.nan)\n", + " surface_px[name][frames, surface_df[\"particle_id\"].to_numpy()] = surface_df[\"x\"]\n", + " # a particle that has left the water stays at its last position\n", + " for frame in range(1, n_surface_frames):\n", + " left_the_water = np.isnan(surface_px[name][frame])\n", + " surface_px[name][frame, left_the_water] = surface_px[name][\n", + " frame - 1, left_the_water\n", + " ]\n", + "\n", + "coarse_z = coarse_ds[\"depth\"].values[\n", + " :, :, centre_row\n", + "] # interface depths, shape (time, depth, x)\n", + "surface_colors = [f\"C{p}\" for p in range(surface_z0.size)]\n", + "\n", + "fig, axes = plt.subplots(1, 2, figsize=(12, 4), sharey=True)\n", + "panels = []\n", + "for ax, name, drawstyle in zip(\n", + " axes, z_interp_methods, [\"default\", \"steps-mid\"], strict=True\n", + "):\n", + " lines = ax.plot(\n", + " coarse_x_centres / 1e3, coarse_z[0].T, \"k\", linewidth=0.8, drawstyle=drawstyle\n", + " )\n", + " dots = ax.scatter(surface_px[name][0] / 1e3, surface_z0, c=surface_colors, zorder=3)\n", + " ax.set_title(name)\n", + " ax.set_xlim(-800, 0)\n", + " ax.set_xlabel(\"x [km]\")\n", + " panels.append((name, lines, dots))\n", + "axes[0].set_ylim(\n", + " 300, -75\n", + ") # the upper part of the water column, with depth increasing downwards\n", + "axes[0].set_ylabel(\"depth [m]\")\n", + "suptitle = fig.suptitle(f\"Particles near the surface at {time[0]}\")\n", + "plt.tight_layout()\n", + "\n", + "\n", + "def animate(frame):\n", + " for name, lines, dots in panels:\n", + " for line, interface_depths in zip(lines, coarse_z[frame], strict=True):\n", + " line.set_ydata(interface_depths)\n", + " dots.set_offsets(np.c_[surface_px[name][frame] / 1e3, surface_z0])\n", + " suptitle.set_text(f\"Particles near the surface at {time[frame]}\")\n", + "\n", + "\n", + "animate(0)\n", + "anim = FuncAnimation(fig, animate, frames=n_surface_frames, interval=300)\n", + "plt.close()\n", + "anim" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "With {py:obj}`parcels.interpolators.XVerticalCoordinateLinear`, the sea surface and the layer interfaces vary smoothly between the cell centres, while with {py:obj}`parcels.interpolators.XVerticalCoordinateConstant` they jump at every cell edge. The discontinuity of piecewise constant vertical coordinate interpolators can cause particles to spuriously leave the domain, as shown above. \n\nBy default, Parcels uses `XVerticalCoordinateLinear` for time-varying z on structured grids, {py:obj}`parcels.interpolators.UxVerticalCoordinateLinearNode` for z on the nodes of an unstructured grid, and {py:obj}`parcels.interpolators.UxVerticalCoordinateConstantFace` for z on the faces of an unstructured grid. To use a different interpolator, set the `z_interp_method` of the grid after building the {py:obj}`parcels.FieldSet`, which changes it for every field on that grid:\n\n```python\nfieldset.U.grid.z_interp_method = parcels.interpolators.XVerticalCoordinateConstant()\n```\n\nParcels raises an error when the interpolator does not fit where z is defined, for example `XVerticalCoordinateConstant` with z on the cell corners." + } + ], + "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" + }, + "mystnb": { + "execution_in_temp": true + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/docs/user_guide/index.md b/docs/user_guide/index.md index fa0c32379..64d23637c 100644 --- a/docs/user_guide/index.md +++ b/docs/user_guide/index.md @@ -56,6 +56,11 @@ group_convert_concepts.md :caption: Work with FieldSets :name: work-with-fieldsets :titlesonly: +examples/explanation_grids.ipynb +examples/tutorial_time_varying_z_coordinates.ipynb +examples/tutorial_velocityconversion.ipynb +examples/tutorial_nestedgrids.ipynb +examples/tutorial_manipulating_field_data.ipynb group_fieldsets.md ``` diff --git a/src/parcels/_core/particleset.py b/src/parcels/_core/particleset.py index 9743f6263..f600af88a 100644 --- a/src/parcels/_core/particleset.py +++ b/src/parcels/_core/particleset.py @@ -86,7 +86,7 @@ def __init__( 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 + isinstance(field.grid, XGrid) and field.grid.depth.ndim == 4 ) if has_time_varying_z: raise ValueError( diff --git a/src/parcels/_core/uxgrid.py b/src/parcels/_core/uxgrid.py index 3927a8cea..39ef536a1 100644 --- a/src/parcels/_core/uxgrid.py +++ b/src/parcels/_core/uxgrid.py @@ -1,16 +1,15 @@ 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, _search_1d_columns, uxgrid_point_in_cell from parcels._core.mesh import BaseMesh, TMesh -from parcels._core.warnings import FieldSetWarning +from parcels.interpolators._base import VerticalCoordinateInterpolator +from parcels.interpolators._uxinterpolators import UxVerticalCoordinateConstantFace, UxVerticalCoordinateLinearNode if TYPE_CHECKING: import uxarray as ux @@ -35,7 +34,7 @@ def __init__(self, grid: ux.grid.Grid, z: ux.UxDataArray, mesh: TMesh) -> None: z : ux.UxDataArray 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. + ("time", "zf", "n_node") or ("time", "zf", "n_face"), varying in time and across the mesh nodes or faces. mesh : FlatMesh or SphericalMesh The mesh of the grid, e.g. ``FlatMesh()`` or ``SphericalMesh(radius=...)``. """ @@ -50,24 +49,30 @@ def __init__(self, grid: ux.grid.Grid, z: ux.UxDataArray, mesh: TMesh) -> None: raise TypeError("z must be an instance of ux.UxDataArray") 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 = mesh self._spatialhash = None + self._z_interp_method = None + if z.ndim == 3: + self._init_time_varying_z() if is_dask_collection(self.uxgrid.face_node_connectivity.data): self.uxgrid.face_node_connectivity.load() + def _init_time_varying_z(self) -> None: + """Checks the dims of a time-varying (3D) z and assigns its default z_interp_method.""" + z = self.z + + if z.dims not in (("time", "zf", "n_node"), ("time", "zf", "n_face")): + raise ValueError( + f"A 3D z must have dims ('time', 'zf', 'n_node') or ('time', 'zf', 'n_face'), got {z.dims}" + ) + + if z.dims[-1] == "n_node": + self.z_interp_method = UxVerticalCoordinateLinearNode() + else: + self.z_interp_method = UxVerticalCoordinateConstantFace() + @property def depth(self): """ @@ -76,16 +81,31 @@ def depth(self): Included for compatibility with v3 codebase. May be removed in future. TODO v4: Evaluate """ - try: - _ = self.z.values - except KeyError: - return np.zeros(1) + # Load time varying z lazily + if self.z.ndim == 3: + return self.z return self.z.values @property def axes(self) -> list[_UXGRID_AXES]: return ["Z", "FACE"] + @property + def z_interp_method(self) -> VerticalCoordinateInterpolator | None: + return self._z_interp_method + + @z_interp_method.setter + def z_interp_method(self, value: VerticalCoordinateInterpolator): + if not isinstance(value, VerticalCoordinateInterpolator): + raise ValueError(f"z_interp_method must be a `VerticalCoordinateInterpolator` object. Got {type(value)=!r}") + if self.depth.ndim != 3: + raise ValueError("z_interp_method can only be set on a grid with a time-varying (3D) z") + if isinstance(value, UxVerticalCoordinateLinearNode) and self.z.dims[-1] != "n_node": + raise ValueError(f"UxVerticalCoordinateLinearNode needs z on the mesh nodes, got dims {self.z.dims}") + if isinstance(value, UxVerticalCoordinateConstantFace) and self.z.dims[-1] != "n_face": + raise ValueError(f"UxVerticalCoordinateConstantFace needs z on the mesh faces, got dims {self.z.dims}") + self._z_interp_method = value + def get_axis_dim(self, axis: _UXGRID_AXES) -> int: if axis not in self.axes: raise ValueError(f"Axis {axis!r} is not part of this grid. Available axes: {self.axes}") @@ -163,32 +183,16 @@ def search(self, z, y, x, ei=None, ti=None, tau=None, tol=1e-6): 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. - """ + """Vertical search of each particle in the z column built by the grid's z_interp_method.""" 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() - - tau_points = xr.DataArray(np.broadcast_to(tau, fi.shape)[found], dims="points") - node_columns = (1 - tau_points) * cols_ti + tau_points * cols_tnext + grid_positions = { + "T": {"index": np.broadcast_to(ti, fi.shape)[found], "bcoord": np.broadcast_to(tau, fi.shape)[found]}, + "FACE": {"index": fi[found], "bcoord": coords[found]}, + } - 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").values + columns[found] = self.z_interp_method.columns(grid_positions, self) 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 c96fee26b..f7d87407d 100644 --- a/src/parcels/_core/xgrid.py +++ b/src/parcels/_core/xgrid.py @@ -1,4 +1,3 @@ -import warnings from collections.abc import Hashable, Sequence from functools import cached_property from typing import Literal, cast @@ -13,9 +12,9 @@ from parcels._core.basegrid import BaseGrid from parcels._core.index_search import _search_1d_array, _search_1d_columns, _search_indices_curvilinear_2d from parcels._core.mesh import BaseMesh -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 +from parcels.interpolators._base import VerticalCoordinateInterpolator +from parcels.interpolators._xinterpolators import XVerticalCoordinateConstant, XVerticalCoordinateLinear _FIELD_DATA_ORDERING: Sequence[ptyping.XgcmAxisDirection] = "TZYX" _XGRID_AXES_ORDERING: Sequence[ptyping.XgridAxis] = "ZYX" @@ -129,6 +128,7 @@ def __init__(self, model_data: xr.Dataset, mesh: ptyping.TMesh): self._ds = model_data self.mesh = mesh self._spatialhash = None + self._z_interp_method = None ds = model_data axes = self.axes @@ -149,27 +149,31 @@ def __init__(self, model_data: xr.Dataset, mesh: ptyping.TMesh): 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, - ) + self._init_time_varying_depth() else: assert_valid_depth(depth) self._ds = ds + def _init_time_varying_depth(self) -> None: + """Checks the dims of a time-varying (4D) depth and assigns its default z_interp_method.""" + depth = self.depth + + 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 + face_registered_dims = ("time", vertical_dim, fnp_y.face, fnp_x.face) + node_registered_dims = ("time", vertical_dim, fnp_y.node, fnp_x.node) + + if depth.dims not in (face_registered_dims, node_registered_dims): + raise ValueError( + f"A 4D depth must have dims {face_registered_dims} (cell centres) or {node_registered_dims} " + f"(cell corners), got {depth.dims}" + ) + + self.z_interp_method = XVerticalCoordinateLinear() + # def __repr__(self): # return xgrid_repr(self) @@ -177,6 +181,24 @@ def __init__(self, model_data: xr.Dataset, mesh: ptyping.TMesh): def axes(self) -> list[ptyping.XgridAxis]: return _get_xgrid_axes(self.sgrid_metadata, set(self._ds.dims)) + @property + def z_interp_method(self) -> VerticalCoordinateInterpolator | None: + return self._z_interp_method + + @z_interp_method.setter + def z_interp_method(self, value: VerticalCoordinateInterpolator): + if not isinstance(value, VerticalCoordinateInterpolator): + raise ValueError(f"z_interp_method must be a `VerticalCoordinateInterpolator` object. Got {type(value)=!r}") + if self.depth.ndim != 4: + raise ValueError("z_interp_method can only be set on a grid with a time-varying (4D) depth") + depth_dims = self.depth.dims + if ( + isinstance(value, XVerticalCoordinateConstant) + and get_dim_position(self.sgrid_metadata, depth_dims[-1]) != "face" + ): + raise ValueError(f"XVerticalCoordinateConstant needs depth on the cell centres, got dims {depth_dims}") + self._z_interp_method = value + @property def lon(self): """ @@ -219,6 +241,9 @@ def depth(self): """ if "Z" not in self.axes: return np.zeros(1) + # Load time varying z lazily + if self._ds["depth"].ndim == 4: + return self._ds["depth"] return self._ds["depth"].values @property @@ -379,7 +404,7 @@ def search(self, z, y, x, ei=None, ti=None, tau=None): raise NotImplementedError("Searching in >2D lon/lat arrays is not implemented yet.") if "Z" in self.axes and ds.depth.ndim == 4: - zi, zeta = self._search_time_varying_depth(z, yi, xi, ti, tau) + zi, zeta = self._search_time_varying_depth(z, yi, eta, xi, xsi, ti, tau) elif "Z" in self.axes: zi, zeta = _search_1d_array(ds.depth.values, z) else: @@ -391,35 +416,18 @@ 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. - """ + def _search_time_varying_depth(self, z, yi, eta, xi, xsi, ti, tau): + """Vertical search of each particle in the z column built by the grid's z_interp_method.""" 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"), + grid_positions = { + "T": {"index": np.broadcast_to(ti, yi.shape)[found], "bcoord": np.broadcast_to(tau, yi.shape)[found]}, + "Y": {"index": yi[found], "bcoord": eta[found]}, + "X": {"index": xi[found], "bcoord": xsi[found]}, } - 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 - # 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 + columns = np.full((yi.size, self.depth.shape[1]), np.nan) + columns[found] = self.z_interp_method.columns(grid_positions, self) zi, zeta = _search_1d_columns(columns, z) # Particles outside the grid are given a 0 vertical index and a NaN vertical barycentric coordinate diff --git a/src/parcels/interpolators/__init__.py b/src/parcels/interpolators/__init__.py index 2f131c546..752a51965 100644 --- a/src/parcels/interpolators/__init__.py +++ b/src/parcels/interpolators/__init__.py @@ -6,6 +6,8 @@ UxConstantFaceLinearZF, UxLinearNodeConstantZC, UxLinearNodeLinearZF, + UxVerticalCoordinateConstantFace, + UxVerticalCoordinateLinearNode, ) from ._xinterpolators import ( CGrid_Tracer, @@ -17,6 +19,9 @@ XLinearInvdistLandTracer, XNearest, XPartialslip, + XVerticalCoordinateConstant, + XVerticalCoordinateLinear, + XVerticalCoordinateNearest, ) __all__ = [ # noqa: RUF022 @@ -30,6 +35,9 @@ "XLinear_Velocity", "XNearest", "XPartialslip", + "XVerticalCoordinateConstant", + "XVerticalCoordinateLinear", + "XVerticalCoordinateNearest", # uxinterpolators "Ux_Velocity", "UxConstantFaceConstantZC", @@ -38,4 +46,6 @@ "UxConstantEdgeLinearZF", "UxLinearNodeConstantZC", "UxLinearNodeLinearZF", + "UxVerticalCoordinateConstantFace", + "UxVerticalCoordinateLinearNode", ] diff --git a/src/parcels/interpolators/_base.py b/src/parcels/interpolators/_base.py index 6a550b240..afb54e267 100644 --- a/src/parcels/interpolators/_base.py +++ b/src/parcels/interpolators/_base.py @@ -1,6 +1,9 @@ from abc import ABC, abstractmethod +from types import SimpleNamespace from typing import Any +import numpy as np + class ScalarInterpolator(ABC): @abstractmethod @@ -18,3 +21,36 @@ def interp(self, particle_positions, grid_positions, vectorfield) -> Any: #! AP def __repr__(self): return f"{self.__class__.__name__}(...)" + + +class VerticalCoordinateInterpolator(ABC): + """Used for models with time varying z.""" + + @abstractmethod + def columns(self, grid_positions, grid) -> np.ndarray: #! API a WIP, returns an (npart, n_interfaces) array + ... + + def _columns_from_scalar(self, interp: ScalarInterpolator, grid_positions, grid) -> np.ndarray: + """Evaluates a scalar interpolator on the grid's time-varying depth at every vertical level, giving each + particle's z column as an (npart, nlev) array. This is effectively only a horizontal and temporal interpolation, any + input vertical position or barycentric coordinate is ignored. + """ + nlev = grid.depth.shape[1] + npart = next(np.shape(pos["index"])[0] for axis, pos in grid_positions.items() if axis != "T") + + def repeat(a): + a = np.broadcast_to(a, (npart, *np.shape(a)[1:])) + return np.repeat(a, nlev, axis=0) + + level_positions = { + axis: {"index": repeat(pos["index"]), "bcoord": repeat(pos["bcoord"])} + for axis, pos in grid_positions.items() + } + level_positions["Z"] = {"index": np.tile(np.arange(nlev), npart), "bcoord": np.zeros(npart * nlev)} + + # Scalar interpolators only use ``field.data`` and ``field.grid``, so a lightweight stand-in for the depth works + depth_field = SimpleNamespace(data=grid.depth, grid=grid) + return np.asarray(interp.interp({}, level_positions, depth_field)).reshape(npart, nlev) + + def __repr__(self): + return f"{self.__class__.__name__}(...)" diff --git a/src/parcels/interpolators/_uxinterpolators.py b/src/parcels/interpolators/_uxinterpolators.py index 6db1000eb..752d3f23b 100644 --- a/src/parcels/interpolators/_uxinterpolators.py +++ b/src/parcels/interpolators/_uxinterpolators.py @@ -8,11 +8,11 @@ import xarray as xr from dask import is_dask_collection -from parcels.interpolators._base import ScalarInterpolator, VectorInterpolator +from parcels.interpolators._base import ScalarInterpolator, VectorInterpolator, VerticalCoordinateInterpolator if TYPE_CHECKING: from parcels._core.field import Field, VectorField - from parcels._core.uxgrid import _UXGRID_AXES + from parcels._core.uxgrid import _UXGRID_AXES, UxGrid class UxConstantFaceConstantZC(ScalarInterpolator): @@ -25,17 +25,29 @@ def interp( field: Field, ): # Broadcast the per-axis indices to a common (npart,) shape (``ti`` may be scalar for time-constant fields) - ti, zi, fi = np.broadcast_arrays( - grid_positions["T"]["index"], grid_positions["Z"]["index"], grid_positions["FACE"]["index"] + ti, tau, zi, fi = np.broadcast_arrays( + grid_positions["T"]["index"], + grid_positions["T"]["bcoord"], + grid_positions["Z"]["index"], + grid_positions["FACE"]["index"], ) tdim, zdim, fdim = field.data.dims + + if np.any(tau > 0): + t_levels = np.stack([ti, np.clip(ti + 1, 0, field.data.sizes[tdim] - 1)]) + t_weights = np.stack([1 - tau, tau]) + else: + t_levels, t_weights = ti[np.newaxis], np.ones((1, ti.size)) selection_dict = { - tdim: xr.DataArray(ti, dims="points"), + tdim: xr.DataArray(t_levels, dims=("tlevels", "points")), zdim: xr.DataArray(zi, dims="points"), fdim: xr.DataArray(fi, dims="points"), } - value = field.data.isel(selection_dict, ignore_grid=True).data + face_data = field.data.isel(selection_dict, ignore_grid=True) + + # Linear interpolation in time between the snapshots ti and ti + 1 + value = (face_data * xr.DataArray(t_weights, dims=("tlevels", "points"))).sum("tlevels").data return value.compute() if is_dask_collection(value) else value @@ -51,28 +63,41 @@ def interp( grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]], field: Field, ): - ti, zi, fi = np.broadcast_arrays( - grid_positions["T"]["index"], grid_positions["Z"]["index"], grid_positions["FACE"]["index"] + ti, tau, zi, fi = np.broadcast_arrays( + grid_positions["T"]["index"], + grid_positions["T"]["bcoord"], + grid_positions["Z"]["index"], + grid_positions["FACE"]["index"], ) zeta = grid_positions["Z"]["bcoord"] tdim, zdim, fdim = field.data.dims - def _zsample(z_index): - """Pointwise ``isel`` of the face values at a single vertical interface level.""" + if np.any(tau > 0): + t_levels = np.stack([ti, np.clip(ti + 1, 0, field.data.sizes[tdim] - 1)]) + t_weights = np.stack([1 - tau, tau]) + else: + t_levels, t_weights = ti[np.newaxis], np.ones((1, ti.size)) + time_weights = xr.DataArray(t_weights, dims=("tlevels", "points")) + + def _z_and_t_sample(z_index): + """Pointwise ``isel`` of the face values at a single vertical interface level, linear in time between the + snapshots ti and ti + 1. + """ selection_dict = { - tdim: xr.DataArray(ti, dims="points"), + tdim: xr.DataArray(t_levels, dims=("tlevels", "points")), zdim: xr.DataArray(z_index, dims="points"), fdim: xr.DataArray(fi, dims="points"), } - value = field.data.isel(selection_dict, ignore_grid=True).data + face_data = field.data.isel(selection_dict, ignore_grid=True) + value = (face_data * time_weights).sum("tlevels").data return value.compute() if is_dask_collection(value) else value # The zi refers to the vertical layer index. The field in this routine are assumed to be defined at the vertical interface levels. # For interface zi, the interface indices are [zi, zi+1], so we need to use the values at zi and zi+1. # First, do barycentric interpolation in the lateral direction for each interface level - fzk = _zsample(zi) - fzkp1 = _zsample(zi + 1) + fzk = _z_and_t_sample(zi) + fzkp1 = _z_and_t_sample(zi + 1) # Then, do piecewise linear interpolation in the vertical direction return fzk * (1 - zeta) + fzkp1 * zeta @@ -92,21 +117,32 @@ def interp( Effectively, it applies barycentric interpolation in the lateral direction and piecewise constant interpolation in the vertical direction. """ - ti, zi, fi = np.broadcast_arrays( - grid_positions["T"]["index"], grid_positions["Z"]["index"], grid_positions["FACE"]["index"] + ti, tau, zi, fi = np.broadcast_arrays( + grid_positions["T"]["index"], + grid_positions["T"]["bcoord"], + grid_positions["Z"]["index"], + grid_positions["FACE"]["index"], ) bcoords = xr.DataArray(grid_positions["FACE"]["bcoord"], dims=("points", "nodes")) node_ids = field.grid.uxgrid.face_node_connectivity[fi, :].values tdim, zdim, ndim = field.data.dims + + if np.any(tau > 0): + t_levels = np.stack([ti, np.clip(ti + 1, 0, field.data.sizes[tdim] - 1)]) + t_weights = np.stack([1 - tau, tau]) + else: + t_levels, t_weights = ti[np.newaxis], np.ones((1, ti.size)) selection_dict = { - tdim: xr.DataArray(ti, dims="points"), + tdim: xr.DataArray(t_levels, dims=("tlevels", "points")), zdim: xr.DataArray(zi, dims="points"), ndim: xr.DataArray(node_ids, dims=("points", "nodes")), } node_data = field.data.isel(selection_dict, ignore_grid=True) - value = (node_data * bcoords).sum("nodes").data # Barycentric interpolation in the lateral direction + # Barycentric interpolation in the lateral direction and linear interpolation in time + weights = bcoords * xr.DataArray(t_weights, dims=("tlevels", "points")) + value = (node_data * weights).sum(("nodes", "tlevels")).data return value.compute() if is_dask_collection(value) else value @@ -124,8 +160,11 @@ def interp( Effectively, it applies barycentric interpolation in the lateral direction and piecewise linear interpolation in the vertical direction. """ - ti, zi, fi = np.broadcast_arrays( - grid_positions["T"]["index"], grid_positions["Z"]["index"], grid_positions["FACE"]["index"] + ti, tau, zi, fi = np.broadcast_arrays( + grid_positions["T"]["index"], + grid_positions["T"]["bcoord"], + grid_positions["Z"]["index"], + grid_positions["FACE"]["index"], ) zeta = grid_positions["Z"]["bcoord"] bcoords = xr.DataArray(grid_positions["FACE"]["bcoord"], dims=("points", "nodes")) @@ -133,22 +172,32 @@ def interp( tdim, zdim, ndim = field.data.dims - def _zsample(z_index): - """Barycentric (lateral) interpolation of the node values at a single vertical interface level.""" + if np.any(tau > 0): + t_levels = np.stack([ti, np.clip(ti + 1, 0, field.data.sizes[tdim] - 1)]) + t_weights = np.stack([1 - tau, tau]) + else: + t_levels, t_weights = ti[np.newaxis], np.ones((1, ti.size)) + # Barycentric weights in the lateral direction combined with linear weights in time + weights = bcoords * xr.DataArray(t_weights, dims=("tlevels", "points")) + + def _z_and_t_sample(z_index): + """Barycentric (lateral) interpolation of the node values at a single vertical interface level, linear in + time between the snapshots ti and ti + 1. + """ selection_dict = { - tdim: xr.DataArray(ti, dims="points"), + tdim: xr.DataArray(t_levels, dims=("tlevels", "points")), zdim: xr.DataArray(z_index, dims="points"), ndim: xr.DataArray(node_ids, dims=("points", "nodes")), } - # Reduce over the "nodes" dimension by name so the result is independent of ``isel`` dim order. + # Reduce over the "nodes" and "tlevels" dimensions by name so the result is independent of ``isel`` dim order. node_data = field.data.isel(selection_dict, ignore_grid=True) - return (node_data * bcoords).sum("nodes").data + return (node_data * weights).sum(("nodes", "tlevels")).data # The zi refers to the vertical layer index. The field in this routine are assumed to be defined at the vertical interface levels. # For interface zi, the interface indices are [zi, zi+1], so we need to use the values at zi and zi+1. # First, do barycentric interpolation in the lateral direction for each interface level - fzk = _zsample(zi) - fzkp1 = _zsample(zi + 1) + fzk = _z_and_t_sample(zi) + fzkp1 = _z_and_t_sample(zi + 1) # Then, do piecewise linear interpolation in the vertical direction value = fzk * (1 - zeta) + fzkp1 * zeta @@ -168,19 +217,31 @@ def interp( field: Field, ): # Broadcast the per-axis indices to a common (npart,) shape (``ti`` may be scalar for time-constant fields) - ti, zi, fi = np.broadcast_arrays( - grid_positions["T"]["index"], grid_positions["Z"]["index"], grid_positions["FACE"]["index"] + ti, tau, zi, fi = np.broadcast_arrays( + grid_positions["T"]["index"], + grid_positions["T"]["bcoord"], + grid_positions["Z"]["index"], + grid_positions["FACE"]["index"], ) edge_ids = field.grid.uxgrid.face_edge_connectivity[fi, :].values tdim, zdim, edim = field.data.dims + + if np.any(tau > 0): + t_levels = np.stack([ti, np.clip(ti + 1, 0, field.data.sizes[tdim] - 1)]) + t_weights = np.stack([1 - tau, tau]) + else: + t_levels, t_weights = ti[np.newaxis], np.ones((1, ti.size)) selection_dict = { - tdim: xr.DataArray(ti, dims="points"), + tdim: xr.DataArray(t_levels, dims=("tlevels", "points")), zdim: xr.DataArray(zi, dims="points"), edim: xr.DataArray(edge_ids, dims=("points", "edges")), } edge_data = field.data.isel(selection_dict, ignore_grid=True) - value = edge_data.mean("edges").data # Face value is the mean of its surrounding edges + face_data = edge_data.mean("edges") # Face value is the mean of its surrounding edges + + # Linear interpolation in time between the snapshots ti and ti + 1 + value = (face_data * xr.DataArray(t_weights, dims=("tlevels", "points"))).sum("tlevels").data return value.compute() if is_dask_collection(value) else value @@ -198,27 +259,39 @@ def interp( field: Field, ): # Broadcast the per-axis indices to a common (npart,) shape (``ti`` may be scalar for time-constant fields) - ti, zi, fi = np.broadcast_arrays( - grid_positions["T"]["index"], grid_positions["Z"]["index"], grid_positions["FACE"]["index"] + ti, tau, zi, fi = np.broadcast_arrays( + grid_positions["T"]["index"], + grid_positions["T"]["bcoord"], + grid_positions["Z"]["index"], + grid_positions["FACE"]["index"], ) zeta = grid_positions["Z"]["bcoord"] edge_ids = field.grid.uxgrid.face_edge_connectivity[fi, :].values tdim, zdim, edim = field.data.dims - def _zsample(z_index): - """Mean of the edge values surrounding each face at a single vertical interface level.""" + if np.any(tau > 0): + t_levels = np.stack([ti, np.clip(ti + 1, 0, field.data.sizes[tdim] - 1)]) + t_weights = np.stack([1 - tau, tau]) + else: + t_levels, t_weights = ti[np.newaxis], np.ones((1, ti.size)) + time_weights = xr.DataArray(t_weights, dims=("tlevels", "points")) + + def _z_and_t_sample(z_index): + """Mean of the edge values surrounding each face at a single vertical interface level, linear in time + between the snapshots ti and ti + 1. + """ selection_dict = { - tdim: xr.DataArray(ti, dims="points"), + tdim: xr.DataArray(t_levels, dims=("tlevels", "points")), zdim: xr.DataArray(z_index, dims="points"), edim: xr.DataArray(edge_ids, dims=("points", "edges")), } edge_data = field.data.isel(selection_dict, ignore_grid=True) - return edge_data.mean("edges").data + return (edge_data.mean("edges") * time_weights).sum("tlevels").data # zi is the vertical layer index, bounded by the interface levels zi and zi + 1 - fzk = _zsample(zi) - fzkp1 = _zsample(zi + 1) + fzk = _z_and_t_sample(zi) + fzkp1 = _z_and_t_sample(zi + 1) value = fzk * (1 - zeta) + fzkp1 * zeta return value.compute() if is_dask_collection(value) else value @@ -243,3 +316,21 @@ def interp( else: w = np.zeros_like(u) return u, v, w + + +class UxVerticalCoordinateConstantFace(VerticalCoordinateInterpolator): + """Piecewise constant over the face (face registered z), linear in time.""" + + def columns( + self, grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]], grid: UxGrid + ) -> np.ndarray: + return self._columns_from_scalar(UxConstantFaceConstantZC(), grid_positions, grid) + + +class UxVerticalCoordinateLinearNode(VerticalCoordinateInterpolator): + """Barycentric from the face's nodes (node registered z), linear in time.""" + + def columns( + self, grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]], grid: UxGrid + ) -> np.ndarray: + return self._columns_from_scalar(UxLinearNodeConstantZC(), grid_positions, grid) diff --git a/src/parcels/interpolators/_xinterpolators.py b/src/parcels/interpolators/_xinterpolators.py index ae32932cc..b8488fc12 100644 --- a/src/parcels/interpolators/_xinterpolators.py +++ b/src/parcels/interpolators/_xinterpolators.py @@ -11,7 +11,8 @@ import parcels._core.utils.interpolation as i_u import parcels._sgrid as sgrid import parcels._typing as ptyping -from parcels.interpolators._base import ScalarInterpolator, VectorInterpolator +from parcels._sgrid.accessor import get_dim_position +from parcels.interpolators._base import ScalarInterpolator, VectorInterpolator, VerticalCoordinateInterpolator if TYPE_CHECKING: from parcels._core._windowed_array import WindowedArray @@ -635,3 +636,81 @@ def interp( values[exact_particles] = exact_vals[exact_particles] return values.compute() if is_dask_collection(values) else values + + +class XVerticalCoordinateConstant(VerticalCoordinateInterpolator): + """Piecewise constant over the grid cell (face registered z), linear in time. + + Each particle uses the z column at the centre of its cell, so z is discontinuous between cells. + """ + + def columns( + self, grid_positions: dict[ptyping.XgridAxis, dict[str, int | float | np.ndarray]], grid: XGrid + ) -> np.ndarray: + # Search indices are cells (indexed by their lower node). Point each particle exactly (bcoord 0) at its cell's + # face, so XNearest picks that single column. + offsets = _get_offsets_dictionary(grid) + yi, xi = grid_positions["Y"]["index"], grid_positions["X"]["index"] + cell_centre_positions = { + "T": grid_positions["T"], + "Y": {"index": yi + offsets["Y"], "bcoord": np.zeros(np.shape(yi))}, + "X": {"index": xi + offsets["X"], "bcoord": np.zeros(np.shape(xi))}, + } + return self._columns_from_scalar(XNearest(), cell_centre_positions, grid) + + +def _position_between_cell_centres(face_index: np.ndarray, bcoord: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Moves an index/barycentric coordinate combination back by half a cell in index space. + + The grid search gives positions relative to the nodes, but face registered data lives at the cell centres, half a + cell further along. Shifting back by half a cell gives the position relative to the cell centres, so node based + interpolators (e.g. XLinear) can interpolate face registered data. Positions before the first cell centre are + clamped onto it. + """ + upper_half = bcoord >= 0.5 + lower_face = np.where(upper_half, face_index, face_index - 1) + centre_bcoord = np.where(upper_half, bcoord - 0.5, bcoord + 0.5) + below_first_centre = lower_face < 0 + return np.where(below_first_centre, 0, lower_face), np.where(below_first_centre, 0.0, centre_bcoord) + + +class XVerticalCoordinateLinear(VerticalCoordinateInterpolator): + """Bilinear between the surrounding z columns (face or node registered z), linear in time. + + Node registered z is interpolated from the corners of the particle's cell. Face registered z is interpolated from + the four cell centres around the particle. Both are continuous in space. Beyond the outermost cell centres z is + constant. + """ + + def columns( + self, grid_positions: dict[ptyping.XgridAxis, dict[str, int | float | np.ndarray]], grid: XGrid + ) -> np.ndarray: + if get_dim_position(grid.sgrid_metadata, grid.depth.dims[-1]) == "face": + offsets = _get_offsets_dictionary(grid) + yi, eta = _position_between_cell_centres( + grid_positions["Y"]["index"] + offsets["Y"], grid_positions["Y"]["bcoord"] + ) + xi, xsi = _position_between_cell_centres( + grid_positions["X"]["index"] + offsets["X"], grid_positions["X"]["bcoord"] + ) + grid_positions = { + "T": grid_positions["T"], + "Y": {"index": yi, "bcoord": eta}, + "X": {"index": xi, "bcoord": xsi}, + } + return self._columns_from_scalar(XLinear(), grid_positions, grid) + + +class XVerticalCoordinateNearest(VerticalCoordinateInterpolator): + """Nearest z column (face or node registered z), linear in time. + + For node registered z this is the nearest corner of the particle's cell. For face registered z the nearest cell + centre is the particle's own cell, so this is the same as XVerticalCoordinateConstant. + """ + + def columns( + self, grid_positions: dict[ptyping.XgridAxis, dict[str, int | float | np.ndarray]], grid: XGrid + ) -> np.ndarray: + if get_dim_position(grid.sgrid_metadata, grid.depth.dims[-1]) == "face": + return XVerticalCoordinateConstant().columns(grid_positions, grid) + return self._columns_from_scalar(XNearest(), grid_positions, grid) diff --git a/tests/test_advection.py b/tests/test_advection.py index e12b05b4b..84c843a09 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -577,7 +577,6 @@ def test_mitgcm(): 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 diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py index 7fd63f5c9..c1b042a43 100644 --- a/tests/test_fieldset.py +++ b/tests/test_fieldset.py @@ -704,7 +704,6 @@ def test_fieldset_describe_backends(tmp_path): # 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", [ diff --git a/tests/test_uxadvection.py b/tests/test_uxadvection.py index 9d7babd18..e0eeccb2d 100644 --- a/tests/test_uxadvection.py +++ b/tests/test_uxadvection.py @@ -30,7 +30,6 @@ 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 b8a748f37..6c0e69b2a 100644 --- a/tests/test_uxarray_fieldset.py +++ b/tests/test_uxarray_fieldset.py @@ -206,6 +206,41 @@ def test_from_ugrid_conventions_selects_correct_interpolator(): assert isinstance(getattr(fieldset, f"{lateral_dim}_{vertical_dim}").interp_method, expected_interpolator) +def test_ux_interpolators_are_linear_in_time(): + """The default interpolator of every registration is linear in time between snapshots (fields uniform in space).""" + ds = convert.icon_to_ugrid(datasets_unstructured["icon_square_delaunay_uniform_z_coordinate"].copy(deep=True)) + rng = np.random.default_rng(0) + value_at_snapshots = rng.random(ds.sizes["time"]) + registrations = [ + ("zc", "n_face"), + ("zf", "n_face"), + ("zc", "n_node"), + ("zf", "n_node"), + ("zc", "n_edge"), + ("zf", "n_edge"), + ] + for vertical_dim, lateral_dim in registrations: + shape = (ds.sizes["time"], ds.sizes[vertical_dim], getattr(ds.uxgrid, lateral_dim)) + ds[f"{lateral_dim}_{vertical_dim}"] = ux.UxDataArray( + np.broadcast_to(value_at_snapshots[:, np.newaxis, np.newaxis], shape), + dims=["time", vertical_dim, lateral_dim], + coords={"time": ds.time, vertical_dim: ds[vertical_dim]}, + uxgrid=ds.uxgrid, + ) + fieldset = FieldSet.from_ugrid_conventions(ds) + + # One particle at each face centre, at a random time (in seconds since the first snapshot) + xc, yc = ds.uxgrid.face_lon.values, ds.uxgrid.face_lat.values + zq = np.full(xc.shape, 300.0) + snapshot_times = (ds.time.values - ds.time.values[0]) / np.timedelta64(1, "s") + tq = rng.uniform(snapshot_times[0], snapshot_times[-1], xc.size) + + expected = np.interp(tq, snapshot_times, value_at_snapshots) + for vertical_dim, lateral_dim in registrations: + values = getattr(fieldset, f"{lateral_dim}_{vertical_dim}").eval(t=tq, z=zq, y=yc, x=xc) + np.testing.assert_allclose(values, expected, err_msg=f"{lateral_dim}_{vertical_dim}") + + _NESTEDGRIDS_NODES = np.array( [ [10.0, 15.0], @@ -264,7 +299,6 @@ 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", [ @@ -284,7 +318,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") 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_vertical_coordinate_interpolators.py b/tests/test_vertical_coordinate_interpolators.py new file mode 100644 index 000000000..01b67bdba --- /dev/null +++ b/tests/test_vertical_coordinate_interpolators.py @@ -0,0 +1,216 @@ +from types import SimpleNamespace + +import numpy as np +import pytest +import uxarray as ux + +from parcels import FieldSet, UxGrid +from parcels._core.mesh import FlatMesh +from parcels._datasets.structured.generated import simple_UV_sigma_dataset +from parcels._datasets.unstructured.generated import sigma_coordinate_lattice_dataset +from parcels.interpolators import ( + UxConstantFaceConstantZC, + UxVerticalCoordinateConstantFace, + UxVerticalCoordinateLinearNode, + XLinear, + XVerticalCoordinateConstant, + XVerticalCoordinateLinear, + XVerticalCoordinateNearest, +) +from parcels.interpolators._base import ScalarInterpolator +from parcels.interpolators._xinterpolators import _get_offsets_dictionary, _position_between_cell_centres + +# Time-varying depth used to test the vertical coordinate interpolators. It is linear in level, lateral index and time, +# so the expected z column of each interpolator can be written down exactly, and strictly increasing along the levels. +SIGMA_SHAPE = {"n_times": 3, "nz": 5, "ny": 8, "nx": 8} +DEPTH_PER_LEVEL, DEPTH_PER_Y, DEPTH_PER_X, DEPTH_PER_TIME = 10.0, 0.3, 0.2, 0.1 + + +def _linear_depth(level, y, x, t): + return DEPTH_PER_LEVEL * level + DEPTH_PER_Y * y + DEPTH_PER_X * x + DEPTH_PER_TIME * t + + +def _sigma_grid(registration): + """XGrid with a time-varying depth on the cell centres ("face") or on the cell corners ("node").""" + ds = simple_UV_sigma_dataset(**SIGMA_SHAPE) + lateral_dims = ("YC", "XC") if registration == "face" else ("YG", "XG") + t, level, y, x = np.meshgrid(*(np.arange(n) for n in SIGMA_SHAPE.values()), indexing="ij") + ds = ds.assign_coords(depth=(("time", "ZF", *lateral_dims), _linear_depth(level, y, x, t), {"axis": "Z"})) + return FieldSet.from_sgrid_conventions(ds).U.grid + + +def _random_grid_positions(npart, seed=0): + """Positions in the interior of the grid, away from the edges where the face registered columns are clamped.""" + rng = np.random.default_rng(seed) + return { + "T": {"index": rng.integers(0, SIGMA_SHAPE["n_times"] - 1, npart), "bcoord": rng.random(npart)}, + "Y": {"index": rng.integers(1, SIGMA_SHAPE["ny"] - 3, npart), "bcoord": rng.random(npart)}, + "X": {"index": rng.integers(1, SIGMA_SHAPE["nx"] - 3, npart), "bcoord": rng.random(npart)}, + } + + +# Expected lateral position of a particle's z column, from its search index (i), barycentric coordinate (b) and the +# padding offset between cell and face indices (o), measured in the indices of the registration's depth +def _cell_centre(i, b, o): + return i + o + + +def _nearest_corner(i, b, o): + return i + (b > 0.5) + + +def _position_in_face_indices(i, b, o): + return i + o + b - 0.5 + + +def _position_in_node_indices(i, b, o): + return i + b + + +@pytest.mark.parametrize( + ("z_interp_method", "registration", "expected_y", "expected_x"), + [ + # Columns at the particle's cell centre (face indices are the cell indices plus the padding offset) + pytest.param(XVerticalCoordinateConstant(), "face", _cell_centre, _cell_centre, id="constant-face"), + pytest.param(XVerticalCoordinateNearest(), "face", _cell_centre, _cell_centre, id="nearest-face"), + # Column at the nearest corner of the particle's cell + pytest.param(XVerticalCoordinateNearest(), "node", _nearest_corner, _nearest_corner, id="nearest-node"), + # Bilinear interpolation reproduces the linear depth at the particle's position, measured in face or node indices + pytest.param( + XVerticalCoordinateLinear(), "face", _position_in_face_indices, _position_in_face_indices, id="linear-face" + ), + pytest.param( + XVerticalCoordinateLinear(), "node", _position_in_node_indices, _position_in_node_indices, id="linear-node" + ), + ], +) +@pytest.mark.parametrize("time_varying", [True, False]) +def test_vertical_coordinate_columns(z_interp_method, registration, expected_y, expected_x, time_varying): + grid = _sigma_grid(registration) + grid.z_interp_method = z_interp_method + offsets = _get_offsets_dictionary(grid) + grid_positions = _random_grid_positions(npart=200) + if not time_varying: + grid_positions["T"] = {"index": 1, "bcoord": 0.0} # scalar time position, as on a time-constant grid + + columns = grid.z_interp_method.columns(grid_positions, grid) + + expected = _linear_depth( + np.arange(SIGMA_SHAPE["nz"])[np.newaxis, :], + expected_y(grid_positions["Y"]["index"], grid_positions["Y"]["bcoord"], offsets["Y"])[:, np.newaxis], + expected_x(grid_positions["X"]["index"], grid_positions["X"]["bcoord"], offsets["X"])[:, np.newaxis], + np.add(grid_positions["T"]["index"], grid_positions["T"]["bcoord"])[..., np.newaxis] if time_varying else 1.0, + ) + assert columns.shape == (200, SIGMA_SHAPE["nz"]) + np.testing.assert_allclose(columns, expected) + + +@pytest.mark.parametrize("registration", ["face", "node"]) +def test_vertical_coordinate_linear_is_continuous(registration): + """z is continuous across cell edges, and across the cell centre (where face registered z changes the bracketing + cell centres it interpolates between).""" + ds = simple_UV_sigma_dataset(**SIGMA_SHAPE) + lateral_dims = ("YC", "XC") if registration == "face" else ("YG", "XG") + rng = np.random.default_rng(1) + depth = ds["depth"].values + rng.random(ds["depth"].shape) * 0.01 # non-linear depth, unlike _sigma_grid + ds = ds.assign_coords(depth=(("time", "ZF", *lateral_dims), depth, {"axis": "Z"})) + grid = FieldSet.from_sgrid_conventions(ds).U.grid + grid_positions = _random_grid_positions(npart=100) + yi, npart = grid_positions["Y"]["index"], len(grid_positions["Y"]["index"]) + + def columns_at(y_index, y_bcoord): + return grid.z_interp_method.columns( + {**grid_positions, "Y": {"index": y_index, "bcoord": np.full(npart, y_bcoord)}}, grid + ) + + np.testing.assert_allclose(columns_at(yi, 1.0), columns_at(yi + 1, 0.0)) + np.testing.assert_allclose(columns_at(yi, 0.5 - 1e-12), columns_at(yi, 0.5)) + + +@pytest.mark.parametrize( + ("face_index", "bcoord", "expected_index", "expected_bcoord"), + [ + pytest.param(3, 0.75, 3, 0.25, id="upper-half-of-cell"), + pytest.param(3, 0.5, 3, 0.0, id="at-cell-centre"), + pytest.param(3, 0.25, 2, 0.75, id="lower-half-of-cell"), + pytest.param(0, 0.25, 0, 0.0, id="clamped-before-first-cell-centre"), + ], +) +def test_position_between_cell_centres(face_index, bcoord, expected_index, expected_bcoord): + index, centre_bcoord = _position_between_cell_centres(np.array([face_index]), np.array([bcoord])) + assert index[0] == expected_index + np.testing.assert_allclose(centre_bcoord[0], expected_bcoord) + + +def test_z_interp_method_defaults_and_validation(): + face_grid, node_grid = _sigma_grid("face"), _sigma_grid("node") + assert isinstance(face_grid.z_interp_method, XVerticalCoordinateLinear) + assert isinstance(node_grid.z_interp_method, XVerticalCoordinateLinear) + + with pytest.raises(ValueError, match="needs depth on the cell centres"): + node_grid.z_interp_method = XVerticalCoordinateConstant() + with pytest.raises(ValueError, match="must be a `VerticalCoordinateInterpolator`"): + face_grid.z_interp_method = XLinear() + + +class _EchoInterpolator(ScalarInterpolator): + """Records the grid positions it is evaluated at, and returns 100 * time index + 10 * face index + level index.""" + + def interp(self, particle_positions, grid_positions, field): + self.grid_positions = grid_positions + return 100 * grid_positions["T"]["index"] + 10 * grid_positions["FACE"]["index"] + grid_positions["Z"]["index"] + + +@pytest.mark.parametrize("time_varying", [True, False]) +def test_columns_from_scalar(time_varying): + npart, nlev = 4, 3 + grid = SimpleNamespace(depth=np.zeros((2, nlev, 5))) # only the number of levels of depth is used + fi, face_bcoords = np.array([4, 0, 2, 3]), np.arange(npart * 3).reshape(npart, 3) / 10 + ti, tau = (np.array([0, 1, 1, 0]), np.array([0.1, 0.2, 0.3, 0.4])) if time_varying else (1, 0.0) + grid_positions = {"T": {"index": ti, "bcoord": tau}, "FACE": {"index": fi, "bcoord": face_bcoords}} + echo = _EchoInterpolator() + + columns = UxVerticalCoordinateLinearNode()._columns_from_scalar(echo, grid_positions, grid) + + np.testing.assert_array_equal(columns, 100 * np.reshape(ti, (-1, 1)) + 10 * fi[:, np.newaxis] + np.arange(nlev)) + + # Each level is evaluated exactly on the level, at its particle's face and time barycentric coordinates + recorded = echo.grid_positions + np.testing.assert_array_equal(recorded["Z"]["bcoord"], 0.0) + np.testing.assert_array_equal( + recorded["FACE"]["bcoord"].reshape(npart, nlev, 3), + np.broadcast_to(face_bcoords[:, np.newaxis], (npart, nlev, 3)), + ) + np.testing.assert_array_equal( + recorded["T"]["bcoord"].reshape(npart, nlev), + np.broadcast_to(np.reshape(tau, (-1, 1)), (npart, nlev)), + ) + + +def _time_varying_uxgrid(lateral_dim): + """UxGrid on a small lattice with a time-varying z, increasing along the levels, on the given lateral dimension.""" + uxgrid = sigma_coordinate_lattice_dataset(3, (0.0, 1.0), (0.0, 1.0), 4, bottom_depth=1.0).uxgrid + z = np.broadcast_to(np.arange(4.0)[:, np.newaxis], (2, 4, getattr(uxgrid, lateral_dim))) + return UxGrid(uxgrid, z=ux.UxDataArray(z, dims=("time", "zf", lateral_dim), uxgrid=uxgrid), mesh=FlatMesh()) + + +def test_ux_z_interp_method_defaults_and_validation(): + face_grid, node_grid = _time_varying_uxgrid("n_face"), _time_varying_uxgrid("n_node") + assert isinstance(face_grid.z_interp_method, UxVerticalCoordinateConstantFace) + assert isinstance(node_grid.z_interp_method, UxVerticalCoordinateLinearNode) + + with pytest.raises(ValueError, match="needs z on the mesh faces"): + node_grid.z_interp_method = UxVerticalCoordinateConstantFace() + with pytest.raises(ValueError, match="needs z on the mesh nodes"): + face_grid.z_interp_method = UxVerticalCoordinateLinearNode() + with pytest.raises(ValueError, match="must be a `VerticalCoordinateInterpolator`"): + face_grid.z_interp_method = UxConstantFaceConstantZC() + + z_1d = ux.UxDataArray(np.arange(4.0), dims=("zf",), uxgrid=face_grid.uxgrid) + grid_with_1d_z = UxGrid(face_grid.uxgrid, z=z_1d, mesh=FlatMesh()) + assert grid_with_1d_z.z_interp_method is None + with pytest.raises(ValueError, match="can only be set on a grid with a time-varying"): + grid_with_1d_z.z_interp_method = UxVerticalCoordinateLinearNode() + + with pytest.raises(ValueError, match="A 3D z must have dims"): + _time_varying_uxgrid("n_edge")