Skip to content

Improving performance of animation example #2798

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
erikvansebille wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Improving performance of animation example #2798

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
700143e
Improving performance of animation example
erikvansebille Aug 4, 2026
2aa5e91
Transposing array for further simplification of tail generation
erikvansebille Aug 4, 2026
41b6ad7
Renaming trail_collection -> trails
erikvansebille Aug 4, 2026
88364d6
Removing for-loops in generation of arrays
erikvansebille Aug 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 42 additions & 52 deletions docs/user_guide/getting_started/tutorial_output.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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\""
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -503,7 +493,7 @@
"metadata": {
"celltoolbar": "Metagegevens bewerken",
"kernelspec": {
"display_name": "Parcels:docs (3.14.6)",
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
Expand Down