I am currently working on simulating a 2D random walk problem, and I am 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 want to display the legend for each trial, and once the animation for one trial is complete, it should show the legend for the next trial. Can you please suggest how I can resolve this problem?
I am attaching below the code and the plot that I got.
%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)
alpha = 1 - trial_num / num_trials
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}")
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 + 1) * num_trials), interval = 10, repeat = False)
plt.show()
This implies the 20th step of every trial since there are num_steps = 20.
Just so that I and others understand your problem correctly, are you stating that you have multiple trials but that you only want to plot the last step of every trial and not necessarily the first 19 steps of every trial? In other words, you only want to plot the 20th step of every trial?
So, here you only want to plot a series of the final 20th step of every trial?
[20th_step_first, 20th_step_second, 20th_step_third, 20th_step_fourth, ...., and so on]
Is this what you would like to plot? If not, can you please clarify if I am not fully understanding your issue.
This is a straightforward logic problem - it’s easier to figure out than the code you’ve already written. It’s not something that needs any special Matplotlib knowledge.
The animate function runs for each step of the animation.
Try to think of a way to make part of that code only happen on certain frames, and think about which part needs to “only happen on certain frames”, and about which frames it needs to happen on.
No, I want to plot every step for each trial. I want to show the legend for each trial after the animation of the random walk for that trial is done.
But in this case, it shows the legends for every step in each trial, i.e. for every frame.
I am attaching the animation that I got for your reference.
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 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()
The FuncAnimation class allows us to create an animation by passing a function that iteratively modifies the data of a plot.
So, it means the ax variable is not created anew each animation frame. This means your plot in the ax is given new data in each call of the animate(...) while keeping the previous object intact.
It may be not visible in your animation, but the dots and lines in the plot are actually draw again and again rendered above old set of dots and lines.
TLDR;
You have to find a way to reset the ax object in a frame.
