python + matplotlib でグラフのアニメーション

前回に引き続き、備忘録の第二弾です!

去年やっていたプロジェクトの最初の頃に、 作ったアルゴリズムの動作確認のために使って、 お世話になりました!

python + matplotlib においては、 グラフのアニメーションもお手の物なんだけど、 いろいろと注意することがあったりする

環境

• python 3.5.0

matplotlib のインストール

pip install matplotlib

プログラム

#!python
#-*-coding:utf-8-*-
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as anim

def main():
    fig = plt.figure()
    ax = fig.add_subplot(111)

    frames = []
    for dx in np.arange(1, 10, 0.1):
        x = np.linspace(0, 5, 100)
        y = np.sin(x + dx)
        tmp_frame, = ax.plot(x, y, "b-")
        frames.append([tmp_frame])

    ax.axis([0, 5, -1.1, 1.1])
    ani = anim.ArtistAnimation(fig, frames, interval=1, repeat=True, repeat_delay=1000)

    plt.show(block=False)
    input("Enter to close")
    plt.close()

    # ani.save("test.mp4", writer="ffmpeg", fps=30, bitrate=1000)

    print("end")

if __name__ == '__main__':
    main()

コード中のコメントアウトしている箇所で、動画を保存できる。
今回のコードでできた動画はこんな感じ↓

注意すべきところなど

matplotlibのアニメーション

matplotlibでは、FuncAnimationとArtistAnimationという、 2通りの関数が用意されている。
FuncAnimationは、動的にグラフを書き換えていくという感じであるが、 アニメーションのフレームが変わったあとに、 前のフレームで書かれたものが残っているため、 それを消してから、次のフレームを書かなければいけない。
と、慣れればそれほどではないのかもしれないが、 とっつきにくい感じがあった。
一方、ArtistAnimationはあらかじめフレームのリストを生成して、 そのリストにフレームごとの描画結果をappendしていく感じ。
私は、こちらの方が好みだったので、今回はこちらの関数を使っている。
FuncAnimationの使い所としては、 今回のプログラムのax.axis([0, 5, -1.1, 1.1])のところに関連して、 今回はこのコードによって、すべてのフレームの描画される範囲を一括で指定しているが、 この範囲がフレームによって変わる場合はFuncAnimationのほうがよいのかもしれない。

plotの返り値について

また、よくつまづくポイントとして、 今回のコード中のtmp_frame, = ax.plot(x, y, "b-")のところ。 意味としては、あるフレームの描画を変数に格納しているだけだが、 ax.plot() の返り値は要素数1個のリストで返ってくるので、 それをunpackしてから、フレームリストに追加しなければいけない。

ちょっとしたTips

最後に、Tipsを。
私は Emacs + quickrun.el という組み合わせで開発をしているが、 plt.show() をしてから、開かれたウィンドウを閉じるのが手間に感じる。
そこで、この手間を多少解消してくれるTipsが stack overflow にあったので紹介する。
stack overflow
今回のコードでは、

plt.show(block=False)
input("Enter to close")
plt.close()

の所。
ただし、アニメーションのときは、フレームリスト一周分が回らないと、
input()の待機状態に入らなかったので、 最初の段階ではフレーム数を少なめにしてグラフの体裁を調整、 最後の段階でフレーム数を増やす、というやり方がよさそう。
また、書いていて思ったが、Jupyter Notebook とかなら、 アニメーションもあのコンソールの中で表示してくれそう。
後日、試してみます♪