matplotlibによるシミュレーションの可視化

matplotlibでシミュレーションを可視化する手法は複数あるが、ここではそれをまとめ、私家版ベストプラクティスを述べる。

手法1

plt.pause()を使う

example

以下では、シミュレーションモデルを表す変数modelは、適切に初期化されていると仮定する。

import matplotlib.pyplot as plt

plt.figure()
for step in range(step_max):
    # シミュレーションを1ステップ進める
    model.step()
    # figureに描画する
    model.draw()
    plt.pause(wait_ms / 1000)

Pros

• シミュレーションループを明示的に書けるので、その他の形式に比してコード全体を理解しやすい

Cons

• plt.pause()を呼ぶたびにmatplotlibのウインドウにフォーカスがあたってうるさい。ターミナルにCtrl-Cを入力しての強制終了がやりづらくなる。
• 動画生成には対応できない

手法2

matplotlib.animation.FuncAnimationとplt.show()を使う

example

from matplotlib.animation import FuncAnimation
import matplotlib.pyplot as plt

def update(step, _step_max):
    model.step()
    model.draw()

    if step >= _step_max:
        plt.close()

fanm = FuncAnimation(
    plt.figure(),
    update,
    fargs=(step_max-1,),
    interval=wait_ms,
    frames=step_max-1,
    repeat=False,
    )
plt.show()

Pros

• matplotlibウインドウへのフォーカスは一度だけなのでうるさくない。
• plt.show()をfanm.save("movie.mp4", writer="ffmpeg")などに変えれば、動画生成にも対応できる

Cons

• シミュレーション終了後にウインドウを閉じるためのif文をupdate()の中に書かなくてはならず、見た目がうるさい。

手法3（オススメ）

matplotlib.animation.FuncAnimationを継承した自作クラス（FuncAnimationWithEndFunc）とplt.show()を使う。 自作クラスでは、シミュレーションが終了したときのcallback関数end_funcを設定できるようにしている。

example

from matplotlib.animation import Animation, FuncAnimation
import matplotlib.pyplot as plt

class FuncAnimationWithEndFunc(FuncAnimation):
    def __init__(
        self,
        fig,
        func,
        frames=None,
        init_func=None,
        fargs=None,
        save_count=None,
        *,
        end_func,
        cache_frame_data=True,
        **kwargs,
    ):
        super().__init__(
            fig,
            func,
            frames,
            init_func,
            fargs,
            save_count,
            cache_frame_data=cache_frame_data,
            **kwargs,
        )
        self._end_func = end_func

    def _step(self, *args):
        still_going = Animation._step(self, *args)
        if not still_going:
            # If self._end_func includes plt.close, returning False raises an exception
            # Belows are workaround for this
            self.event_source.remove_callback(self._step)
            self._end_func()
        return True

def update(step):
    model.step()
    model.draw()

fanm = FuncAnimationWithEndFunc(
    plt.figure(),
    update,
    interval=wait_ms,
    frames=step_max-1,
    end_func=plt.close,
    )
plt.show()

Pros

• update()を簡潔に記述できる
• matplotlibウインドウへのフォーカスは一度だけなのでうるさくない
• plt.show()をfanm.save("movie.mp4", writer="ffmpeg")などに変えれば、動画生成にも対応できる

Cons

• matplotlibの内部仕様が変わった際には動作しなくなる可能性がある
• コードの行数が長くなる。ただし、これはFuncAnimationWithEndFuncを別ファイルにすれば解決する。

手法3を使ったシミュレーション（鳥の運動モデルのシミュレーション）のコードを以下のリポジトリに置いた。

github.com

python ./src/run.pyで実行できる。