diff --git a/docs/user_guide/getting_started/tutorial_output.ipynb b/docs/user_guide/getting_started/tutorial_output.ipynb index e6444473d..c18594bb0 100644 --- a/docs/user_guide/getting_started/tutorial_output.ipynb +++ b/docs/user_guide/getting_started/tutorial_output.ipynb @@ -398,6 +398,7 @@ "import cartopy.feature as cfeature\n", "import matplotlib\n", "from matplotlib.animation import FuncAnimation\n", + "from matplotlib.collections import LineCollection\n", "\n", "# for interactive display of animation\n", "plt.rcParams[\"animation.html\"] = \"jshtml\"" @@ -411,20 +412,13 @@ "source": [ "time_step = np.timedelta64(2, \"h\") # time step for animation frames\n", "\n", + "# Build the time axis for the animation.\n", "timerange = np.arange(\n", " np.nanmin(df_particles[\"t\"]),\n", " np.nanmax(df_particles[\"t\"]) + time_step,\n", " time_step,\n", ")\n", "\n", - "# set up a unique color for each trajectory\n", - "colormap = matplotlib.colormaps[\"tab20b\"]\n", - "trajectory_to_color = {}\n", - "for i, trajectory in enumerate(df_particles[\"particle_id\"].unique()):\n", - " trajectory_to_color[trajectory] = colormap(\n", - " i / max(len(df_particles[\"particle_id\"].unique()) - 1, 1)\n", - " )\n", - "\n", "# figure setup\n", "fig, ax = plt.subplots(figsize=(6, 5), subplot_kw={\"projection\": ccrs.PlateCarree()})\n", "ax.set_xlim(30, 33)\n", @@ -436,61 +430,57 @@ "ax.set_ylabel(\"Latitude (deg S)\")\n", "ax.coastlines()\n", "ax.add_feature(cfeature.LAND)\n", + "cmap = matplotlib.colormaps[\"tab20b\"]\n", + "\n", + "trails = LineCollection([], linewidths=0.6, alpha=0.3)\n", + "ax.add_collection(trails)\n", "\n", - "# --> plot first timestep\n", - "particles = df_particles.filter(pl.col(\"t\") == pl.lit(timerange[0]))\n", - "scatter = ax.scatter(\n", - " particles[\"x\"],\n", - " particles[\"y\"],\n", - " s=10,\n", - " c=[trajectory_to_color[p] for p in particles[\"particle_id\"]],\n", + "particle_ids = df_particles[\"particle_id\"].unique().to_list()\n", + "\n", + "# Convert the trajectory data to a time-by-particle array and assign colors.\n", + "colors = np.asarray(\n", + " [cmap(i / max(len(particle_ids) - 1, 1)) for i in range(len(particle_ids))]\n", ")\n", "\n", - "# --> initialize trails\n", - "trail_plot = []\n", + "# Create arrays to hold the x and y positions of each particle at each time step.\n", + "x = np.full((len(particle_ids), len(timerange)), np.nan)\n", + "y = np.full((len(particle_ids), len(timerange)), np.nan)\n", + "\n", + "traj = df_particles.with_columns(\n", + " pl.col(\"particle_id\")\n", + " .replace(particle_ids, range(len(particle_ids)))\n", + " .alias(\"p_idx\"),\n", + " pl.int_range(pl.len()).alias(\"row_idx\"),\n", + ")\n", + "\n", + "# Map each observation to its nearest animation time index.\n", + "traj = traj.with_columns(\n", + " pl.Series(\"t_idx\", np.searchsorted(timerange, traj[\"t\"].to_numpy())).alias(\"t_idx\")\n", + ")\n", + "\n", + "x[traj[\"p_idx\"].to_numpy(), traj[\"t_idx\"].to_numpy()] = traj[\"x\"].to_numpy()\n", + "y[traj[\"p_idx\"].to_numpy(), traj[\"t_idx\"].to_numpy()] = traj[\"y\"].to_numpy()\n", + "\n", + "# Plot first timestep\n", + "scatter = ax.scatter(x[:, 0], y[:, 0], s=10, c=colors, zorder=2)\n", "\n", "# Set initial title\n", - "t_str = pd.to_datetime(timerange[0]).strftime(\n", - " \"%Y-%m-%d %H:%M:%S\"\n", - ") # Format datetime nicely\n", + "t_str = pd.to_datetime(timerange[0]).strftime(\"%Y-%m-%d %H:%M:%S\")\n", "title = ax.set_title(f\"Particles at t = {t_str}\")\n", "\n", "\n", - "# loop over for animation\n", "def animate(i):\n", " t_str = pd.to_datetime(timerange[i]).strftime(\"%Y-%m-%d %H:%M:%S\")\n", " title.set_text(f\"Particles at t = {t_str}\")\n", "\n", - " # Find particles at current time\n", - " particles = df_particles.filter(pl.col(\"t\") == pl.lit(timerange[i]))\n", - "\n", - " if len(particles) > 0:\n", - " scatter.set_offsets(np.c_[particles[\"x\"], particles[\"y\"]])\n", - " scatter.set_color([trajectory_to_color[p] for p in particles[\"particle_id\"]])\n", - "\n", - " # --> reset trails\n", - " for trail in trail_plot:\n", - " trail.remove()\n", - " trail_plot.clear()\n", - " trail_length = min(10, i) # trails will have max length of 10 time steps\n", - " if trail_length > 0:\n", - " for traj in particles[\"particle_id\"].unique():\n", - " traj_trail = df_particles.filter(\n", - " (pl.col(\"particle_id\") == traj)\n", - " & (pl.col(\"t\") >= pl.lit(timerange[max(0, i - trail_length)]))\n", - " & (pl.col(\"t\") <= pl.lit(timerange[i]))\n", - " )\n", - " if len(traj_trail) > 1:\n", - " (trail,) = ax.plot(\n", - " traj_trail[\"x\"],\n", - " traj_trail[\"y\"],\n", - " color=trajectory_to_color[traj],\n", - " linewidth=0.6,\n", - " alpha=0.6,\n", - " )\n", - " trail_plot.append(trail)\n", - " else:\n", - " scatter.set_offsets(np.empty((0, 2)))\n", + " scatter.set_offsets(np.column_stack((x[:, i], y[:, i])))\n", + "\n", + " trail_length = min(10, i) # trails have max length of 10 time steps\n", + " start = max(0, i - trail_length)\n", + " x_slice = x[:, start : i + 1]\n", + " y_slice = y[:, start : i + 1]\n", + " trails.set_segments(np.dstack((x_slice, y_slice)))\n", + " trails.set_color(colors)\n", "\n", "\n", "# Create animation\n", @@ -503,7 +493,7 @@ "metadata": { "celltoolbar": "Metagegevens bewerken", "kernelspec": { - "display_name": "Parcels:docs (3.14.6)", + "display_name": "Python 3", "language": "python", "name": "python3" },