Skip to content

Instantly share code, notes, and snippets.

@arshednabeel
Created April 24, 2024 08:35
Show Gist options
  • Save arshednabeel/51dff402cdbf775a469ca9abdb3d4e87 to your computer and use it in GitHub Desktop.
Save arshednabeel/51dff402cdbf775a469ca9abdb3d4e87 to your computer and use it in GitHub Desktop.
Animate collective movement of flocks from trajectories.
'''
A simple animation function to animate collective movement of flocks, given trajectory data.
(c) Arshed Nabeel, 2024
'''
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
def animate_trajectories(
x, # (T, N, 2) array of individual positions
e, # (T, N) array of individual orientations
trail_length=10, # Number of time steps to show in the past
L=10, # Size of the domain
n_frames=None, # Number of frames to animate
filename=None, # Name of the file to save the animation
):
fig, ax = plt.subplots(figsize=(6, 6))
T, N, _ = x.shape
if n_frames is None:
n_frames = T - trail_length
start_t = 0
end_t = start_t + trail_length
fish_traj = np.empty(N, dtype=object)
fish_curr = np.empty(N, dtype=object)
for fish in range(N):
fish_traj[fish], = ax.plot(x[start_t:end_t, fish, 0],
x[start_t:end_t, fish, 1], '.',
ms=2, c='k', alpha=0.2, )#solid_joinstyle='miter')
fish_curr[fish], = ax.plot(x[end_t, fish, 0], x[end_t, fish, 1], '.', ms=5, c='r')
ax.set(aspect='equal', xlim=[0, L], ylim=[0, L],)
ax.set_axis_off()
def animate(i):
start_t = i
end_t = i + trail_length
for fish in range(N):
fish_traj[fish].set_data(x[start_t:end_t, fish, 0], x[start_t:end_t, fish, 1])
fish_curr[fish].set_data(x[end_t, fish, 0], x[end_t, fish, 1])
return *fish_traj, *fish_curr
ani = animation.FuncAnimation(fig, animate, frames=n_frames, interval=33, blit=False)
if filename is not None:
ani.save(filename, writer='ffmpeg', fps=30)
else:
plt.show()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment