matplotlibでアニメーション

ある画像データをパラメータの変化に合わせてアニメーションをさせながら表示させたい。 そのような場合のメモをしておく。

pcolormesh()を用いて画像データをプロットし、 その下にplot()を用いたパラメータの変化を示す。 さらに、pcolormesh()で表示されている場所をパラメータの変化のグラフ上に線で示す。 pcolormeth()に対するカラーバーの位置も制御したい。 レイアウトの自由度を持たせるために、add_axes()メソッドを用いている。

アニメーションは、

matplotlib.animation.ArtistAnimation
matplotlib.animation.FuncAnimation
のどちらかを用いて行える。

ArtistAnimationでは、各フレームのArtistオブジェクトを保存して、 Animationを作成する。 フレーム数が多くなる場合も考えると、FuncAnimationを利用したい。 ArtistAnimation, FuncAnimationとも、blitキーワードをTrueに設定すると、 指定した場所のみ書き換えるようになるため速く描画できると思われるが、 グラフ（Axes）の外は書き換えが行われなかったりと制約もできてしまう。 pcolormesh()の書き換えでは、軸の枠線がプロットに上書きされて消えてしまうようなことも起こる。 この場合は、blit=Falseを指定することをお勧めする。

MakeData()で、(nx, ny, nz)のshapeをもつデータを作成し、 各zに対する幅nx、高さnyのデータとみて、プロットをしてみる。

FuncAnimation()には、各フレームごとに呼び出される関数を渡す必要がある。 その関数には、フレーム番号を受け取る一つの引数が必要となる。 実際のupdate()関数では、関数内部で用いるfigure、 data、az とフレーム番号を受け取るように書きたい。 そうすると、引数が多くなってしまうので、lambdaを用いて、引数をフレーム番号のみに減らし、 各フレームごとに呼び出される関数としている。

最後のほうにある、

anim.save('anim.gif', writer='pillow')

のコメントを外せば、アニメーションGIFファイルとして保存できる。 writer='pillow'を指定するためにpillow(PIL)をインストールしておく必要がある。

モジュールのバージョン
python : 3.8.3
matplotlib : 3.2.2
numpy : 1.19.0
PIL(pillow): 7.1.2

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

def MakeData(nx, ny, nz):
    data = np.ones((nx, ny, nz))
    ax = np.linspace(-3, 3, nx)
    ay = np.linspace(-3, 3, ny)
    az = np.linspace(0, np.pi, nz)
    for k, z in enumerate(az):
        for j, y in enumerate(ay): 
            for i, x in enumerate(ax):
                data[i, j, k] = np.sin(x+z) * np.cos(y+z)
    return data, az

def update(fig, data, az, iz):
    fig.clear()
    # create axes
    ax0 = fig.add_axes([0.1, 0.4, 0.7, 0.5])
    cbar_ax0 = fig.add_axes([0.82, 0.4, 0.05, 0.5])
    ax1 = fig.add_axes([0.1, 0.1, 0.7, 0.2])
    
    # plot data
    im0 = ax0.pcolormesh(data[:, :, iz].T)
    fig.colorbar(im0, cax=cbar_ax0)
    ax1.plot(az, '-', color='#0000FF')
    ax1.plot([iz, iz], [0, np.pi], '-', color='#FF0000')

    # put text labels
    label = ax1.text(0.1, 0.8, '%.5g (rad)' % az[iz], 
            transform=ax1.transAxes)
    ax0.set_title('Frame %d' % iz)


nx, ny, nz = 30, 60, 100
data, az = MakeData(nx, ny, nz)

fig = plt.figure(figsize=(5, 6))

anim = FuncAnimation(fig, lambda iz: update(fig, data, az, iz), nz, 
        interval=100, blit=False, repeat=False)

#anim.save('anim.gif', writer='pillow')

plt.show()