CODEWITHSUNDEEP
pythonmatplotlibanimation
KindaTechy
KindaTechy

Reputation: 1173

Multiple animations in python with matplotlib

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:

  1. The next state of the Hallway generated with go
  2. The frames are generated with draw
  3. The done condition indicates when the simulation should end
  4. 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:

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.

Upvotes: 1

Views: 2009

Answers (1)

tmdavison
tmdavison

Reputation: 69076

If you add a plt.close() call inside the if done clause of the animate function, the plot window will close when the simulation finishes, and the next window will open with the next simulation.

So that the next animation doesn't require any interaction with the mouse, we also need to add block=False to the plt.show; we can the check whether all_done is true or false, and plt.pause() if the animation is not done.

For example:

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)
            plt.close(fig)
    
    anim = FuncAnimation(fig, animate, frames=gen_frame_until_done, repeat=False)
    
    plt.show(block=False)
    while all_done is False:
        plt.pause(1)

Upvotes: 1

Related Questions