Python + matplotlib の二次元ヒストグラムでレゴブロックっぽい3D図を作る

ROOTのTH2DやTH2Fなどのクラスでは、2次元ヒストグラムを、legoっぽく図示できる。

ご存知の通り、Python + matplotlib の hist2dには、このような描画をするコマンドはない(はずである)。

というわけで、作ってみた。3次元バーを描画する関数であるbar3dを使った。これは、直方体を描画できる関数で、2つの対角線の座標を指定すれば良い。

ax.bar3d(x=xpos, y=ypos, z=zpos, dx=dx, dy=dy, dz=dz, shade=True)

コード全体を表示する

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(0)
x0 = np.random.randn(10000)*0.5 + 4 #標準偏差0.5, 平均4の正規分布
y0 = np.random.randn(10000)*2   + 3 #標準偏差2.0, 平均3の正規分布
x1 = np.random.randn(1000)*0.5 + 0 #標準偏差0.5, 平均0の正規分布
y1 = np.random.randn(1000)*2   + 0 #標準偏差2.0, 平均0の正規分布

x=np.concatenate((x0,x1))
y=np.concatenate((y0,y1))

xrange=[-10,10]
yrange=[-10,10]
bins=[50,100]
N=(bins[0])*(bins[1])
wbins = [(xrange[1]-xrange[0])/bins[0],(yrange[1]-yrange[0])/bins[1]]

hist, edgesx, edgesy = np.histogram2d(x,y,bins=bins,range=[xrange,yrange])
hist = np.flipud(np.rot90(hist)) # histogram2d 出力の仕様由来のおまじない

xpos, ypos = np.meshgrid(edgesx[:-1], edgesy[:-1])
xpos = xpos.flatten()
ypos = ypos.flatten()
zpos = np.array([0]*N)

dx = np.array([wbins[0]]*N)
dy = np.array([wbins[1]]*N)
dz = hist.flatten()

fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(projection='3d')

ax.bar3d(x=xpos[dz>0], y=ypos[dz>0], z=zpos[dz>0], dx=dx[dz>0], dy=dy[dz>0], dz=dz[dz>0], shade=True)
# ax.bar3d(x=xpos, y=ypos, z=zpos, dx=dx, dy=dy, dz=dz, shade=True) # dzが0のbinも描画する場合

ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
ax.set_xlim(*xrange)
ax.set_ylim(*yrange)
plt.show()

color引数に色の配列を与えることができるので、色を付けてみた。zsort="max"というのがミソである。だが思ったような図にならなかった。

cm = plt.get_cmap('jet')
ax.bar3d(x=xpos[dz>0], y=ypos[dz>0], z=zpos[dz>0], dx=dx[dz>0], dy=dy[dz>0], dz=dz[dz>0], color=cm(dz[dz>0]/max(dz)),alpha=0.9, zsort="max")

ROOTのlego描画は、単色の直方体ではなく、高さで色を変化させているが、bar3d では色の変化には対応していないためである。

コード全体を表示する

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(0)
x0 = np.random.randn(10000)*0.5 + 4 #標準偏差0.5, 平均4の正規分布
y0 = np.random.randn(10000)*2   + 3 #標準偏差2.0, 平均3の正規分布
x1 = np.random.randn(1000)*0.5 + 0 #標準偏差0.5, 平均0の正規分布
y1 = np.random.randn(1000)*2   + 0 #標準偏差2.0, 平均0の正規分布

x=np.concatenate((x0,x1))
y=np.concatenate((y0,y1))

xrange=[-10,10]
yrange=[-10,10]
bins=[50,100]
N=(bins[0])*(bins[1])
wbins = [(xrange[1]-xrange[0])/bins[0],(yrange[1]-yrange[0])/bins[1]]

hist, edgesx, edgesy = np.histogram2d(x,y,bins=bins,range=[xrange,yrange])
hist = np.flipud(np.rot90(hist)) # histogram2d 出力の仕様由来のおまじない

xpos, ypos = np.meshgrid(edgesx[:-1], edgesy[:-1])
xpos = xpos.flatten()
ypos = ypos.flatten()
zpos = np.array([0]*N)

dx = np.array([wbins[0]]*N)
dy = np.array([wbins[1]]*N)
dz = hist.flatten()

fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(projection='3d')

cm = plt.get_cmap('jet')
ax.bar3d(x=xpos[dz>0], y=ypos[dz>0], z=zpos[dz>0], dx=dx[dz>0], dy=dy[dz>0], dz=dz[dz>0], color=cm(dz[dz>0]/max(dz)))

ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
ax.set_xlim(*xrange)
ax.set_ylim(*yrange)
plt.show()

そこで、高さごとにバーの色を変えるコードを書いた。最後に一気にbar3dするのがミソである。コードとしては美しくないが許して欲しい。

ついでに、カラーバーを表示するコードも加えた。

mappable = plt.cm.ScalarMappable(cmap=cm)
mappable.set_array(dz)
cbar = plt.colorbar(mappable, shrink=0.5, aspect=20) #shrinkは高さ、aspectは高さ/横幅
cbar.set_label('Z')

コード全体を表示する

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(0)
x0 = np.random.randn(10000)*0.5 + 4 #標準偏差0.5, 平均4の正規分布
y0 = np.random.randn(10000)*2   + 3 #標準偏差2.0, 平均3の正規分布
x1 = np.random.randn(1000)*0.5 + 0 #標準偏差0.5, 平均0の正規分布
y1 = np.random.randn(1000)*2   + 0 #標準偏差2.0, 平均0の正規分布

x=np.concatenate((x0,x1))
y=np.concatenate((y0,y1))

xrange=[-10,10]
yrange=[-10,10]
bins=[50,100]
N=(bins[0])*(bins[1])
wbins = [(xrange[1]-xrange[0])/bins[0],(yrange[1]-yrange[0])/bins[1]]

hist, edgesx, edgesy = np.histogram2d(x,y,bins=bins,range=[xrange,yrange])
hist = np.flipud(np.rot90(hist)) # histogram2d 出力の仕様由来のおまじない

xpos, ypos = np.meshgrid(edgesx[:-1], edgesy[:-1])
xpos = xpos.flatten()
ypos = ypos.flatten()
zpos = np.array([0]*N)

dx = np.array([wbins[0]]*N)
dy = np.array([wbins[1]]*N)
dz = hist.flatten()

fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(projection='3d')

cm = plt.cm.get_cmap('jet')
dz_max = max(dz)
zbins = 20
zwbin = dz_max/(zbins)

# 0以下
flag=dz<=0
xall = xpos[flag]
yall = ypos[flag]
zall = zpos[flag]
dxall = dx[flag]
dyall = dy[flag]
dzall = dz[flag]
call = [0]*sum(flag)

for zbin in range(zbins):
    z0 = zwbin*(zbin+0)
    z1 = zwbin*(zbin+1)
    
    #z1以下
    flag=(dz>z0)&(dz<=z1)
    xall = np.concatenate((xall,xpos[flag]))
    yall = np.concatenate((yall,ypos[flag]))
    zall = np.concatenate((zall,[z0]*sum(flag)))
    dxall = np.concatenate((dxall,dx[flag]))
    dyall = np.concatenate((dyall,dy[flag]))
    dzall = np.concatenate((dzall,dz[flag]-z0))
    call = np.concatenate((call,[z1/dz_max]*sum(flag)))

    #z1以上
    flag=dz>z1
    xall = np.concatenate((xall,xpos[flag]))
    yall = np.concatenate((yall,ypos[flag]))
    zall = np.concatenate((zall,[z0]*sum(flag)))
    dxall = np.concatenate((dxall,dx[flag]))
    dyall = np.concatenate((dyall,dy[flag]))
    dzall = np.concatenate((dzall,[zwbin]*sum(flag)))
    call = np.concatenate((call,[z1/dz_max]*sum(flag)))
    
ax.bar3d(x=xall, y=yall, z=zall, dx=dxall, dy=dyall, dz=dzall, color=cm(call),alpha=0.9)

ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
ax.set_xlim(*xrange)
ax.set_ylim(*yrange)

mappable = plt.cm.ScalarMappable(cmap=cm)
mappable.set_array(dz)
cbar = plt.colorbar(mappable, shrink=0.5, aspect=20) #shrinkは高さ、aspectは高さ/横幅
cbar.set_label('Z')

plt.savefig("temp.png",dpi=300)
plt.show()