🐍

実験データ解析のためのPython

Pythonで物理実験のデータ解析を行うためのメモ

はじめに
データ読み込み/書き出し
dummy_data.txt
events の中身
波形解析
波形プロット
波形特徴量計算
ヒストグラム
フィッティング
ROOTファイルの書き込み/読み込み
書き込み
読み込み
おわりに

はじめに

宇宙/素粒子/原子核実験においては, 業界標準としてC++/ROOTが用いられていることが多いと思います. 一方で, 共同研究者がROOTに詳しくなかったり, pythonは簡単だと聞きかじった指導教員の圧力だったり, あるいは他の政治的理由によりpythonを使わなければならないこともあるでしょう. より積極的にpythonを使う理由としては, ROOTに対して少ない行数でコードを書けることと, 最終的に出来上がるプロットがROOTよりかなり綺麗であることがあげられます. また, ROOT全体をbuildするのに対して軽量です. しかし, pythonとよく使われるモジュールはデータ解析のためだけに作られたわけではないので, 実験データ解析のための文書はあまりありません. また, python本体は非常に低速なので, C++やFortranで実装されているnumpyの関数を多用する必要が (将来numbaなどを使うためにも) あります. 本メモではコードを含めてひととおりpythonでのやり方を詳解します.

numpy, matplotlib, pandas などの純粋なpython モジュールのみを使って解析を行います.

データ解析には, コードを書き, その場でプロットを見て修正できる jupyter lab が非常に便利です. 作成したノートブックにはコードを図が両方含まれているので, 共同研究者と共有するのにも便利です.

最初に使うモジュール類をインポートします.

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

データ読み込み/書き出し

ROOTにはTTreeという実験データの保存システムがあります. これは唯一無二で代替不可能な実験の大規模データを扱うためのフォーマットです. 例えば, TTreeのインスタンス tree があって, 各イベントに energy という変数 (正確にはBranch) が入っていれば, tree->Draw("energy") だけでヒストグラムを非常に高速に書くことができます. これは, Treeのデータ保持フォーマットが工夫を凝らして作成されているためです. 一方で, numpyで似たようなことがしたければ, numpy.structured_array にデータを集約するのがいいでしょう. ただし, せいぜいコンピューターのメモリに載る程度のデータ量しか一度に扱えません.

任意の波形データを取得したとします. この際のフォーマットは検出器とData acquisition (DAQ) system により千差万別ですが, とにかくファイルを読み込んで, python上で扱えるようにします.

例えば, 以下のような波形データがあったとします.

‣

dummy_data.txt

1列目は unix timestamp (秒), 2列目はマイクロ秒, その後, 波形データが16列並んでいます. 計32行あり, これは32イベント取得されたという想定です. このくらいだと, pandas.readline() で読んだ方が早いと思われるかもしれませんが, 現実には波形データは16列よりかなり多く, DataFrameとして扱うのが大変なため, 愚直に読み込みます. DataFrameに波形のようなリストを入れることはできません.

time_s, time_us, waveforms = [], [], []
with open('dummy_data.txt', 'r') as f:
    line_list = f.readlines()
    for line in line_list:
        l = line.split(',')
        time_s.append(int(l[0]))
        time_us.append(int(l[1]))
        waveforms.append([int(v) for v in l[2:]])

いったん pythonのリストに格納します.

waveforms はこんな感じの2次元リストになっているはずです.

[[131, 126, 134, ..., 126, 129, 127],
 [127, ......................., 125],
 [.................................],
 .................略................,
 [134, 129, 128, ..., 130, 127, 121]]

この段階では, 各変数が同じイベントに紐づいていません. そこで, numpy.structured_array を用いて, イベント毎に紐づけます.

number_of_events = len(time_s)
waveform_length = len(waveforms[0])

dtype = [('id', int), ('time', int), ('time_us', int), ('wf', int, waveform_length)]
events = np.zeros(shape=number_of_events, dtype=dtype)
events['id'] = np.arange(number_of_events)
events['time'] = np.array(time_s)
events['time_us'] = np.array(time_us)
events['wf'] = np.array(waveforms)

この時の events の中身は

‣

events の中身

このようになっており, 適切にイベント毎に構造化されてデータを保存できています.

ファイルに保存する際は, np.save('data.npy', events) ,

ファイルから読み込むときは events = np.load('data.npy') で簡単に取り扱えます.

pythonで解析を行う際は, テキストデータをリストに格納するという手順が重たいので, とりあえず生データをnumpyのarrayにすることをめざすと良いでしょう.

波形解析

波形プロット

まずは, 波形を描いてみましょう. 前項により numpy.structured_array ができていれば,

plt.plot(events['wf'][0])

これで0イベントめの波形を書くことができます.

複数のイベントの波形をまとめて描きたければ, plt.plot(events['wf'][0:10].T);とするだけです. .T は転置のためにつけています.

平均波形を描きたければ, plt.plot(events['wf'].mean(axis=0)) で簡単に描くことができます.このような処理を非常に簡単に書くことができるのは, 行列としてデータを扱えるnumpyの大きな利点です.

plt.plot(events['wf'][0:3].T, label='Example')
plt.plot(events['wf'].mean(axis=0), label='Average', color='r', lw=5)
plt.fill_between(np.arange(waveform_length), 
                 np.percentile(events['wf'], 25, axis=0), 
                 np.percentile(events['wf'], 75, axis=0),
                 color='r', alpha=0.2, label='Quantile region')
plt.xlabel('Time bin')
plt.ylabel('Amplitude')
plt.legend()

数行でそれっぽいプロットがかけました.

波形特徴量計算

実際の実験のように, 本データのベースラインにはオフセットがのっています. まずはこれをイベント毎に計算して引き算しましょう.

baselines = events['wf'][:,0:2].mean(axis=1)

これで, 波形の最初の0ビン目から1ビン目まで (pythonのスライスは後ろ側を含まないため2-1) を平均したベースラインを各イベント毎に得ます. これを波形から引き算すれば, ベースラインを差し引いた波形を書くことができます.

wfs = events['wf'] - np.expand_dims(baselines, 1)

ここで, events[’wf’] は (32, 16) の形状, baselines は(32,　) なので, np.expand_dims() を使ってnumpyのbroadcast rule に載るようにしてやります.

plt.plot(wfs[0:10,:].T);

実際の実験では, 最初の2ビンよりもっと多くのビンをベースラインを引くのに使うはずです. 今回は非常に少ないビンで無理やりベースラインを引いたため, ノイズによるバイアスが大きくなっています.

ここまでくると, 波形の大きさや最大値, 最大値の位置を計算するのは非常に簡単になって,

area = wfs.sum(axis=1) # 波形の面積
amplitude = wfs.max(axis=1) # 波形の最大値
t_amp = wfs.argmax(axis=1) # 波形の最大値のBin

で, それぞれ1行で求めることができます. (そして, 高速です)

ここまでくると, 感のいい人は先ほど eventsを作ったときにこれらも dtype として作っておいた方が良かったということがわかると思います.

例えば,

number_of_events = len(time_s)
waveform_length = len(waveforms[0])

dtype = [('id', int), ('time_s', int), ('time_us', int), 
         ('time', np.float64),
         ('baseline', np.float64),
         ('area', np.float64),
         ('amplitude', np.float64),
         ('t_amp', np.int32),
         ('wf', int, waveform_length)]
events = np.zeros(shape=number_of_events, dtype=dtype)
events['id'] = np.arange(number_of_events)
events['time_s'] = np.array(time_s)
events['time_us'] = np.array(time_us)
events['time'] = events['time_s'] + 1.e-6*events['time_us']
wfs = np.array(waveforms)
events['baseline'] = wfs[:,0:2].mean(axis=1)
events['wf'] = wfs - np.expand_dims(events['baseline'], 1)
events['area'] = events['wf'].sum(axis=1)
events['amplitude'] = events['wf'].max(axis=1)
events['t_amp'] = events['wf'].argmax(axis=1)

あるいは, 波形から必要なデータを抽出してしまえば, pandasのDataFrameに入れた方が楽です.

df = pd.DataFrame(events['id'], columns=['id'])
df['time'] = events['time']
df['time_us'] = events['time_us']
df['area'] = wfs.sum(axis=1)
df['amplitude'] = wfs.max(axis=1)
df['t_amp'] = wfs.argmax(axis=1)

一般に, 波形などの低レベルのデータ処理は高い計算量を要求します. ここまでの部分をきちんと書いてしまえば, 残りの部分は適当に書いてもなんとかなることがほとんどだと思います.

ヒストグラム

前項の最後に求めた areaなどは1次元でイベント数に等しい要素数のnumpyの配列です.

これをヒストグラムとして表示するためには,

plt.hist(area, range=(-30, 270), bins=30, histtype='step')

この一行で済みます. range, bins, histtypeは解析と好みに応じてつけてください.

また, 複数のヒストグラムを作ることを見越して,

option = dict(histtype='step', range=(0, 300), bins=30)
plt.hist(events['area'], **option)

と引数を括り出しておくと便利です.

pandasのDataFrameに入れていた場合はもっと簡単で,df.hist('area') だけで済みます.

2次元ヒストグラムは

plt.hist2d(events['area'], events['amplitude'])

のようにして書きます.

フィッティング

非常に残念なことに, 具体的に以下のことを要求すると, 満足できるフィッティングツールはpythonにはありません.

簡単なインストール
簡単なインターフェース
豊富なプリセットモデル (gaussian, exponential, 多項式, etc.) + ユーザー定義モデル
物理でよく使われるコスト関数 ( $\chi^2$ , (Un)binned likelihood)とpull-term追加の余地
可視化のサポート
継続した開発体制
高速 (backendがC++またはFortran)

Scipyのcurve_fitはおもちゃなので物理解析を真面目に行うためには適しません.

部分的にこれらの要求を満たすのは probfit (公式), もしくはzfit (公式)です.

probfitはかなり要求を満たしていますが, 継続した開発は既に行われていないようです.

probfitのチュートリアルを置いておきます.

probfit_tutorial_ja/probfit_tutorial1.ipynb at master · mzks/probfit_tutorial_ja

probfitによるFittingのチュートリアル. Contribute to mzks/probfit_tutorial_ja development by creating an account on GitHub.

github.com

probfit_tutorial_ja/probfit_tutorial1.ipynb at master · mzks/probfit_tutorial_ja

zfitはインターフェースが難しく, 可視化の支援もありません. そのため, wrapperを書いておきました.

GitHub - mzks/mzfit: zfit wrapper for lazy analysts

zfit wrapper for lazy analysts GitHub Author: Keita Mizukoshi (Kobe Univ. mzks@stu.kobe-u.ac.jp, @mzks) is a nice fitting tools on python, built on object-oriented interface. However, I would sometime like to fit easily. I do not always want to care minimizers and cost function. I usually take an approach for good fitting, try-error-retry with visual environment.

github.com

GitHub - mzks/mzfit: zfit wrapper for lazy analysts

ただし, しばらくメンテナンスする時間がないため, 放置してしまっています.

ROOTファイルの書き込み/読み込み

uproot という純粋なpythonライブラリがその機能を提供しています.

import uproot で使えるはずです.

uproot5

scikit-hep ⋅ a year ago

書き込み

df という pandas.DataFrameがあれば,

with uproot.recreate('test1.root') as f:
    f['tree'] = df

これだけで書き込めます. branchの型はdataframeのものがそのまま使われます.

test1.root というファイルに, tree という名前でTTreeが書き込まれます.

また, 上で作成したような, numpy.structured_array (ここでは events)があれば, 同様に,

with uproot.recreate('test2.root') as f:
    f['tree'] = events

とすることが可能です. これが最もシンプルな方法だと思います.

または, ROOTっぽくTTreeを作ってから, 1イベントずつ順次書き込むこともできます.

with uproot.recreate('test3.root') as f:
    f.mktree('tree', {'x': np.int32, 'y': np.float64})
    
    for i in range(100):
        f['tree'].extend({'x': np.array([i]), 'y': np.array([i*2.])})

ただし, pythonのloopは遅いのであまり現実的ではありません. extend がarrayを値としてとる辞書を要求することから, ある程度まとまった処理をbatchで走らせるようなことが想定されていると思います. この機能が必要になることはおそらくないでしょう.

TTree以外にも, np.histogram や, string も書き込むことが可能です.

with uproot.recreate('test2.root') as f:
    f['tree'] = events
		f['author'] = 'Keita Mizukoshi'

読み込み

ファイルを開いて, 中身を確認するためには,

uproot.open('test2.root').keys()

種類も確認したい時は,

uproot.open('test2.root').classnames()

で .ls 相当のことができます.

uproot.open('test2.root').classnames()

{'tree;1': 'TTree', 'author;1': 'TObjString'} このような出力が得られます.

TTreeの名前がわかれば (ここでは tree), 簡単にpandas.DataFrameにすることが可能です.

df = uproot.open('test2.root')['tree'].arrays(library='pd')

ただし, pip install awkward-pandas しておく必要があるかもしれません.

最近のpandas.DataFrameはcolumnにリストが入ることが許されているので, かなりシンプルに描けるようになりました. 一方で, numpyの行列演算を活かすことができないので, 波形解析の時はnumpy, 特徴量に落ちてきてヒストグラムを中心に解析するようになるのであればDataFrameとうまく使い分ける必要があります.

特定のBranchだけnumpyで落としてくるのは

wf = uproot.open('test2.root')['tree']['wf'].array(library='np')

このような感じで, library=’np’ を指定すればできます.

お行儀よく,

with uproot.open('test2.root') as f:
    df = f['tree'].arrays(library='pd')

としてもいいと思いますが, 未知のTTreeを読む時は大抵色々やってみてデータを見ながらごにょごにょすると思うので, 私はいつもclose処理をサボっています.

おわりに

慣れればpythonも便利に使っていくことができると思います. 他に取り扱ってほしい/取り扱うべきテーマがあったり, 誤りやより良い方法を見つけた際はご連絡ください.