Colaboratory環境下で、matplotlibを使いgif動画を保存する

Colaboratory環境下で、強化学習っぽいものを作ろうと思い、
matplotlibを使ってgif動画を作ろうとしたら、ちょっとハマりました。

メモ代わりに残しておきます。

Colaboratoryでgif保存

「matplotlib gif 保存」でググると、一般的に、matplotlibから
imagemagickというソフトを経由してgifを保存するのが簡単とのこと。

imagemagickは、pipではインストールできないものらしく、Colaboratory環境下では
以下のようにコマンドを打つと、インストールできました。

!apt-get update && apt-get install imagemagick

scatterの注意点

Colaboratoryとは関係ありませんが、plt.plotを使うと、以下のようにプロットを
追加していけば、後でgifに変換できます。

im = plt.plot(x)
ims.append(im)

しかし、plt.scatterだと以下のように、[]を付けないとエラーが出ました。

im = plt.scatter(x, y)
ims.append([im])

ボールの放物線を描く

import numpy as np
import os

import matplotlib.pyplot as plt
import matplotlib.animation as animation

#t秒後のボールの位置
def position(theta, v0, t):
    x_return = v0 * np.cos(theta) * t
    y_return = -0.5*9.81*(t**2) + v0*np.sin(theta)*t

    return x_return, y_return

xt = 0
yt = 0
t = 0
v0 = 10 #初速
theta = 45 #投げる角度
theta = theta*3.14/180 #ラジアンに変換

#gifファイルの保存先
path = 'movies/'
if not os.path.exists(path):
      os.mkdir(path)

#グラフの体裁
fig = plt.figure()
plt.ylim(0,6)
plt.xlim(0,10)
plt.grid(True)

ims = []

#ボールの演算
while(True):
    xt, yt = position(theta, v0, t)   # ボール位置の更新

    im = plt.scatter(xt, yt, c="b")   # ボールの位置をグラフにする
    ims.append([im])                  # グラフを配列 ims に追加

    #地面に付いたら終了
    if yt<0:
        break

    t+=0.05

# ロットを 50ms ごとに表示
ani = animation.ArtistAnimation(fig, ims, interval=50)
ani.save(path + "output.gif", writer="imagemagick")

無事に保存できました。