まとめると

二次元データを画像化して保存したいとき、matplotlib.pyplot.imshow() は遅いので、matplotlibのColormapを使って直接保存できます。

問題設定

ある二次元データが（たくさん）あって、画像化したいとします。

import numpy as np
from matplotlib import pyplot as plt
data = np.arange(100).reshape(10,10)
plt.imshow(data)
plt.savefig("test.png")

しかしご存知の通りmatplotlibは遅いですね。遅いのはいちいちグラフィックレンダラーを呼んでレンダリングしてその結果を保存しているからです。
しかしやりたいことは、dataをなんらかの変換でRGBに押し込んで保存するだけです。

import PIL.Image
def quantize(arr):
    arr = np.floor((arr - arr.min())  * 255 / (arr.max() - arr.min())).astype(np.uint8)

img = PIL.Image.fromarray(quantize(data))
img = img.resize([400,400], resample=PIL.Image.NEAREST) #拡大しているだけ
img.save("test.png")

img = np.zeros(data.shape+(3,), dtype=np.uint8) # キャンバスの用意
img[:,:,0] = 255 # 適当に1
img[:,:,1] =  255 - quantize(data) # 適当に2
img[:,:,2] = quantize(data) # 適当に3

img = PIL.Image.fromarray(img)
img = img.resize([400,400], resample=PIL.Image.NEAREST) #拡大しているだけ
img.save("test.png")

Colormapの使い方

ご存知の通り(?)matplotlibにはカラーマップという概念があります。
値から色(RGB値)への写像です。plt.imshow(data, cmap=hoge)などと指定することができますよね。

カラーマップのバリエーションについてはblog:matplotlibのcmap(colormap)パラメータの一覧。など参照。

matplotlib内で、カラーマップはmatplotlib.colors.Colormapオブジェクトとして表現されています¹。

Colormapオブジェクトは、2d-arrayを受け取ってRGBA値の入った3d (2d*4ch) arrayを返す関数（に毛が生えたもの）として扱うことができます。
matplotlib.pyplot.get_cmap()にカラーマップの名前を指定することでColormapが返ってきます。つまり、

cmap = matplotlib.pyplot.get_cmap(NAME_OF_COLORMAP)
image = Colormap(raw_array)

という感じで、画像化が出来上がりです。
この画像をレンダラに表示させるのがplt.imshow()なので、画像が欲しいだけなら

PIL.Image.fromarray(image).save("filename.png")

の方がずっと早いということになりますね。例を示します。

import matplotlib.pyplot as plt

def minmax(arr):
    arr = (arr - arr.min()) / (arr.max() - arr.min())
    return arr

cmap = plt.get_cmap("viridis")
data = minmax(data)
img = cmap(data, bytes=True)
img = PIL.Image.fromarray(img)
img = img.resize([220,220], resample=PIL.Image.NEAREST) #拡大しているだけ
img.save("test.png")

• 注意点1.

  • viridisというのはmatplotlibのデフォルトのカラーマップの名前です。

• 注意点2.

  • 入力としては[0.,1.](float)か[0, 255](uint)の範囲のデータが期待されているので、先にminmax正規化によって値の範囲を調整しています。

• 注意点3.

  • bytes=Trueを指定することで[0, 255]のデータが返ってきます。デフォルトでは[0.,1.]のデータが返ってきます。PIL.ImageではRGBAはでしか[0, 255]扱えないのでちょっと注意します。

以上で、plt.imshow()と同等のことがずいぶん速く実現できます。

参考: 自作colormapの作り方や既存カラーマップの中身を分析するやり方について以下が詳しいです。

• matplotlibのcolormapのRGB情報取得と関連操作