Adding Particle Advection to nesting tutorial

Author: erikvansebille

docs/examples/documentation_nesting.ipynb

@@ -357,6 +357,114 @@
     "plt.tight_layout()\n",
     "plt.show()"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "19",
+   "metadata": {},
+   "source": [
+    "## Advecting particles with nest transitions\n",
+    "\n",
+    "We can now set up a particle advection simulation using the nested grids. We first combine all the Fields into a single FieldSet. \n",
+    "\n",
+    "We rename the individual Fields by appending the nest index to their names, so that we can easily identify which Field belongs to which nest. We also add the `NestID` Field to the FieldSet (note that Parcels v4 supports combining structured and unstructured Fields into one FieldSet, which is very convenient for this usecase)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "20",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fields = [NestID]\n",
+    "for i, ds in enumerate(ds_in.values()):\n",
+    "    # TODO : remove depth dimension when Parcels supports 2D copernicusmarine datasets\n",
+    "    ds = ds.assign_coords(depth=np.array([0]))\n",
+    "    ds[\"depth\"].attrs[\"axis\"] = \"Z\"\n",
+    "\n",
+    "    fset = parcels.FieldSet.from_copernicusmarine(ds)\n",
+    "    for fld in fset.fields.values():\n",
+    "        fld.name = f\"{fld.name}{i}\"\n",
+    "        fields.append(fld)\n",
+    "fieldset = parcels.FieldSet(fields)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "21",
+   "metadata": {},
+   "source": [
+    "We then define a custom Advection kernel that advects particles using the appropriate velocity Field based on the `NestID` at the particle's location. Note that for simplicity, we use Eulerian advection here, but in a real-world application you would typically use a higher-order scheme."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "22",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def AdvectEE_Nests(particles, fieldset):\n",
+    "    particles.nestID = fieldset.NestID[particles]\n",
+    "\n",
+    "    # TODO because of KernelParticle bug (GH #2143), we need to copy lon/lat/time to local variables\n",
+    "    time = particles.time\n",
+    "    z = particles.z\n",
+    "    lat = particles.lat\n",
+    "    lon = particles.lon\n",
+    "    u = np.zeros_like(particles.lon)\n",
+    "    v = np.zeros_like(particles.lat)\n",
+    "\n",
+    "    unique_ids = np.unique(particles.nestID)\n",
+    "    for nid in unique_ids:\n",
+    "        mask = particles.nestID == nid\n",
+    "        UVField = getattr(fieldset, f\"UV{nid}\")\n",
+    "        (u[mask], v[mask]) = UVField[time[mask], z[mask], lat[mask], lon[mask]]\n",
+    "\n",
+    "    particles.dlon += u * particles.dt\n",
+    "    particles.dlat += v * particles.dt\n",
+    "\n",
+    "\n",
+    "pset = parcels.ParticleSet(\n",
+    "    fieldset, pclass=NestParticle, lon=X.flatten(), lat=Y.flatten()\n",
+    ")\n",
+    "ofile = parcels.ParticleFile(\"nest_particles.zarr\", outputdt=np.timedelta64(1, \"h\"))\n",
+    "pset.execute(\n",
+    "    AdvectEE_Nests,\n",
+    "    runtime=np.timedelta64(24, \"h\"),\n",
+    "    dt=np.timedelta64(1, \"h\"),\n",
+    "    output_file=ofile,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "23",
+   "metadata": {},
+   "source": [
+    "And finally we plot the particles moving through the nested grids, confirming that they correctly transition between the nests based on their location."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "24",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(1, 1, figsize=(10, 5))\n",
+    "\n",
+    "ds_out = xr.open_zarr(\"nest_particles.zarr\")\n",
+    "\n",
+    "cmap = plt.get_cmap(\"viridis\", 3)\n",
+    "sc = ax.scatter(ds_out.lon, ds_out.lat, c=ds_out.nestID, s=4, cmap=cmap, vmin=0, vmax=2)\n",
+    "\n",
+    "cbar = plt.colorbar(sc, ticks=[0, 1, 2], ax=ax)\n",
+    "cbar.set_label(\"Nest ID\")\n",
+    "ax.set_title(\"Particle advection through nests\")\n",
+    "plt.show()"
+   ]
   }
  ],
  "metadata": {