add Vesna's animation and move to examples folder

Author: reint-fischer

docs/examples_v3/tutorial_output.ipynb renamed to docs/examples/tutorial_output.ipynb

@@ -21,7 +21,7 @@
     "- [**Plotting**](#Plotting)\n",
     "- [**Animations**](#Animations)\n",
     "\n",
-    "For more advanced reading and tutorials on the analysis of Lagrangian trajectories, we recommend checking out the [Lagrangian Diagnostics](https://lagrangian-diags.readthedocs.io/en/latest/index.html) project."
+    "For more advanced reading and tutorials on the analysis of Lagrangian trajectories, we recommend checking out the [Lagrangian Diagnostics](https://lagrangian-diags.readthedocs.io/en/latest/index.html) project. The [TrajAn package]() can be used to read and plot datasets of Lagrangian trajectories."
    ]
   },
   {
@@ -128,12 +128,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "data_xarray = xr.open_zarr(\"Output.zarr\", decode_times=False)\n",
-    "# TODO: remove once decode_times issue is fixed\n",
-    "data_xarray[\"time\"].values = (\n",
-    "    data_xarray[\"time\"].values.astype(\"timedelta64[s]\").astype(\"timedelta64[ns]\")\n",
-    ")\n",
-    "data_xarray[\"time\"].values[data_xarray[\"time\"].values < 0] = np.timedelta64(\"NaT\", \"ns\")\n",
+    "data_xarray = xr.open_zarr(\"Output.zarr\")\n",
     "\n",
     "print(data_xarray)"
    ]
@@ -168,10 +163,9 @@
    "source": [
     "np.set_printoptions(linewidth=160)\n",
     "one_hour = np.timedelta64(1, \"h\")  # Define timedelta object to help with conversion\n",
+    "time_from_start = data_xarray[\"time\"].values - fieldset.time_interval.left\n",
     "\n",
-    "print(\n",
-    "    data_xarray[\"time\"].data / one_hour\n",
-    ")  # timedelta / timedelta -> float number of hours"
+    "print(time_from_start / one_hour)  # timedelta / timedelta -> float number of hours"
    ]
   },
   {
@@ -206,9 +200,9 @@
     "    np.sqrt(np.square(np.diff(x)) + np.square(np.diff(y))), axis=1\n",
     ")  # d = (dx^2 + dy^2)^(1/2)\n",
     "\n",
-    "real_time = data_xarray[\"time\"] / one_hour  # convert time to hours\n",
+    "real_time = time_from_start / one_hour  # convert time to hours\n",
     "time_since_release = (\n",
-    "    real_time.values.transpose() - real_time.values[:, 0]\n",
+    "    real_time.transpose() - real_time[:, 0]\n",
     ")  # substract the initial time from each timeseries"
    ]
   },
@@ -380,17 +374,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from IPython.display import HTML\n",
-    "from matplotlib.animation import FuncAnimation\n",
-    "\n",
-    "outputdt = timedelta(hours=2)\n",
-    "\n",
-    "# timerange in nanoseconds\n",
-    "timerange = np.arange(\n",
-    "    np.nanmin(data_xarray[\"time\"].values),\n",
-    "    np.nanmax(data_xarray[\"time\"].values) + np.timedelta64(outputdt),\n",
-    "    outputdt,\n",
-    ")"
+    "import cartopy.crs as ccrs\n",
+    "import cartopy.feature as cfeature\n",
+    "import matplotlib"
    ]
   },
   {
@@ -399,47 +385,157 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "%%capture\n",
-    "fig = plt.figure(figsize=(5, 5), constrained_layout=True)\n",
-    "ax = fig.add_subplot()\n",
-    "\n",
-    "ax.set_ylabel(\"Meridional distance [m]\")\n",
-    "ax.set_xlabel(\"Zonal distance [m]\")\n",
-    "ax.set_xlim(31, 33)\n",
-    "ax.set_ylim(-33, -30)\n",
-    "\n",
-    "# Indices of the data where time = 0\n",
-    "time_id = np.where(data_xarray[\"time\"] == timerange[0])\n",
-    "\n",
-    "scatter = ax.scatter(\n",
-    "    data_xarray[\"lon\"].values[time_id], data_xarray[\"lat\"].values[time_id]\n",
-    ")\n",
-    "\n",
-    "t = str(timerange[0].astype(\"timedelta64[h]\"))\n",
-    "title = ax.set_title(\"Particles at t = \" + t)\n",
-    "\n",
-    "\n",
-    "def animate(i):\n",
-    "    t = str(timerange[i].astype(\"timedelta64[h]\"))\n",
-    "    title.set_text(\"Particles at t = \" + t)\n",
-    "\n",
-    "    time_id = np.where(data_xarray[\"time\"] == timerange[i])\n",
-    "    scatter.set_offsets(\n",
-    "        np.c_[data_xarray[\"lon\"].values[time_id], data_xarray[\"lat\"].values[time_id]]\n",
+    "# for interactive display of animation\n",
+    "plt.rcParams[\"animation.html\"] = \"jshtml\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Number of timesteps to animate\n",
+    "nframes = 13  # use less frames for testing purposes\n",
+    "nreducedtrails = 1  # every 10th particle will have a trail (if 1, all particles have trails. Adjust for faster performance)\n",
+    "\n",
+    "\n",
+    "# Set up the colors and associated trajectories:\n",
+    "# get release times for each particle (first valide obs for each trajectory)\n",
+    "release_times = data_xarray[\"time\"].min(dim=\"obs\", skipna=True).values\n",
+    "\n",
+    "# get unique release times and assign colors\n",
+    "unique_release_times = np.unique(release_times[~np.isnat(release_times)])\n",
+    "n_release_times = len(unique_release_times)\n",
+    "print(f\"Number of unique release times: {n_release_times}\")\n",
+    "\n",
+    "# choose a continuous colormap\n",
+    "colormap = matplotlib.colormaps[\"tab20b\"]\n",
+    "\n",
+    "# set up a unique color for each release time\n",
+    "release_time_to_color = {}\n",
+    "for i, release_time in enumerate(unique_release_times):\n",
+    "    release_time_to_color[release_time] = colormap(i / max(n_release_times - 1, 1))\n",
+    "\n",
+    "\n",
+    "# --> Store data for all timeframes (this is needed for faster performance)\n",
+    "print(\"Pre-computing all particle positions...\")\n",
+    "all_particles_data = []\n",
+    "for i, target_time in enumerate(timerange):\n",
+    "    time_id = np.where(data_xarray[\"time\"] == target_time)\n",
+    "    lons = data_xarray[\"lon\"].values[time_id]\n",
+    "    lats = data_xarray[\"lat\"].values[time_id]\n",
+    "    particle_indices = time_id[0]\n",
+    "    valid = ~np.isnan(lons) & ~np.isnan(lats)\n",
+    "\n",
+    "    all_particles_data.append(\n",
+    "        {\n",
+    "            \"lons\": lons[valid],\n",
+    "            \"lats\": lats[valid],\n",
+    "            \"particle_indices\": particle_indices[valid],\n",
+    "            \"valid_count\": np.sum(valid),\n",
+    "        }\n",
     "    )\n",
     "\n",
     "\n",
-    "anim = FuncAnimation(fig, animate, frames=len(timerange), interval=100)"
+    "# figure setup\n",
+    "fig, ax = plt.subplots(figsize=(6, 5), subplot_kw={\"projection\": ccrs.PlateCarree()})\n",
+    "ax.set_xlim(30, 33)\n",
+    "ax.set_xticks(np.arange(30, 33.5, 0.5))\n",
+    "ax.set_xlabel(\"Longitude (deg E)\")\n",
+    "ax.set_ylim(-33, -30)\n",
+    "ax.set_yticks(ticks=np.arange(-33, -29.5, 0.5))\n",
+    "ax.set_yticklabels(np.arange(33, 29.5, -0.5).astype(str))\n",
+    "ax.set_ylabel(\"Latitude (deg S)\")\n",
+    "ax.coastlines(color=\"saddlebrown\")\n",
+    "ax.add_feature(cfeature.LAND, alpha=0.5, facecolor=\"saddlebrown\")\n",
+    "\n",
+    "# --> Use pre-computed data for initial setup\n",
+    "initial_data = all_particles_data[0]\n",
+    "initial_colors = []\n",
+    "for particle_idx in initial_data[\"particle_indices\"]:\n",
+    "    rt = release_times[particle_idx]\n",
+    "    if rt in release_time_to_color:\n",
+    "        initial_colors.append(release_time_to_color[rt])\n",
+    "    else:\n",
+    "        initial_colors.append(\"blue\")\n",
+    "\n",
+    "# --> plot first timestep\n",
+    "scatter = ax.scatter(initial_data[\"lons\"], initial_data[\"lats\"], s=10, c=initial_colors)\n",
+    "\n",
+    "# --> initialize trails\n",
+    "trail_plot = []\n",
+    "\n",
+    "# Set initial title\n",
+    "t_str = str(timerange[0])[:19]  # Format datetime nicely\n",
+    "title = ax.set_title(f\"Particles at t = {t_str}\")\n",
+    "\n",
+    "\n",
+    "# loop over for animation\n",
+    "def animate(i):\n",
+    "    print(f\"Animating frame {i + 1}/{len(timerange)} at time {timerange[i]}\")\n",
+    "    t_str = str(timerange[i])[:19]\n",
+    "    title.set_text(f\"Particles at t = {t_str}\")\n",
+    "\n",
+    "    # Find particles at current time\n",
+    "    current_data = all_particles_data[i]\n",
+    "\n",
+    "    if current_data[\"valid_count\"] > 0:\n",
+    "        current_colors = []\n",
+    "        for particle_idx in current_data[\"particle_indices\"]:\n",
+    "            rt = release_times[particle_idx]\n",
+    "            current_colors.append(release_time_to_color[rt])\n",
+    "\n",
+    "        scatter.set_offsets(np.c_[current_data[\"lons\"], current_data[\"lats\"]])\n",
+    "        scatter.set_color(current_colors)\n",
+    "\n",
+    "        # --> add trails\n",
+    "\n",
+    "        for trail in trail_plot:\n",
+    "            trail.remove()\n",
+    "        trail_plot.clear()\n",
+    "\n",
+    "        trail_length = min(10, i)  # trails will have max length of 10 time steps\n",
+    "\n",
+    "        if trail_length > 0:\n",
+    "            sampled_particles = current_data[\"particle_indices\"][\n",
+    "                ::nreducedtrails\n",
+    "            ]  # use all or sample if you want faster computation\n",
+    "\n",
+    "            for particle_idx in sampled_particles:\n",
+    "                trail_lons = []\n",
+    "                trail_lats = []\n",
+    "                for j in range(i - trail_length, i + 1):\n",
+    "                    past_data = all_particles_data[j]\n",
+    "                    if particle_idx in past_data[\"particle_indices\"]:\n",
+    "                        idx = np.where(past_data[\"particle_indices\"] == particle_idx)[\n",
+    "                            0\n",
+    "                        ][0]\n",
+    "                        trail_lons.append(past_data[\"lons\"][idx])\n",
+    "                        trail_lats.append(past_data[\"lats\"][idx])\n",
+    "                if len(trail_lons) > 1:\n",
+    "                    rt = release_times[particle_idx]\n",
+    "                    color = release_time_to_color[rt]\n",
+    "                    (trail,) = ax.plot(\n",
+    "                        trail_lons, trail_lats, color=color, linewidth=0.6, alpha=0.6\n",
+    "                    )\n",
+    "                    trail_plot.append(trail)\n",
+    "\n",
+    "    else:\n",
+    "        scatter.set_offsets(np.empty((0, 2)))\n",
+    "\n",
+    "\n",
+    "# Create animation\n",
+    "anim = matplotlib.animation.FuncAnimation(fig, animate, frames=nframes, interval=100)\n",
+    "anim"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "HTML(anim.to_jshtml())"
-   ]
+   "source": []
   }
  ],
  "metadata": {