【python】matplotlibで3次元データを描画し、回転アニメーションにする

　3次元くらいのデータを描画したいときがある。簡単に散布図にできると便利。

データの用意

　sklearnのload_irisなどで取得できるデータセットを入力にする前提の次のような関数を作った。

from sklearn.decomposition import PCA

def gen_3d_data(dataset):
    pca = PCA(n_components=3)
    return pca.fit_transform(dataset.data), dataset.target

　あとはirisなり何なりを入れてやる。

スポンサーリンク

3次元プロット

　とりあえずプロットしたいだけならこれだけ。

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d

def matplotlib_plt(X, y, filename):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")

    ax.scatter(X[:,0], X[:,1], X[:,2], c=y/len(set(y)))
    plt.savefig(filename)
    plt.show()

　点の色の指定がちょっとセコいが、まあ良いこととする。axes3dはパっと見使ってないように見えるが、importしないとエラーになるので必要と思われる。

　呼び出し元のmainも書く。

from sklearn.datasets import load_iris

def main():
    iris_X, iris_y = gen_3d_data(load_iris())
    matplotlib_plt(iris_X, iris_y, "iris.png")

　実行すると次のような画像が出力される。

irisの3次元描画

　あと、ぐりぐり回せるグラフのようなものが別ウィンドウで開く（plt.sow()に対応）。

回転させたアニメーションを表示

　このような例が公式で示されている。

for angle in range(0, 360):
    ax.view_init(30, angle)
    plt.draw()
    plt.pause(.001)

mplot3d example code: rotate_axes3d_demo.py — Matplotlib 2.0.2 documentation

　やると確かにぐるぐる回るアニメーションが表示される。こりゃあええわ、ということで、次はこれをgifアニメにすることを考える。

　matplotlibにもanimationというモジュールがあり、色々できるようだが使い方を覚えるのが大変そうだった。なので、「上のforループ内で一枚ずつ画像出力してffmpegで繋げば良いだろ」という手抜きの方針で行くことにする。

def matplotlib_rotate(X, y, dataname):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")

    ax.scatter(X[:,0], X[:,1], X[:,2], c=y/len(set(y)))

    for angle in range(0, 360):
        ax.view_init(30, angle)
        plt.savefig("figs/{0}_{1:03d}.jpg".format(dataname, angle))

　呼び方は、

matplotlib_rotate(iris_X, iris_y, "iris")

　こうするとfigs/以下に画像が360枚吐かれるので、ffmpegでつなぐ。

$  ffmpeg -r 10 -i figs/iris_%03d.jpg -pix_fmt rgb24 -f gif out.gif

　とりあえずこれで行けた。画質が悪い割に容量が重いので、どこか上手くない指定になってるのかもしれないけど。

回るiris

　上出来ではないだろうか。