Animate collective movement of flocks from trajectories.

animate_flock.py

'''
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()