From 700143e381f8669f5110d75c1b6dc80df4508fe9 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 4 Aug 2026 09:54:17 +0200 Subject: [PATCH 1/4] Improving performance of animation example By avoiding for-loop in polars frame for tails In large output files, this can reduce animation generation time from 2 hours to 6 minutes --- .../getting_started/tutorial_output.ipynb | 86 ++++++++----------- 1 file changed, 34 insertions(+), 52 deletions(-) diff --git a/docs/user_guide/getting_started/tutorial_output.ipynb b/docs/user_guide/getting_started/tutorial_output.ipynb index e6444473d..fe81dca9a 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,49 @@ "ax.set_ylabel(\"Latitude (deg S)\")\n", "ax.coastlines()\n", "ax.add_feature(cfeature.LAND)\n", + "cmap = matplotlib.colormaps[\"tab20b\"]\n", + "\n", + "trail_collection = LineCollection([], linewidths=0.6, alpha=0.3)\n", + "ax.add_collection(trail_collection)\n", + "\n", + "particle_ids = df_particles[\"particle_id\"].unique().to_list()\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", + "# 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", + "x = np.full((len(timerange), len(particle_ids)), np.nan)\n", + "y = np.full((len(timerange), len(particle_ids)), np.nan)\n", + "\n", + "for ti, t in enumerate(timerange):\n", + " frame = df_particles.filter(pl.col(\"t\") == pl.lit(t)).sort(\"particle_id\")\n", + " for row in frame.iter_rows(named=True):\n", + " pi = particle_ids.index(row[\"particle_id\"])\n", + " x[ti, pi] = row[\"x\"]\n", + " y[ti, pi] = row[\"y\"]\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", + " trail_segments = np.stack((x_slice, y_slice), axis=-1).transpose(1, 0, 2)\n", + " trail_collection.set_segments(trail_segments)\n", + " trail_collection.set_color(colors)\n", "\n", "\n", "# Create animation\n", @@ -503,7 +485,7 @@ "metadata": { "celltoolbar": "Metagegevens bewerken", "kernelspec": { - "display_name": "Parcels:docs (3.14.6)", + "display_name": "Python 3", "language": "python", "name": "python3" }, From 2aa5e9105b4e8e1146b2e957dc3020af5aca1118 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 4 Aug 2026 13:18:25 +0200 Subject: [PATCH 2/4] Transposing array for further simplification of tail generation Avoids transpose on tail data --- .../getting_started/tutorial_output.ipynb | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/docs/user_guide/getting_started/tutorial_output.ipynb b/docs/user_guide/getting_started/tutorial_output.ipynb index fe81dca9a..70f59185b 100644 --- a/docs/user_guide/getting_started/tutorial_output.ipynb +++ b/docs/user_guide/getting_started/tutorial_output.ipynb @@ -442,18 +442,18 @@ " [cmap(i / max(len(particle_ids) - 1, 1)) for i in range(len(particle_ids))]\n", ")\n", "\n", - "x = np.full((len(timerange), len(particle_ids)), np.nan)\n", - "y = np.full((len(timerange), len(particle_ids)), np.nan)\n", + "x = np.full((len(particle_ids), len(timerange)), np.nan)\n", + "y = np.full((len(particle_ids), len(timerange)), np.nan)\n", "\n", "for ti, t in enumerate(timerange):\n", " frame = df_particles.filter(pl.col(\"t\") == pl.lit(t)).sort(\"particle_id\")\n", " for row in frame.iter_rows(named=True):\n", " pi = particle_ids.index(row[\"particle_id\"])\n", - " x[ti, pi] = row[\"x\"]\n", - " y[ti, pi] = row[\"y\"]\n", + " x[pi, ti] = row[\"x\"]\n", + " y[pi, ti] = row[\"y\"]\n", "\n", "# Plot first timestep\n", - "scatter = ax.scatter(x[0, :], y[0, :], s=10, c=colors, zorder=2)\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(\"%Y-%m-%d %H:%M:%S\")\n", @@ -464,14 +464,13 @@ " 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", - " scatter.set_offsets(np.column_stack((x[i, :], y[i, :])))\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", - " trail_segments = np.stack((x_slice, y_slice), axis=-1).transpose(1, 0, 2)\n", - " trail_collection.set_segments(trail_segments)\n", + " x_slice = x[:, start : i + 1]\n", + " y_slice = y[:, start : i + 1]\n", + " trail_collection.set_segments(np.dstack((x_slice, y_slice)))\n", " trail_collection.set_color(colors)\n", "\n", "\n", From 41b6ad771b11ee34ebea7e049acca0282247b83b Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 4 Aug 2026 13:21:11 +0200 Subject: [PATCH 3/4] Renaming trail_collection -> trails --- docs/user_guide/getting_started/tutorial_output.ipynb | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/user_guide/getting_started/tutorial_output.ipynb b/docs/user_guide/getting_started/tutorial_output.ipynb index 70f59185b..2f4748e70 100644 --- a/docs/user_guide/getting_started/tutorial_output.ipynb +++ b/docs/user_guide/getting_started/tutorial_output.ipynb @@ -432,8 +432,8 @@ "ax.add_feature(cfeature.LAND)\n", "cmap = matplotlib.colormaps[\"tab20b\"]\n", "\n", - "trail_collection = LineCollection([], linewidths=0.6, alpha=0.3)\n", - "ax.add_collection(trail_collection)\n", + "trails = LineCollection([], linewidths=0.6, alpha=0.3)\n", + "ax.add_collection(trails)\n", "\n", "particle_ids = df_particles[\"particle_id\"].unique().to_list()\n", "\n", @@ -470,8 +470,8 @@ " start = max(0, i - trail_length)\n", " x_slice = x[:, start : i + 1]\n", " y_slice = y[:, start : i + 1]\n", - " trail_collection.set_segments(np.dstack((x_slice, y_slice)))\n", - " trail_collection.set_color(colors)\n", + " trails.set_segments(np.dstack((x_slice, y_slice)))\n", + " trails.set_color(colors)\n", "\n", "\n", "# Create animation\n", From 88364d632518109f49204d9ca85c2d437833379e Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 4 Aug 2026 13:35:53 +0200 Subject: [PATCH 4/4] Removing for-loops in generation of arrays --- .../getting_started/tutorial_output.ipynb | 21 +++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/docs/user_guide/getting_started/tutorial_output.ipynb b/docs/user_guide/getting_started/tutorial_output.ipynb index 2f4748e70..c18594bb0 100644 --- a/docs/user_guide/getting_started/tutorial_output.ipynb +++ b/docs/user_guide/getting_started/tutorial_output.ipynb @@ -442,15 +442,24 @@ " [cmap(i / max(len(particle_ids) - 1, 1)) for i in range(len(particle_ids))]\n", ")\n", "\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", - "for ti, t in enumerate(timerange):\n", - " frame = df_particles.filter(pl.col(\"t\") == pl.lit(t)).sort(\"particle_id\")\n", - " for row in frame.iter_rows(named=True):\n", - " pi = particle_ids.index(row[\"particle_id\"])\n", - " x[pi, ti] = row[\"x\"]\n", - " y[pi, ti] = row[\"y\"]\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",