二項ツリーモデルをmatplotlibでアニメーション化する

matplotlibのanimationを用いた記事が意外に少なかったので，自分で作ったものを共有したいと思って書きました．

二項ツリー

二項ツリーモデルは上昇確率$p$と下降確率$1-p$の二項分布を用いた不確実性を記述するモデルである．今回は，二項分布から得られた値をモンテカルロ・シミュレーションにより発生させ，そのパスで記述していく．

matplotlibでアニメーション

参考文献: matplotlibでアニメーションを作る
https://qiita.com/yubais/items/c95ba9ff1b23dd33fde2

ArtistAnimation

あらかじめグラフの配列を用意することで，intervalごとに1つずつデータを取り出し，アニメーション化してくれる．

FuncAnimation

データを動的に変化させる必要がある．データが大きい場合にはデータを全て保持しないため，メモリ効率が良い．

プログラムコード

インポート

import math
from numpy import *
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import pylab

描画の設定

特に%matplotlib nbaggを加えることで，Jupyter上でインタラクティブな画像表示が可能になります．

N = 100   # 描画個数

%matplotlib nbagg
fig = plt.figure()
ax = fig.add_subplot(111)
ax.set_xlim(0, 100)    # x軸の範囲
ax.set_ylim(-30, 30)   # y軸の範囲
ax.set_xlabel('x')     # x軸ラベル
ax.set_ylabel('y')     # y軸ラベル
ax.grid()   # gridを描画する

描画するリストを準備

line = []   # 描画するグラフを格納する
for j in range(N):
    line.append(ax.plot([], [])[0])

x_list = []
y_list = []

データの生成

$x, y$の逐次返す関数を指定します．
FuncAnimationの第3引数に指定します．

def gen():
    y = [0 for i in range(N)]
    for x in np.linspace(0, 100, 100): # [start, end, 間隔]
        for i in range(N):
            y[i] += np.random.choice([-1,1]) 
        yield x, y

コールバック関数

FuncAnimationの第2引数に指定する関数です．
gen関数から受け取ったデータを描画するline配列に格納します．

def func(data):   # コールバック関数
    x, y = data
    x_list.append(x)

    for i in range(N):
        y_list.append(y[i])
        line[i].set_data(x_list, [y_list[j] for j in range(i, len(y_list), N)])

アニメーション関数を定義

ani = animation.FuncAnimation(fig,
        func, gen, blit=False, interval=100, repeat=False)

最終的なプログラム

import math
from numpy import *
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import pylab

N = 30   # 描画工数

%matplotlib nbagg
fig = plt.figure()
ax = fig.add_subplot(111)
ax.set_ylim(-30, 30)
ax.set_xlim(0, 100)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.grid()

line = []   # 描画するグラフ
for j in range(N):
    line.append(ax.plot([], [])[0])

x_list = []
y_list = []

def gen():
    y = [0 for i in range(N)]
    for x in np.linspace(0, 100, 100): # [start, end, 間隔]
        for i in range(N):
            y[i] += np.random.choice([-1,1])   # -1か1を選択
        yield x, y

def func(data):   # コールバック関数
    x, y = data
    x_list.append(x)

    for i in range(N):
        y_list.append(y[i])
        line[i].set_data(x_list, [y_list[j] for j in range(i, len(y_list), N)])

ani = animation.FuncAnimation(fig,
        func, gen, blit=False, interval=100, repeat=False)

二項ツリーのシミュレーションをすることができました．

参考文献

matplotlibでアニメーションを作る
https://qiita.com/yubais/items/c95ba9ff1b23dd33fde2
Jupyter上でmatplotlibのアニメーションを再生する
https://qiita.com/Tatejimaru137/items/6083e2e3a4e618da6274