Fixing ErrorParticles and improving final plot

Author: erikvansebille

docs/examples/documentation_nesting.ipynb

@@ -259,7 +259,7 @@
     "            face_poly[np.newaxis, np.newaxis, :],\n",
     "            {\n",
     "                \"long_name\": \"Nest id\",\n",
-    "                \"description\": \"0=regional, 1=coastal\",\n",
+    "                \"description\": \"2=regional, 1=coastal, 0=harbour\",\n",
     "                \"location\": \"face\",\n",
     "                \"mesh\": \"delaunay\",\n",
     "            },\n",
@@ -425,14 +425,31 @@
     "    particles.dlon += u * particles.dt\n",
     "    particles.dlat += v * particles.dt\n",
     "\n",
+    "    # TODO particle states have to be updated manually because UVField is not called with `particles` argument (becaise of GH #2143)\n",
+    "    particles.state = np.where(\n",
+    "        np.isnan(u) | np.isnan(v),\n",
+    "        parcels.StatusCode.ErrorInterpolation,\n",
+    "        particles.state,\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "def DeleteErrorParticles(particles, fieldset):\n",
+    "    any_error = particles.state >= 50  # This captures all Errors\n",
+    "    particles[any_error].state = parcels.StatusCode.Delete\n",
+    "\n",
+    "\n",
+    "X, Y = np.meshgrid(\n",
+    "    np.linspace(1.22, 1.26, 5),\n",
+    "    np.linspace(41.02, 41.08, 4),\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",
+    "    [AdvectEE_Nests, DeleteErrorParticles],\n",
+    "    runtime=np.timedelta64(5, \"D\") - np.timedelta64(1, \"h\"),\n",
     "    dt=np.timedelta64(1, \"h\"),\n",
     "    output_file=ofile,\n",
     ")"
@@ -453,16 +470,27 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "fig, ax = plt.subplots(1, 1, figsize=(10, 5))\n",
+    "fig, ax = plt.subplots(1, 1, figsize=(10, 3))\n",
     "\n",
     "ds_out = xr.open_zarr(\"nest_particles.zarr\")\n",
     "\n",
     "cmap = plt.get_cmap(\"viridis\", 3)\n",
+    "plt.plot(ds_out.lon.T, ds_out.lat.T, \"k\", linewidth=0.5)\n",
     "sc = ax.scatter(ds_out.lon, ds_out.lat, c=ds_out.nestID, s=4, cmap=cmap, vmin=0, vmax=2)\n",
+    "xl, yl = ax.get_xlim(), ax.get_ylim()\n",
+    "\n",
+    "for rect in rectangle.values():\n",
+    "    ax.plot(\n",
+    "        np.append(rect[:, 0], rect[0, 0]), np.append(rect[:, 1], rect[0, 1]), \"-r\", lw=1\n",
+    "    )\n",
+    "ax.set_xlim(xl)\n",
+    "ax.set_ylim(yl)\n",
+    "ax.set_aspect(\"equal\")\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.tight_layout\n",
     "plt.show()"
    ]
   }