I would like to use matplotlib's animation capabilities to display and save multiple animations. Previously, I was just using pyplot with a short pause at each step to fake the animation, but I did not find a way to save these "animations" as videos, so I'm switching to using the real animations. Here is a dummy version of my code (which will run) when I started:
from matplotlib import pyplot as plt import numpy as np class Hallway: def __init__(self): self.end_pos = 5 self.cur_pos = 0 def setup(self): self.cur_pos = 0 def go(self, decision): if decision == 0 and self.cur_pos > 0: self.cur_pos -= 1 elif decision == 1: self.cur_pos += 1 done = self.cur_pos >= self.end_pos return done def draw(self, fig): fig.clear() ax = fig.gca() ax.set(xlim=(-0.5, 5.5), ylim=(-0.5, 0.5)) ax.scatter(self.cur_pos, 0., s=350) plt.draw() plt.pause(0.01) sim = Hallway() for num_sim in range(5): print("running simulation {}".format(num_sim)) sim.setup() sim.draw(plt.gcf()) while True: done = sim.go(np.random.randint(0,2)) sim.draw(plt.gcf()) if done: break # Save animation here Key things to note in here:
- The next state of the Hallway generated with
go - The frames are generated with
draw - The done condition indicates when the simulation should end
- Once an animation ends, I want to save it, but we're not done! After saving the animation, I want to launch a new one. This will happen 5 times with the outer loop.
So I changed my code around so that I could use an animation object, and this is what it is now:
from matplotlib import pyplot as plt from matplotlib.animation import FuncAnimation import numpy as np class Hallway: def __init__(self): self.end_pos = 5 self.cur_pos = 0 def setup(self): self.cur_pos = 0 def go(self, decision): if decision == 0 and self.cur_pos > 0: self.cur_pos -= 1 elif decision == 1: self.cur_pos += 1 done = self.cur_pos >= self.end_pos return done def draw(self, fig): fig.clear() ax = fig.gca() ax.set(xlim=(-0.5, 5.5), ylim=(-0.5, 0.5)) ax.scatter(self.cur_pos, 0., s=350) plt.draw() plt.pause(0.01) sim = Hallway() for num_sim in range(5): print("running simulation {}".format(num_sim)) sim.setup() all_done = False fig = plt.figure() def gen_frame_until_done(): # Using a generator to give me frames as long as not all_done global all_done i = 0 while not all_done: i += 1 yield i def animate(i): # Animate function takes the place of the while loop sim.draw(fig) done = sim.go(np.random.randint(0,2)) if done: global all_done all_done = True sim.draw(fig) anim = FuncAnimation(fig, animate, frames=gen_frame_until_done, repeat=False) plt.show() # anim.save(...) This will run, but it won't quite give me what I want. It will show only one animation and the terminal will show running simulation 0. When all_done is triggered and the simulation is over, the program will wait for me to exit the plot window. Once I exit, the program will continue to the next simulation and repeat.
I don't like that I have to manually exit the window. I got a little hack semi-working by replacing the blocking plt.show() with
plt.show(block=False) plt.pause(3) plt.close() This will allow the program to continue without having to manually exit the window. However, it will only allow 3 seconds of the animation to display before going on to the next one.
What I want:
- I want to be able to display the simulation until it is over. When it is over, I want the window to automatically close.
- I want the next simulation to run with a new animation window right after the previous one.
Again, I'm using the animation objects because I need to be able to save the animations as videos. But if there's another way to do this, I'm definitely open to it.
