matplotlibでランダムウォークをアニメーション

Python

matplotlibのanimation.FuncAnimationを利用すると（それなりに）簡単にアニメーションを作ることができる。

animation ― Matplotlib v1.1.0 documentation

使い方は次のようになっている。

class matplotlib.animation.FuncAnimation(fig, func, frames=None, init_func=None, fargs=None, save_count=None, **kwargs)

Bases: matplotlib.animation.TimedAnimation
Makes an animation by repeatedly calling a function func, passing in (optional) arguments in fargs.
frames can be a generator, an iterable, or a number of frames.
init_func is a function used to draw a clear frame. If not given, the results of drawing from the first item in the frames sequence will be used.

第一引数に描画するfigure, 第二引数にframe毎に呼び出すアニメーション関数を渡す。framesにジェネレータを渡すとyieldした値をアニメーション関数に引数として渡す。
これを利用してランダムウォークをアニメーションさせてみた。Youtubeの動画アップロードを試してみたかったんだけど、それにしても地味。

ソースは、

animation example code: animate_decay.py ― Matplotlib v1.1.0 documentation

をちょろっと書き換えただけなので不自然なところがあるかも。アニメーション関数の中で利用する変数がグローバルになってると内蔵が飛び出してるみたいだったのでクロージャの中に収めてみたけど余計なこだわりっぽい。

from matplotlib import pyplot as plt
from matplotlib import animation as animation
import numpy as np
import random

def gen():
    x, y = 0, 0
    cnt = 0
    while cnt < 101:
        yield x, y
        x += 1
        if random.random() < 0.5:
            y += 1
        else:
            y += -1
        cnt += 1

fig = plt.figure()
def func():
    ax = fig.add_subplot(111)
    line, = ax.plot([], [], lw=2)
    ax.grid()
    xdata, ydata = [], []
    ax.set_ylim(-1, 1)
    ax.set_xlim(0, 25)
    def f(data):
        x, y = data
        xdata.append(x)
        ydata.append(y)
        # グラフがはみ出そうになったら上（下）限値をセットし直す
        xmin, xmax = ax.get_xlim()
        ymin, ymax = ax.get_ylim()
        if x >= xmax:
            ax.set_xlim(xmin, 2*xmax)
            ax.figure.canvas.draw()
        if y >= ymax:
            ax.set_ylim(ymin, 2*ymax)
            ax.figure.canvas.draw()
        if ymin >= y:
            ax.set_ylim(2*ymin, ymax)
            ax.figure.canvas.draw()
        line.set_data(xdata, ydata)
        return line,
    return f

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

#ani.save("random_walk.mp4")
plt.show()