僕のあまり妥協せずに散らかさないMatplotlibの考察

Matplotlibが辛い

僕はmatplotlibで沢山の図を描きます。とても辛いです。いつもググってますが、人によって使い方が違います。沢山描くと一々コードが散らかります。補完が効かないこともあります。
そこで、どうして僕のコードがゴチャつくのか考えました。

この記事について

これから散らからないように、色々考えて僕のmatplotlibの使い方を定めました。この方法が良いかどうかは知りません。今のところ良いかんじだけれど、使いこんではいないです。
要約するとwrapして継承とかダックタイピングをユーザーに強制します。こんな感じのコードになります。ninmplはこのために書いた自分用パッケージ。コメント除けば50行ほどの、本当に自分用のもの。

from ninmpl import FigureBase, PlotBase

class PlotAngle(PlotBase):
    """角度のプロット"""
    options = dict(projection='polar')

    def adjust(self):
        self.ax.set_title('Circle')

    def plot(self, x, y):
        self.ax.plot(x, y)
        
class PlotLine(PlotBase):
    """普通のプロット"""
    def adjust(self):
        self.ax.set_title('Line')

    def plot(self, x):
        self.ax.plot(x)

class NinFigure(FigureBase):
    """Figureの設定とか"""
    def adjust(self):
        self.fig.set_figwidth(5)
        self.fig.set_figheight(5)
        self.gridspec = plt.GridSpec(2, 4, hspace=0.4)

r = np.arange(0, 2, 0.01)
theta = 2 * np.pi * r
line = [[1, 2], [2, 1]]
line2 = [[1, 1], [3, 4]]

with NinFigure() as nf:
    nf[0, 0] = PlotAngle(r, theta)
    nf[1, 0] = PlotLine(line)
    nf[0, 1:3] = PlotLine(line2)

この記事はソフトの宣伝じゃなく、自分の考えの垂れ流しです。というか、このくらいのソフトは自分で書けばよいのです。

方針1. OOPする

matplotlibはすでにOOPしているように思えるかもしれません。しかし、それはmatplotlib自身であって、ユーザーがOOPしないと散らかるのだろうと思います。
matplotlibのfigure関数が返すのはaxesのインスタンスです。つまり…ユーザーは「折れ線グラフクラス」とか「箱髭図クラス」とかを使って継承していくことが出来ないわけです。ユーザーは実質手続き型をしているように思えてきました。(僕が学び足りないだけかもしれないけど)
さらに「指数のクラス」とか「枠を表示しないクラス」とか、色々自分好みの設定を組み合せたかった…。
ユーザーレベルでmatplotlibをOOPをしたい。なんなら強制したい。強制しないと使い捨てのものを大量に書いてしまうのが無能系データサイエンティスト(つまり僕)だと思うからです。これが散らからない方法として考えた骨子です。

方針2. Gridspecは使う

柔軟にplotしたいです。Gridspecを使うことで非常に柔軟性が高まりますが、反面煩雑です。matplotlibの公式は「普通の人はGridSpecまで必要になることは少ないだろ」と言っています。
Gridspecをいれるとゴチャゴチャ度も高まりますしね。だけど、これ入れるだけでものすごく柔軟度が高まるので僕は妥協できませんでした。妥協しないとなると散らかるのですが、OOPでなんとかなってほしかった。

方針3. axesをwrapする

subplots_mosaicを使わずにmatplotlibを素直に使うとコードがすごく散らかります。例えばこうなるのです。辛い。

from matplotlib import pyplot as plt
import numpy as np

# ここから書かないといけない所
fig = plt.figure()
ax1 = fig.add_subplot(311)
ax1.set_ylim(2, 5)
ax1.plot(np.arange(10))
ax2 = fig.add_subplot(312)
ax2.set_ylim(2, 5)
ax2.plot(np.arange(10))
ax3 = fig.add_subplot(313)
ax3.set_ylim(2, 5)
ax3.plot(np.arange(10))

ゲシュタルト崩壊します。for文である程度なんとかなりそうにも思えるけど、どちらにしろ散らかります。subplot_mosaicは下記のようにできる神がかった関数です。

import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplot_mosaic([['hoge', 'fuga'],
                              ['piyo', 'foo']])
ax['hoge'].plot(np.arange(10))

二次元リストを元にaxesの辞書を作ってくれるのです。素敵だ…。では、何故これを使わないのでしょうか？それはこれが直でインスタンスを返すからです。インスタンスを返すのが辛いのは方針1の通りです。

では、どうするのがいいのでしょうか？
Axesクラスを直で継承するのは無しと思います。あれはfigure関数やsubplot_mosaic関数で生成されるけれど、そのときにGridSpecを組みあわせないといけないです。そんな仕様では直接継承するのはダルすぎます。
Wrapするしかあるまい…そう思いました。

方針4. axesを引数としてとるUIに対応する

Axesを引数に取ってAxesを返す関数やクラスを書くという方法が公式や有名プロジェクトで使われています。

そうすれば、Axesの設定と描画を分けることが出来るのです。そういう関数は、例えばこんな感じに書けるのだと思います。

from matplotlib import pyplot as plt
from hoge import hoge_plot
import numpy as np

# ここから書かないといけない所
def adjust(ax):
    ax.set_ylim(0, 100)
    return ax

def plot(data, cmap='rainbow', ax=ax):
    ax.plot(data)

ax = {}
fig = plt.figure()
ax[0] = fig.add_subplot(311)
ax[1] = fig.add_subplot(312)
ax[2] = fig.add_subplot(313)

hoge_plot(np.arange(10), adjust(ax[0]))
hoge_plot(np.arange(20), adjust(ax[1]))
hoge_plot(np.arange(30), adjust(ax[2]))

こんな感じにAxes自体を使いまわしていきたいですね。Wrapperの中でそんな風に出来るようにしましょう。

方針5. Context manager使う

Context managerというのは、要するにwith文ですね。こういうやつ。

with open('hoge.txt', 'r') as fp:
    data = fp.read()

これは必ずやらないといけない作業をカプセル化することが出来るありがたいものです。
では、matplotlibの場合に必ずするのは何でしょうか？最後に保存したり、表示したりするのをわすれがちです。それより嫌なのは、matplotlibで大量のplotをするときにメモリ食いまくるのに、それを解放しわすれることです。上記の事は一つのブロック内で完結してほしいです。だから、最初にwithで色々設定いたいです。

方針6. Context manager自体の継承

Axesだけじゃなく、figも継承でやっていきたいですよね。上記のcontext managerの中でfigをwrapすればいいですね。これをすると、継承によって図の仕様を使いまわす事が出来るかと思います。

というわけで

冒頭のコードにおちつきました。この程度の規模ですと最終的なコードの量は増えるのでしょうが、大規模になるとコピペが減ったり…したらいいな。