jupyter notebook + matplotlibで非同期アニメーション #Python

以下、 jupyter notebook を notebook と表記します。

TL;DR

notebook で matplotlib の描画coroutineを非同期に回してアニメーションする。
（下のgifはfpsすごい出ている感がありますが、こんなんじゃないです…）

モチベーション

matplotlib のアニメーションの使い方を思い出すのが、面倒。
描画の動きを大雑把に確認したいだけなので、フレームレートは求めない。(私の環境ではfps10くらいでCPUリソースがかわいそうになるくらい)
- フレームレートを求めるなら、このブログ記事などをご参考に。
- ファイル等に書き出すなら、matplotlibのアニメーション機能を使うべき。
ブロッキングな処理(while True: sleep)をやると、 kernel の処理がブロックされてしまうので、おそらく描画されない。

環境

python
- 3.6.5
- 2.7.15 (「python2は小学生にもバカにされる」は割とマジ)
ライブラリ
- pip install jupyter numpy matplotlib

notebook は tornado

notebook のコアは非同期処理ウェブフレームワークの tornado
→ ここで tornado の ioloop が起動されている。

描画coroutineをioloopに処理させる

add_callback版

※以下、 notebook の .ipynb です。

%matplotlib notebook
import matplotlib.pyplot as plt
import numpy as np

from tornado.ioloop import IOLoop
from tornado import gen

@gen.coroutine
def coro_anim():
    fig, ax = plt.subplots()
    x = np.arange(0, 2 * np.pi, 0.01)
    f = lambda x: np.sin(3 * x)
    line, = ax.plot(x, f(x))

    while True:
        x += 0.1
        line.set_data(x, f(x))
        ax.set_xlim((min(x), max(x)))
        plt.show()
        yield gen.sleep(1.0)

IOLoop.current().add_callback(coro_anim)

一見、良さそうに見えますが、何回かセルを走らせていると、coroutineが複数走って、CPUリソースを食っていきます。
再実行時に、 Kernel Restart しても良いと思いますが、ここでは少し粘って、他に使えそうなのを探ります。

PeriodicCallback版

とは言っても、ioloopに入れたcoroutineのkillの仕方が分からなかったので、
初期描画処理をcoroutineの外に出して、再描画処理だけをcoroutineのまま残して、
PeriodicCallback で interval=1000ms で回しました。

以下も、 .ipynb です。

%matplotlib notebook
import matplotlib.pyplot as plt
import numpy as np

from tornado.ioloop import PeriodicCallback
from tornado import gen
cb = None

fig, ax = plt.subplots()
x = np.arange(0, 2 * np.pi, 0.01)
f = lambda x: np.sin(3 * x)
line, = ax.plot(x, f(x))

@gen.coroutine
def coro_anim():
    global x
    x += 0.1
    line.set_data(x, f(x))
    ax.set_xlim((min(x), max(x)))
    plt.show()

if cb and cb.is_running():
    cb.stop()
cb = PeriodicCallback(coro_anim, 1000)
cb.start()

これでセルを何回走らせたとしても、coroutineを stop してからインスタンスし直しているので、複数のcoroutineが走ってCPUリソースを食いつぶすことはなさそうです。
(実のところ、 Run Cells を魂を込めて連打すると、coroutineが複数できます。)
最悪、coroutineが複数できてしまっても、Kernel Restart で、その kernel で回っているioloopを初期化できるので、問題はなさそうです。

coroutineが複数走っているかの判断

notebook のノートの右上の方に python バージョンが出ているのですが、
その横に ○ があります。 coroutine が処理されているときはここが ● になるので、
これがcoroutineのインターバルより頻繁に点滅していれば、coroutineが複数走っている可能性があります。

描画の焼付きみたいなもの

描画させながらいじってると、下図のような焼付きみたいなものがたまに出ますが、グラフの右下のところをドラッグしてグラフサイズを変更すると、
描画しなおしてくれるので、coroutineを再起動しなくても直せます。

asyncioでの動作

私は標準ライブラリの asycnio の方が馴染みがあり、先に asyncio で書いて動作確認しているので、
asyncio でも書けば普通に動きます。もちろん、 python のバージョンに対してのコード依存性が出てきます。
今回は、 event_loop のことを書き始めると、本題からズレていきそうなので、tornado ということもあり、こちらにしました。