Блог

Я хотел бы использовать анимационные возможности matplotlib для отображения и сохранения нескольких анимаций. Ранее я просто использовал pyplot с короткой паузой на каждом шаге, чтобы подделать анимацию, но я не нашел способа сохранить эти «анимации» в виде видео, поэтому я переключаюсь на использование реальных анимаций. Вот фиктивная версия моего кода (которая будет выполняться) при запуске:

 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
  

Здесь следует отметить ключевые моменты:

1. Следующее состояние коридора, сгенерированное с go
2. Кадры генерируются с помощью draw
3. Условие «Выполнено» указывает, когда симуляция должна завершиться
4. Как только анимация заканчивается, я хочу сохранить ее, но мы еще не закончили! После сохранения анимации я хочу запустить новую. Это произойдет 5 раз с внешним циклом.

Итак, я изменил свой код, чтобы я мог использовать объект анимации, и это то, чем он сейчас является:

 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(...)
  

Это будет запущено, но это не совсем даст мне то, что я хочу. Он покажет только одну анимацию, и терминал покажет running simulation 0 . Когда all_done сработает и симуляция закончится, программа будет ждать, пока я выйду из окна графика. Как только я выйду, программа перейдет к следующему моделированию и повторится.

Мне не нравится, что мне приходится вручную выходить из окна. Я получил небольшой полурабочий хак, заменив блокировку plt.show() на

 plt.show(block=False)
plt.pause(3)
plt.close()
  

Это позволит программе продолжить работу без необходимости вручную выходить из окна. Однако это позволит отобразить анимацию только на 3 секунды, прежде чем переходить к следующей.

Чего я хочу:

• Я хочу иметь возможность отображать симуляцию, пока она не закончится. Когда это закончится, я хочу, чтобы окно автоматически закрылось.
• Я хочу, чтобы следующее моделирование запускалось с новым окном анимации сразу после предыдущего.

Опять же, я использую объекты анимации, потому что мне нужно иметь возможность сохранять анимации в виде видео. Но если есть другой способ сделать это, я определенно открыт для этого.

Если вы добавите plt.close() вызов внутри if done предложения animate функции, окно построения закроется, когда симуляция завершится, и откроется следующее окно со следующей симуляцией.

Чтобы следующая анимация не требовала какого-либо взаимодействия с мышью, нам также нужно добавить block=False в plt.show ; мы можем проверить, all_done является ли это true или false, и plt.pause() если анимация не выполнена.

Например:

 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)