Merge branch 'docs/tutorial_fesom' of github.com:Parcels-Code/Parcels into docs/tutorial_fesom

Author: fluidnumericsJoe

docs/user_guide/examples/tutorial_fesom.ipynb

@@ -26,7 +26,17 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "from pathlib import Path\n\nimport matplotlib.pyplot as plt\nimport matplotlib.tri as mtri\nimport numpy as np\nimport uxarray as ux\n\nimport parcels\nimport parcels.tutorial"
+   "source": [
+    "from pathlib import Path\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import matplotlib.tri as mtri\n",
+    "import numpy as np\n",
+    "import uxarray as ux\n",
+    "\n",
+    "import parcels\n",
+    "import parcels.tutorial"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -46,14 +56,15 @@
    "outputs": [],
    "source": [
     "for name in [\n",
-    "    \"FESOM_periodic_channel/fesom_channel\",    # grid description\n",
+    "    \"FESOM_periodic_channel/fesom_channel\",  # grid description\n",
     "    \"FESOM_periodic_channel/u.fesom_channel\",  # zonal velocity (face-registered)\n",
     "    \"FESOM_periodic_channel/v.fesom_channel\",  # meridional velocity (face-registered)\n",
     "    \"FESOM_periodic_channel/w.fesom_channel\",  # vertical velocity (node-registered)\n",
     "]:\n",
     "    parcels.tutorial.open_dataset(name)\n",
     "\n",
     "from parcels._datasets.remote import _DATA_HOME\n",
+    "\n",
     "data_dir = Path(_DATA_HOME) / \"data\" / \"FESOM_periodic_channel\"\n",
     "\n",
     "grid_path = str(data_dir / \"fesom_channel.nc\")\n",
@@ -88,7 +99,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ds = ux.open_mfdataset(grid_path, data_paths).rename_vars({\"u\": \"U\", \"v\": \"V\", \"w\": \"W\"})\n",
+    "ds = ux.open_mfdataset(grid_path, data_paths).rename_vars(\n",
+    "    {\"u\": \"U\", \"v\": \"V\", \"w\": \"W\"}\n",
+    ")\n",
     "ds"
    ]
   },
@@ -113,7 +126,10 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "ds = parcels.convert.fesom_to_ugrid(ds)\nprint(\"dims:\", dict(ds.sizes))"
+   "source": [
+    "ds = parcels.convert.fesom_to_ugrid(ds)\n",
+    "print(\"dims:\", dict(ds.sizes))"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -155,7 +171,35 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "lon_grid, lat_grid = np.meshgrid(\n    np.linspace(0.5, 4.0, 10),\n    np.linspace(3.0, 15.0, 4),\n)\nlon = lon_grid.ravel()\nlat = lat_grid.ravel()\nz = np.full(lon.size, 50.0)  # release at 50 m depth\n\npset = parcels.ParticleSet(\n    fieldset=fieldset, pclass=parcels.Particle, lon=lon, lat=lat, z=z,\n)\n\noutput_file = parcels.ParticleFile(\"output-fesom.parquet\", outputdt=np.timedelta64(1, \"h\"))\n\npset.execute(\n    [parcels.kernels.AdvectionRK4],\n    runtime=np.timedelta64(2, \"D\"),\n    dt=np.timedelta64(5, \"m\"),\n    output_file=output_file,\n    verbose_progress=False,\n)"
+   "source": [
+    "lon_grid, lat_grid = np.meshgrid(\n",
+    "    np.linspace(0.5, 4.0, 10),\n",
+    "    np.linspace(3.0, 15.0, 4),\n",
+    ")\n",
+    "lon = lon_grid.ravel()\n",
+    "lat = lat_grid.ravel()\n",
+    "z = np.full(lon.size, 50.0)  # release at 50 m depth\n",
+    "\n",
+    "pset = parcels.ParticleSet(\n",
+    "    fieldset=fieldset,\n",
+    "    pclass=parcels.Particle,\n",
+    "    lon=lon,\n",
+    "    lat=lat,\n",
+    "    z=z,\n",
+    ")\n",
+    "\n",
+    "output_file = parcels.ParticleFile(\n",
+    "    \"output-fesom.parquet\", outputdt=np.timedelta64(1, \"h\")\n",
+    ")\n",
+    "\n",
+    "pset.execute(\n",
+    "    [parcels.kernels.AdvectionRK4],\n",
+    "    runtime=np.timedelta64(2, \"D\"),\n",
+    "    dt=np.timedelta64(5, \"m\"),\n",
+    "    output_file=output_file,\n",
+    "    verbose_progress=False,\n",
+    ")"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -167,7 +211,70 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "df = parcels.read_particlefile(\"output-fesom.parquet\")\n\ntriang = mtri.Triangulation(\n    ds.uxgrid.node_lon.values,\n    ds.uxgrid.node_lat.values,\n    triangles=ds.uxgrid.face_node_connectivity.values,\n)\n\ndepth_idx = int(np.argmin(np.abs(ds.zc.values - 50.0)))\nU_face = np.asarray(ds[\"U\"].isel(zc=depth_idx)).squeeze()\nV_face = np.asarray(ds[\"V\"].isel(zc=depth_idx)).squeeze()\nspeed = np.hypot(U_face, V_face)\n\nfig, ax = plt.subplots(figsize=(11, 5))\n\n# Background: speed at the release depth.\ntpc = ax.tripcolor(triang, facecolors=speed, shading=\"flat\", cmap=\"Blues\")\nfig.colorbar(tpc, ax=ax, label=\"speed [m/s]\", location=\"left\", shrink=0.85)\n\n# Velocity arrows at face centres. Drawing in lon/lat space (angles/scale_units\n# \"xy\") keeps the arrows aligned with the trajectories; length tracks speed.\nstep = max(1, U_face.size // 400)\nxf = ds.uxgrid.face_lon.values\nyf = ds.uxgrid.face_lat.values\nmax_speed = float(np.nanmax(speed))\nq = ax.quiver(\n    xf[::step], yf[::step], U_face[::step], V_face[::step],\n    angles=\"xy\", scale_units=\"xy\", scale=max_speed / 0.3,\n    color=\"k\", width=0.0018, pivot=\"tail\",\n)\nax.quiverkey(q, 0.86, 1.04, max_speed, f\"{max_speed:.2f} m/s\", labelpos=\"E\", coordinates=\"axes\")\n\n# Particle paths (grey lines) and positions coloured by time.\nfor traj in df.sort(\"time\").partition_by(\"particle_id\"):\n    ax.plot(traj[\"lon\"], traj[\"lat\"], color=\"0.4\", linewidth=0.6, alpha=0.7, zorder=2)\nax.scatter(lon, lat, facecolors=\"none\", edgecolors=\"k\", s=60, label=\"release\", zorder=3)\nsc = ax.scatter(\n    df[\"lon\"], df[\"lat\"], c=df[\"time\"].dt.total_seconds(), s=6, cmap=\"viridis\", zorder=3,\n)\nfig.colorbar(sc, ax=ax, label=\"time since release [s]\", shrink=0.85)\n\nax.set_xlabel(\"Longitude [deg E]\")\nax.set_ylabel(\"Latitude [deg N]\")\nax.set_title(f\"FESOM2 velocity at z ≈ {float(ds.zc.values[depth_idx]):.1f} m with particle trajectories\")\nax.legend(loc=\"upper right\")\nplt.show()"
+   "source": [
+    "df = parcels.read_particlefile(\"output-fesom.parquet\")\n",
+    "\n",
+    "triang = mtri.Triangulation(\n",
+    "    ds.uxgrid.node_lon.values,\n",
+    "    ds.uxgrid.node_lat.values,\n",
+    "    triangles=ds.uxgrid.face_node_connectivity.values,\n",
+    ")\n",
+    "\n",
+    "depth_idx = int(np.argmin(np.abs(ds.zc.values - 50.0)))\n",
+    "U_face = np.asarray(ds[\"U\"].isel(zc=depth_idx)).squeeze()\n",
+    "V_face = np.asarray(ds[\"V\"].isel(zc=depth_idx)).squeeze()\n",
+    "speed = np.hypot(U_face, V_face)\n",
+    "\n",
+    "fig, ax = plt.subplots(figsize=(11, 5))\n",
+    "\n",
+    "# Background: speed at the release depth.\n",
+    "tpc = ax.tripcolor(triang, facecolors=speed, shading=\"flat\", cmap=\"Blues\")\n",
+    "fig.colorbar(tpc, ax=ax, label=\"speed [m/s]\", location=\"left\", shrink=0.85)\n",
+    "\n",
+    "# Velocity arrows at face centres. Drawing in lon/lat space (angles/scale_units\n",
+    "# \"xy\") keeps the arrows aligned with the trajectories; length tracks speed.\n",
+    "step = max(1, U_face.size // 400)\n",
+    "xf = ds.uxgrid.face_lon.values\n",
+    "yf = ds.uxgrid.face_lat.values\n",
+    "max_speed = float(np.nanmax(speed))\n",
+    "q = ax.quiver(\n",
+    "    xf[::step],\n",
+    "    yf[::step],\n",
+    "    U_face[::step],\n",
+    "    V_face[::step],\n",
+    "    angles=\"xy\",\n",
+    "    scale_units=\"xy\",\n",
+    "    scale=max_speed / 0.3,\n",
+    "    color=\"k\",\n",
+    "    width=0.0018,\n",
+    "    pivot=\"tail\",\n",
+    ")\n",
+    "ax.quiverkey(\n",
+    "    q, 0.86, 1.04, max_speed, f\"{max_speed:.2f} m/s\", labelpos=\"E\", coordinates=\"axes\"\n",
+    ")\n",
+    "\n",
+    "# Particle paths (grey lines) and positions coloured by time.\n",
+    "for traj in df.sort(\"time\").partition_by(\"particle_id\"):\n",
+    "    ax.plot(traj[\"lon\"], traj[\"lat\"], color=\"0.4\", linewidth=0.6, alpha=0.7, zorder=2)\n",
+    "ax.scatter(lon, lat, facecolors=\"none\", edgecolors=\"k\", s=60, label=\"release\", zorder=3)\n",
+    "sc = ax.scatter(\n",
+    "    df[\"lon\"],\n",
+    "    df[\"lat\"],\n",
+    "    c=df[\"time\"].dt.total_seconds(),\n",
+    "    s=6,\n",
+    "    cmap=\"viridis\",\n",
+    "    zorder=3,\n",
+    ")\n",
+    "fig.colorbar(sc, ax=ax, label=\"time since release [s]\", shrink=0.85)\n",
+    "\n",
+    "ax.set_xlabel(\"Longitude [deg E]\")\n",
+    "ax.set_ylabel(\"Latitude [deg N]\")\n",
+    "ax.set_title(\n",
+    "    f\"FESOM2 velocity at z ≈ {float(ds.zc.values[depth_idx]):.1f} m with particle trajectories\"\n",
+    ")\n",
+    "ax.legend(loc=\"upper right\")\n",
+    "plt.show()"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -198,4 +305,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}