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.
from
Multiple animations in python with matplotlib
