はじめに

python で 2D ヒストグラムを生成する方法には、大きくは、imshow と pcolormesh を使う方法があるが、微妙な違いがあって、時々忘れてしまうのでまとめておきました。

サンプルプログラム

縦と横を間違えやすいので、ここでは x 軸方向の方が長いデータを擬似的に生成します。

頻度分布を出しても良いのですが、ここでは、z = f(x,y) の z の平均値を 2D ヒストグラムで生成するサンプルで imshow と pcolormesh の比較をしてみます。

# 必要なライブラリをインポート
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm

# xとyの範囲を定義
xmin, xmax = -10, 10
ymin, ymax = -5, 5

# 10000のランダムなxとyのデータポイントを生成
x, y = 20 * (np.random.rand(10000) - 0.5), 10 * (np.random.rand(10000) - 0.5)

# zのデータをxとyを基にして計算
z = 100 * np.random.normal(0.5*x + 0.5*y, 0.1, 10000)**2 + 1000

NBINS = 30
x_bins = np.linspace(xmin, xmax, NBINS)
y_bins = np.linspace(ymin, ymax, NBINS)

# xとyの範囲ごとのzの平均値を計算
z_avg = np.zeros((len(x_bins)-1, len(y_bins)-1))
for i in range(len(x_bins)-1):
    for j in range(len(y_bins)-1):
        mask = (x >= x_bins[i]) & (x < x_bins[i+1]) & (y >= y_bins[j]) & (y < y_bins[j+1])
        z_avg[i, j] = np.mean(z[mask])

# imshowを使ってヒートマップをプロット
plt.tight_layout()
plt.imshow(z_avg, cmap='viridis', origin='lower', 
           extent=[x_bins[0], x_bins[-1], y_bins[0], y_bins[-1]], norm=LogNorm())
plt.colorbar(shrink=0.4, label="a.u.")
plt.xlabel('x')
plt.ylabel('y')
plt.title('z Average Heatmap using imshow')
plt.savefig("test_imshow.png")
plt.show()

# pcolormeshを使ってヒートマップをプロット
plt.tight_layout()
mesh = plt.pcolormesh(x_bins, y_bins, z_avg.T, shading='auto', cmap='viridis', norm=LogNorm())
plt.colorbar(mesh, shrink=0.4, label="a.u.")
plt.xlabel('x')
plt.ylabel('y')
plt.title('z Average Heatmap using pcolormesh')
plt.savefig("test_pcolormesh.png")
plt.show()

結果

生成された図を見てみましょう。こちらが imshow で生成した場合です。

これが pcolormesh を用いた場合です。

中身はどちらも同じですが、imshow の方は x : y が 2:1 の要素数であることを反映して横長になるのに対して、pcolormesh は画像サイズに従って調整していますね。

pcolormesh の使用時の注意点:

pcolormeshは、行と列の境界の配列を引数として取ります。これにはx_binsとy_binsが使用されます。
pcolormeshは、Z値（ここではz_avg）を行ごとに読み取るため、z_avgを転置する必要があります。これは、z_avg.T で行います。

mesh = plt.pcolormesh(x_bins, y_bins, z_avg.T, shading='auto', cmap='viridis', norm=LogNorm())

この部分で、z_avg.T の.T を忘れないようにしましょう。

その他

from matplotlib.colors import LogNorm

を用いて、 norm=LogNorm() というオプションを用いて z軸方向を log スケールにするのは pcolormesh も imshow も同じです。

関連ページ

secondary_axis の利用については下記の記事を参照ください。