渡米生活。(日記)

渡米生活。本家から切り離しました。あまり渡米生活に関係のないプログラムネタや音楽ネタなど。

ROOT から matplotlib(pylab) + numpy へ(1次元&2次元ヒストグラム)

このエントリーは、ROOT一辺倒だった人間が時代の波に流されてしょーがないのでmatplotlib+numpyで図を書くのに「これどうすんの?!」と困った項目の備忘録です。何分初心者なので、ダサダサかつ間違ってるかも。

お品書き

  1. まとめ
  2. TH1D
    1. numpy.array () or list ?
  3. TCanvas
    1. Colorbar
    2. canvas1->cd()?
  4. pylab.hist? numpy.histogram?
  5. bin centerをとってくる
  6. Setlogy(), SetRangeUser()
  7. lineの幅などのデフォルト設定
  8. Legend
  9. Draw, gPad->Update()
  10. TH2D

まとめ

高エネ・宇宙線業界でフツーに見る1次元ヒストグラムをpythonで書こう、というのが目標。
勿論、例題はいろいろ落ちているのだけど、どうも痒いところに手がとどかないというか……
若い人達がpylabで似たような絵を量産しているので、簡単にできるのだと思っていたら、dashiとかいう自前ライブラリーを使って書いていた。(ちなみに同名の公開モジュールがあるけど、そっちは別物)
大変便利そうだけど、きちんと公開されていない誰かのツールを使って絵を書くのは怖すぎるので断念。

で、そのdashiの中をみてみたのだけど、これが結構膨大なライブラリ群で、なんだ、ROOTと同じことするのにこんなに色々準備が必要なのか、と、すっかりpythonに完全移行する気は萎えてしまった(笑)。
というわけで、histogramを書いてprojectionとかガンガンやりたいけど、cintは使いたくない、というなら、大人しくPyROOTを使った方がいいんじゃないかという気がする。どうしてもmatplotlibっぽい絵にしたければ、rootplotあたり使うとか?(使ったことないけど)

しかし、まあ、せめてもっと自由に絵のお化粧くらいさせてくれ、と思ったので、とりあえずその方法のメモです。

(蛇足だが、matplotlibで描かれた絵がかなりの割合で軸が読みにくいのはなんとかならんか。プレゼンではほとんど見えない。セリフフォントは見た目綺麗だが、プレゼンにはむかないと思う。
若い人達、年寄りにもっと優しいフォント設定にしてくれ!
ROOTのデフォルト設定はHelbeticaだし、tickもラベルも割合大きめなので、手元で見ると無骨に見えちゃうんだけど、これは遠くからでも見やすいんだよ。)

TH1D

TH1Dでhist作って、Fillして…という流れは、matplotlibではどうなるか。

  1. fillするデータを配列もしくはnumpy.array()で準備。
  2. weightの値をデータと同じ次元の配列で用意。
  3. 最後にこれをpylab.histなどでプロット

といった手順になる。無骨にやると、for文でデータをひとつずつ読みながら、同時にweightも計算する、という感じになるのだけど、pythonだとこんな書き方もアリ。

(自分の練習プログラムをそのまま載せてるので、関係ない部分もいっぱいあるけどご容赦。最後の1行が該当部分。)

#!/usr/bin/python 

import numpy as np
import matplotlib as mpl

nevts1 = 10000

gammagen = 1
gammasig = 2
testflux = 1000

elog10low1 = 1
elog10high1 = 11
nbins1 = 20

#
# util function for weight calculation
#
def weightfunc(eintg, elog10, nevts=1) :
    ene = 10.0**elog10
    return eintg / (ene**(-gammagen)) * testflux * ene**(-gammasig) / nevts

#
# calculate integral of generation spectrum E^-(gammagen)
#
eintg1 = 0
if (gammagen == 1) :
    eintg1 = np.log((10**elog10high1)/(10**elog10low1))
else :
    eintg1 = ((10**elog10high1)**(1-gammagen) - (10**elog10low1)**(1-gammagen)) / (1-gammagen)

#
# prepare data with E^(-gammagen) 
#
data1 = np.random.uniform(elog10low1, elog10high1, nevts1)

#
# calculate weights for data1. 
# use np.array so that we may scale the weights later!
#
flux1 = np.array([weightfunc(eintg1, ene) for ene in data1])

data1がデータ(入力値)の配列で、flux1がその重み(weight)。
計算したいweightが沢山ある場合は、forで回した方がいいかも。

numpy.array () or list ?

ところで、forで回しながらdataをappendするのに、numpyの配列を使うと必ず失敗するのだけど、何か根本的な使い方が間違っているのだろうか。
numpyは行列演算用の配列だからあとから行を増やすとかフツーはやらないのだと知った。。

pythonの配列は行列要素のスカラー倍演算すら受け付けてくれないので(まあ、そういう用途じゃないから当たり前か)、なるべくnumpy.arrayを使うようにしているのだけど、forで回しながらデータを追加するときは今のところしょうがないのでlistを使っている。
こんな感じ。

data1 = np.random.uniform(elog10low1, elog10high1, nevts1)
flux1 = []
for i, ene in enumerate(data1) :
    flux1.append(weightfunc(eintg1, ene))

もっとも、追加じゃなくて、最初から配列のサイズが分かっている場合は、numpy.arrayを使って以下でもOK.
この方が、何度もメモリアロケートしないだけ早いかも。

data1 = np.random.uniform(elog10low1, elog10high1, nevts1)
flux1 = np.zeros(len(data1))
for i, ene in enumerate(data1) :
    flux1[i] = weightfunc(eintg1, ene)

TCanvas

実は、matplotlibで一番よくわからんのが、matplotlibとpylabの違いだったりする。
matplotlib.pyplotに対して呼ばれる関数と、pylabに対して呼ばれる関数は何か違うのか? そこの使い分けはどうなってるの??
なんかpylabはインタラクティブインターフェース、、とか書いてあるんだけど、ドキュメントを読まずに例題をひろって無理矢理絵を描いているレベルではなんのこっちゃい。。。

というわけで、もうそこを理解するのは諦めて、とりあえずTCanvas相当のものが書ければいいや!ということで、例題。

2014/8/16 訂正
どうやらpylab.figureを使うと複数のfigureを開いた時に期待通りの動作をしてくれないことが分かったので、Pylab の代わりにmatplotlib.pyplot.figureを使うようにコードを変更。

#import Pylab as P 
import matplotlib.pyplot as plt

# generate a canvas with an instance name "fig"
# fig = P.figure()
fig = plt.figure()

# canvas->divide()?
ax1 = fig.add_subplot(2,2,1)
ax2 = fig.add_subplot(2,2,2)
ax3 = fig.add_subplot(2,2,3)
ax4 = fig.add_subplot(2,2,4)

plt.figure()がTCanvasの生成で、TCanvas::Divide(nx, ny)がadd_subplot(nx, ny, index).
ちょっと違うのは、add_subplotの三番目の引数が分割された窓のindexを表していて、戻り値にその窓への参照が返って来る、というところ。
TCanvasだと、canvas->cd(index) みたいなことをする必要があったのだけど、これは各パッドに名前がついてるから、いきなりそこに絵をかけて便利。

ちなみに、1枚しか要らない、という場合は、そもそもfigure関数を呼ぶ必要もないけど、そうするとその場合はy軸をlogにしたりy軸の範囲を制限したりするのに別の関数を使わないといけない(add_subplot →subplot, set_ylim →ylim、という感じで、どうも対応に一定のルールはあるようだけど)…という感じで混乱するので、面倒だから1枚だけのsubplotを作る、でいいと思う。

ax1 = fig.add_subplot(1,1,1)

このへんの統一感のなさが、matplotlibがとっつきにくい理由の一つのような気がするんだよなー。
公式ページの例題みても、同じ結果を得るのに違うやり方が多過ぎるよ。。。

2014/8/16追記:

どうやら、事情がのみこめてきた。
matplotlib.pyplotモジュールというのはとにかく何でも屋さんで、ユーザーがオブジェクト指向を気にせずに使えるようになっているらしい。
つまり、例題でよくあるplt.Plot(hoge)を呼ぶと、figure(TCanvas相当)を作って、subplotを作って、というのを勝手にやってくれて、それからhogeをDrowする、ということ。
一方、Pylab.figureはその背後で動いている関数なので、もっとC++ライクに、ちゃんとオブジェクトを作ってオブジェクトに対して関数を呼ぶ必要がある。

で、どうもPythonのインターフェース(この場合はmatplotlibだけど)には、クラス関数やメンバーのアクセスインターフェースを作る時の関数名のNaming Conventionがある模様。
つまり、Get関数、Set関数なんかは、インターフェースではGetやSetがおちて、その後ろだけでアクセスできるようになる。
だから、、set_ylim が matplotlib.pyplotではただのylimになってしまう、、と理解すると、いちいち二種類の関数を覚える必要はなくなるわけです。
(C++にどっぷり浸かった人間としては、折角カプセル化してprivateにして隠したデータメンバにset/getメソッドつけて、ってやってるのに、要するにこれってprivateメンバをpublicにしちゃうのと同じ??、、という抵抗感はあるのだが。。)

Colorbar

これで一件落着、と思ったら、2次元プロットのcolorbarを描くのに困ったことが発生。plotだとかhistだとか、2次元ヒストグラムのpcolorだとかはaxisに対して呼ぶことができるのだけど、その絵の横につけるcolorbarはaxisに対して呼ぶことが出来ず、plt.colorbar(ヒストグラム)、みたいな文法になる。

引数にヒストグラムオブジェクトをとっているのだから、そのヒストグラムの横に描いてくれるのかと思ったら、どうもそうではないらしい。
実は、matplotlibには、現在のfigure、現在のaxis、という概念がある。
つまりrootのgPadとまったく同じ概念で、どの絵のcolorbarでも現在のgPad(じゃないけど)の上に描いてしまうのだ。

まあ、実はcax=axisオブジェクト、とか引数を足してやれば、希望のaxisに描いてくれる(というか、多分もともと何処にcolorbarを描くかは至極自由に決められる仕様になっているということか)。…のだけど、この手動でaxisを決めてやるのが、結構めんどい。

ここに至り、どうやらPylab.figureを直接使うのは得策ではなくて、matplotlib.pyplot.figureを使うべきなのだと判明。
変えるのは、importのラインと、P.figure() の代わりにplt.figure()を呼ぶだけ。
こうやってfigureオブジェクトを作ってやれば、あとは関数名の方はset_ylimとかC++ライクな関数名が使えます。

canvas1->cd()?

で、どうやって問題のgPad (…じゃないけど)を移動するのか、という話ですが、こうするらしい。

ax1 = plt.subplot(2,2,1)

え、新しくsubplot作っちゃうんじゃないの、とか気になりますが、どうもこの関数は、その位置のsubplotが存在していなければ新しく作り、既にあったらそれをgPad(…じゃないけど…ってしつこい)にセットするものらしい。

ところで、この場合はまだfigureが1個(つまりTCanvasが1個)しかないから良いのだけれど、勿論複数のCanvasを描きたいことはある。

で、plt.figure()を4回くらい呼んで、それぞれfig1~fig4の変数名をつけたとする。で、一度fig4でお絵描きしてから、fig1に戻って来てcolorbarをつける、なんて場合には上の一行ではダメで、こうする。

#最初のfigureに移動。indexは0じゃなくて1であるのに注意。
fig1 = plt.figure(1)
#fig1の1番目のaxisに移動
fig1ax1 = plt.subplot(2,2,1)

というわけで、rootなら canvas1->cd(1) の1行で済む作業が、matplotlibでは2行が必要になる、ということの模様。

ちなみに、indexが0からじゃなくて1から始まってるのは、rootのようにTCanvas自体が0番を指すから、じゃなくて、MATLABの伝統を受け継いでいるかららしい…。


pylab.hist? numpy.histogram?

さて、いよいよヒストグラムを書こう!
matplotlib,histと検索すると、pylab.histの例題が沢山ひっかかるのだけど、こいつがどうも使いにくい。
デフォルトでfillオプションになってヒストグラムを塗りつぶされちゃう。
オプションで histtype='step'をつければ、それっぽくはなるけど、今度は線の幅を指定できない。legendもなんか箱型になっちゃう。

というわけで、どうやら(高エネ・宇宙線業界で)フツーによく見るヒストグラムが書きたければ、numpy.histogramを使った方がよさそう。

まず、pylab.histの例題。
さっきdivideした窓の1番目(ax1)に書くことにする。

n, bins, patches = ax1.hist(data1, nbins1, weights=flux1/nevts1, histtype='step', color="magenta", label="pylab.hist")
ax1.legend()

weightsのところで、配列flux1の要素を一律nevts1で割る、というのをやっているけど、これは先ほども書いたように、flux1がnp.arrayだから可能。
histtypeは'step'を指定しないと、デフォルトではfillされる。
絵を描くのにこれ1行で済んでしまうので、簡単といえば簡単なんだけど、このstepオプションの時の線幅や線種の変更の仕方がどうしても分からない。
というか、そもそもAPIドキュメントをみると、そういうオプションがないんだよね……。

というわけで、例題2。
np.histogramを使ってヒストグラムを生成→pylabでお絵描き。

yval, binEdges = np.histogram(data1, nbins1, weights=flux1/nevts1)
ax2.step(binEdges[0:-1], yval, where='post', color="green", label="np step",linewidth=3)
ax2.legend()

さっきつくった2番目の窓に描いてみました。
ポイントは、binEdgesはデータ数nに対してn+1個の要素数があるので、最後のデータを捨てること。
これだと、とりあえず線の太さは変えられるので、pylab.histよりはちょっとまし。
しかし、線の種類(破線にするとか)はどうも変えられないっぽいです(linestyle="dashed"を指定しても無視される)。

というわけで、次はダサダサだけど、データを変形してpylab.plotで書いちゃえ!という例題。

#
# prepare data for plotting
#
def format (yval, binedges) :
    x = []
    y = []
    for index, value in enumerate(yval) :
        print "y = %f, binedge_low = %f, binedge_hi = %f" % (value, binedges[index],binedges[index+1])
        x.append(binedges[index])
        x.append(binedges[index+1])
        y.append(value)
        y.append(value)

    # append very small value to close histos
    x.insert(0, x[0])
    y.insert(0, 1e-30)
    x.append(x[-1])
    y.append(1e-30)

    return x, y

yval, binEdges = np.histogram(data1, nbins1, weights=flux1/nevts1)
x, y = format(yval, binEdges)
ax2.plot(x, y, '--', color="brown", label="numpy.histogram step",linewidth=1.5)
ax2.legend()

もう完全にお絵描きのためのdirty fixです。これに使った配列で統計的な計算は何一つできません。
(まあ、オリジナルのhistogramが残ってるからいいんですが)

bin centerをとってくる

ヒストグラムのお絵描きにはあまり関係ないが、np.histogramは横軸の値としてbin_lowとbin_highを返してきて、binの中央値は戻ってこないので、bin centerを計算するときはどうするか。
この1行で済みます。このへんは確かにpython使いやすい。

bincenters = 0.5*(binEdges[1:]+binEdges[:-1])

Setlogy(), SetRangeUser()

ROOTでいえば、TAxisで設定するような項目です。
Y軸をlogにする方法はなんだか色々あって、もうどれを使っていいのやら、という感じなのだけど、ROOT風にやるなら、それぞれのパッドで指定できるのが有り難いので、subplotの各インスタンスの関数を使う。

ax3.set_yscale('log')

で、このlogyをやると、ヒストグラムの両端の縦棒が消えてしまうわけです……。なんだかな。
というわけで、無理矢理両端も閉じたかったら、やっぱりpylab.plotを使うしかないっぽい。
上のナンチャッテformat関数では、両端にy=1e-30なんて数字を埋め込んでいるので、このままではレンジが広過ぎてみにくいから、Y軸の範囲を制限します。
TAxisならSetRangeUserです。

ax3.set_ylim([1e-9, 1e3])

lineの幅などのデフォルト設定

いちいちlinewidthを指定するのが面倒な人は、この1行を書きます。

plt.rcParams['lines.linewidth'] = 2

デフォルトの線幅では、大抵細すぎて遠くからでは見えない。
線幅の他にも、いろいろ設定できるはず(多分…)。

Legend

Legendの指定も、ここではsubplotに対して行います。
これも、P.legend()みたいに呼べるんだけど、そうするとどの絵につけるのかが指定できないので。

ax4.legend()

Draw, gPad->Update()

絵を実際に表示するには、以下の一行が必要です。

plt.show()

これをスクリプトの途中で挟むと、そこでスクリプトが止まるっぽいので、最後にこの一行をおいて一気に描く。

例題を全部絵にすると、こんな感じです。

スクリーンショット 2013-09-06 5.37.02 PM

この絵を書いたスクリプトは以下の通り。debug printがそこかしこにありますが煩かったらコメントアウトして下さい。

#!/usr/bin/python 

import numpy as np
import matplotlib.pyplot as plt

# change default linewidth
plt.rcParams['lines.linewidth'] = 2

nevts1 = 10000

gammagen = 1
gammasig = 2
testflux = 1000

elog10low1 = 1
elog10high1 = 11
nbins1 = 20

#
# util function for weight calculation
#
def weightfunc(eintg, elog10, nevts=1) :
    ene = 10.0**elog10
    return eintg / (ene**(-gammagen)) * testflux * ene**(-gammasig) / nevts

#
# calculate integral of generation spectrum E^-(gammagen)
#
eintg1 = 0
if (gammagen == 1) :
    eintg1 = np.log((10**elog10high1)/(10**elog10low1))
else :
    eintg1 = ((10**elog10high1)**(1-gammagen) - (10**elog10low1)**(1-gammagen)) / (1-gammagen)

#
# prepare data with E^(-gammagen) 
#
data1 = np.random.uniform(elog10low1, elog10high1, nevts1)

#
# calculate fill weight for data1. 
# use np.array so that we may scale the weights later!
# (python's list doesn't support even a scalar scale...)
#
flux1 = np.array([weightfunc(eintg1, ene) for ene in data1])

#
# prepare drawing canvas (figure)
#
fig = plt.figure()

#
# option 1: use pylab
#
print "option 1"
ax1 = fig.add_subplot(2,2,1)
n, bins, patches = ax1.hist(data1, nbins1, weights=flux1/nevts1, histtype='step', color="brown", label="pylab.hist")
ax1.legend()

# draw same fig with log scale
ax3 = fig.add_subplot(2,2,3)
ax3.set_yscale('log')
n, bins, patches = ax3.hist(data1, nbins1, weights=flux1/nevts1, histtype='step', color="brown", label="pylab.hist")

#
# option 2: use numpy.histogram
#
yval, binEdges = np.histogram(data1, nbins1, weights=flux1/nevts1)

print "option 2"
print yval
print binEdges
print binEdges[1:]
print binEdges[:-1]

#
# use pylab.step
#
ax2 = fig.add_subplot(2,2,2)
ax2.step(binEdges[0:-1], yval, where='post', color="green", label="np step",linestyle="dashed", linewidth=3)


#
# prepare data for line plotting
#
def format (yval, binedges) :
    x = []
    y = []
    for index, value in enumerate(yval) :
        print "y = %f, binedge_low = %f, binedge_hi = %f" % (value, binedges[index],binedges[index+1])
        x.append(binedges[index])
        x.append(binedges[index+1])
        y.append(value)
        y.append(value)

    # append zeros to close histos
    x.insert(0, x[0])
    y.insert(0, 1e-30)
    x.append(x[-1])
    y.append(1e-30)

    return x, y

x, y = format(yval, binEdges)
print x
print y

ax2.plot(x, y, '--', color="magenta", label="np line",linewidth=1.5)
ax2.legend()

# show in log scale
ax4 = fig.add_subplot(2,2,4)
ax4.set_yscale('log')
ax4.set_ylim([1e-9, 1e3])
ax4.plot(x, y, '--', color="magenta", label="np line",linewidth=1.5)
ax4.legend()

#
# option 3: bin center 
#

bincenters = 0.5*(binEdges[1:]+binEdges[:-1])

print "option 3"
print bincenters

ax3.plot(bincenters, yval,'o-', color="blue", label="np bincenter", linewidth=1.5)
ax3.legend()

plt.show()


TH2D

基本的に、np.arrayで1次元の代わりに2次元の配列を作れば良い。

import numpy as np
import matplotlib.pyplot as plt
 
n = 100
x = np.random.rand(n)
y = x + 0.5*np.random.normal(size=n)

# 1) forで回してfillする場合

nx = 3
ny = 6
xmin = 0.
xmax = 1.
dx = (xmax - xmin)/nx
ymin = -2. 
ymax = 3.
dy = (ymax - ymin)/ny
hist1 = np.zeros((nx, ny))
for index, xvalue in enumerate(x) :
    yvalue = y[index]
    if xvalue < xmin or xvalue > xmax :
        continue
    if yvalue < ymin or yvalue > ymax :
        continue
    xi = int((xvalue - xmin)/dx)
    yi = int((yvalue - ymin)/dy)
    hist1[xi,yi] += 1.0

# 軸を用意。pcolor, pcolormesh関数はnx+1 x ny+1の二次元配列を
# それぞれxaxis, yaxisに指定しなくてはならない。
# 以下はたまにround offのせいかnx+2の要素数を返してくるのでダメ
# xaxis,yaxis = np.mgrid[slice(xmin, xmax+dx, dx), slice(ymin,  ymax+dy, dy)] 

xaxis,yaxis = np.mgrid[xmin:xmax:(nx+1)*1j, ymin:ymax:(ny+1)*1j]

print xaxis
print yaxis
print hist1

fig1 = plt.figure()
ax1 = fig1.add_subplot(1,2,1)
phist1 = ax1.pcolormesh(xaxis,yaxis,hist1)
#phist1 = ax1.pcolor(xaxis,yaxis,hist1)
plt.colorbar(phist1)


# 2) np.histogram2Dを使う

ax2 = fig1.add_subplot(1,2,2)
hist2, xedges, yedges = np.histogram2d(x, y, bins=(nx, ny), range=([xmin,xmax], [ymin,ymax]))

# 軸を用意。xedges, yedges は元々nx+1, ny+1の配列を返してくるので、
# これを2次元に変換。
# ただし、引数の順番に注意!!
yaxis2, xaxis2 = np.meshgrid(yedges, xedges)
print xaxis2
print yaxis2
print hist2
phist2 = ax2.pcolormesh(xaxis2, yaxis2, hist2)
#phist2 = ax2.pcolor(xaxis2, yaxis2, hist2)
plt.colorbar(phist2)

plt.show()

少々面倒な点は、軸の配列の用意。2次元の絵を書くのだから、軸は1次元かというとそうではなくて、軸も2次元で作ってやる必要がある。
しかも、i番目のデータはXi, Xi+1の区間にプロットされるので、軸は常にnx+1, ny+1個のデータが必要になる。

その方法として、np.mgridを使う方法と、np.meshgridを使う方法を挙げた。
mgridの方はまだ直感的に分かり易いのだけれど、[]の中身の指定の仕方がかなり特殊。
普通にslice(xmin, xmax+dx, dx)(xmin:xmax+dx:dxと同じ意味)とかやっちゃうと、dxが無理数の場合roundoff errorがでちゃって、nx+1個どころかnx+2になってしまい、pcolorで描こうとすると配列の次元が合わないと怒られる。

散々ネットで探した挙げ句、sliceの第三引数は通常stepサイズだけど、複素数で指定してやると分割数になることが分かった。
というわけで、分割数はnx+1になるので、それに複素数の1jをかけてやれば複素数で指定できる。

一方、meshgridではyaxis, xaxisの順に引数を指定してやり、戻り値もその順番で受け取らないと、何故か例題1)と同じにならない。

巷の例題では、結果の配列の方を上下左右入れ替える(np.rot90(hist2), np.flipud(hist2))としているものが多いのだけれど、そんな解決方法はイヤで、ちゃんとx軸は横に、y軸は縦に描いてほしいのである(笑)
というわけで、meshgridに食わせる引数と戻り値の順番の方を弄った方が良い、と思った次第。

ちなみに、pcolorとpcolormeshはほとんど同じことをやるのだけど、pcolormeshの方が早いらしい。(あと戻り値オブジェクトのクラスが違う)