Plot legends are being repeated for every frame during matplotlib animation

Community
1

I’m currently working on simulating a 2D random walk problem and facing an issue with the legend display. The legends for each trial are shown for every step I take instead of only once the animation for one trial is finished. I would like to display the legend for each trial only once the animation for that specific trial is complete. Can you please suggest how I can resolve this problem?
I have attached the code and animation below.

%clear
%matplotlib notebook
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

stepsize = 0.5
num_steps = 20
num_trials = 5

final_position = []

for _ in range(num_trials):
    pos = np.array([0, 0])
    path = []
    for i in range(num_steps):
        pos = pos + np.random.normal(0, stepsize, 2)
        path.append(pos)
    final_position.append(np.array(path))
    
x = [final_position[i][:,0] for i in range(len(final_position))]
y = [final_position[j][:,1] for j in range(len(final_position))]

fig = plt.figure(figsize=(10,6))
ax = fig.add_subplot()


cmap = plt.get_cmap('tab10')

def animate(frame):
    step_num = frame % (num_steps+1)
    trial_num = frame//(num_steps+1)
    color = cmap(trial_num % 10)
    ax.plot(x[trial_num][:step_num], y[trial_num][:step_num], color = color, ls = '-',linewidth = 0.5,
            marker = 'o', ms = 8, mfc = color, mec ='k', zorder = trial_num, label = f"Trial = {trial_num+1}")
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_title(f"Number of trials = {trial_num+1} \nNumber of steps = {step_num+1}")  
    ax.legend(fontsize=10, bbox_to_anchor=(1, 1))
    ax.grid(True)
    
    return ax

fig.suptitle(f"2D random walk simulation for {num_steps} steps over {num_trials} trials.")
ani = FuncAnimation(fig, animate, frames= np.arange(0, (num_steps +  * num_trials)), interval = 100, repeat = False)
plt.show()

false_trial1

I tried to put up a condition that whenever (num_step+1) == step_num the legends will appear. However, the plot displays 20 legends after each trial animation instead of only one legend indicating the random walk for that particular trial. This means it displays 20 legends (the number of steps) every time after the animation for each trial is completed.

I am attaching the code animation for this case below.

def animate(frame):
    global num_steps, num_trials
    step_num = frame % (num_steps)
    trial_num = frame//(num_steps)
    color = cmap(trial_num % 10)
    
    ax.plot(x[trial_num][:step_num], y[trial_num][:step_num], color = color, ls = '-',linewidth = 0.5,
            marker = 'o', ms = 8, mfc = color, mec ='k', zorder = trial_num, label = f"Trial = {trial_num+1}")  
    if (step_num+1) == num_steps:
        ax.legend(fontsize=10, bbox_to_anchor=(1, 1))
        
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_title(f"Number of trials = {trial_num+1} \nNumber of steps = {step_num+1}")  
    ax.grid(True)
    
    return ax

fig.suptitle(f"2D random walk simulation for {num_steps} steps over {num_trials} trials.")
ani = FuncAnimation(fig, animate, frames= np.arange(0,num_steps * num_trials), interval = 50, repeat = False)
plt.show()

false_tria

2

One option is to only specify the label to plot when plotting the last step of each trial. E.g.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

stepsize = 0.5
num_steps = 20
num_trials = 5

final_position = []

for _ in range(num_trials):
    pos = np.array([0, 0])
    path = []
    for i in range(num_steps):
        pos = pos + np.random.normal(0, stepsize, 2)
        path.append(pos)
    final_position.append(np.array(path))
    
x = [final_position[i][:,0] for i in range(len(final_position))]
y = [final_position[j][:,1] for j in range(len(final_position))]

fig = plt.figure(figsize=(10,6))
ax = fig.add_subplot()
fig.subplots_adjust(left=0.1, right=0.85)

cmap = plt.get_cmap('tab10')

def animate(frame):
    step_num = frame % (num_steps)
    trial_num = frame//(num_steps)
    color = cmap(trial_num % 10)
    if step_num == num_steps-1:
        label = f"Trial = {trial_num+1}"
    else:
        label = None
    ax.plot(x[trial_num][:step_num], y[trial_num][:step_num], color = color, ls = '-',linewidth = 0.5,
            marker = 'o', ms = 8, mfc = color, mec ='k', zorder = trial_num, label = label)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_title(f"Number of trials = {trial_num+1} \nNumber of steps = {step_num+1}")  
    if step_num == num_steps-1:
        ax.legend(fontsize=10, loc='upper left', bbox_to_anchor=(1, 1))
    ax.grid(True)
    
    return ax

fig.suptitle(f"2D random walk simulation for {num_steps} steps over {num_trials} trials.")
ani = FuncAnimation(fig, animate, frames= np.arange(0, (num_steps * num_trials)), interval = 100, repeat = False)
ani.save('random_walk.gif')

random_walk

1 Like
3

Great! Thanks for the help :slight_smile: