物理の駅 Physics station by 現役研究者

テクノロジーは共有されてこそ栄える

Python+uproot+Awkward ArrayでCERN ROOTファイルを読み、書き出す

この記事の Awkward のバージョンは1.X.Xです

業界にも依るとは思うが、イベント単位になっているデータの標準形式がROOTファイルというところは多いだろう。ただ、数10GB級のROOTファイルを一斉に処理したいような需要ではなければ、あえてROOT(というかC++)を使って解析する必要性は感じられないので、PythonネイティブでCERN ROOTファイルの読み取りが可能なuprootを使ってみた。(こんかいは読むだけ。書き出しは今度)

uprootは、numpyなど、一部のライブラリに依存しているが、CERNのROOTとは完全に独立しているので、Windowsユーザーでも素直にインストールできるだろう。

pip3 install uproot

ROOTファイルを読み込んでみる。

import uproot
file = uproot.open("rootfiles/run0254.root")
print(file.keys())

# ['tree;1']

今回はtreeという名前のTTreeが1つ入っているファイルを扱うので、次のように読み込んだほうが良い。

tree = uproot.open("rootfiles/run0254.root:tree")
tree.show()
# name                 | typename                 | interpretation                
# ---------------------+--------------------------+-------------------------------
# RUNNUMBER            | int32_t                  | AsDtype('>i4')
# EVENTNUMBER          | int32_t                  | AsDtype('>i4')
# TIMESTAMP            | uint64_t                 | AsDtype('>u8')
# 以下略

ここで得られるTreeはROOTファイルのバージョンにも依る?がuprootのTTreeというクラスっぽいものになっていてやや扱いが難しい。

print(type(tree))
# <class 'uproot.models.TTree.Model_TTree_v19'>

Pythonで配列といえばnumpyかpandasかを使うと思う。例えばnumpy配列にしてグラフ化するには以下のようにする。

print(type(tree["TIMESTAMP"].array(library="np")))
#print(type(tree["TIMESTAMP"].array(library="ak")))
#print(type(tree["TIMESTAMP"].array(library="pd")))

import matplotlib.pyplot as plt
plt.hist(tree["TIMESTAMP"].array(library="np"), bins=100)
plt.show()

# <class 'numpy.ndarray'>

ただ、データとしてかなり扱いづらかったので、Awkward Arrayというライブラリを使ってみることにする。

https://awkward-array.readthedocs.io/en/latest/index.htmlawkward-array.readthedocs.io

一連のファイルを連続で読み取れるuproot.concatenateを使って、必要なメンバー変数だけをAwkward Array形式にする。ファイル名にはワイルドカードを使えたり、配列を使えたりする。今回は読み込む名前として3つのメンバー変数を指定し、Awkward Arrayで読むためにlibrary = "ak"を指定した。例では書いていないが、カット条件を指定しながら必要なデータだけ読むことも可能 cut="(AX!=-9999)&(AY!=-9999)&(TX!=-9999)&(TY!=-9999)"

tree = uproot.concatenate("rootfiles/run0254.root:tree",
                   filter_name=["RUNNUMBER", "EVENTNUMBER", "TIMESTAMP"],
                   library ="ak")
print(type(tree))
print(tree[0]) #0番目のイベントを取得 (データ構造上遅い)
print(list(tree[0].to_list().keys())) # メンバー一覧 (データ構造上遅い)
print(tree["RUNNUMBER"]) # RUNNUMBERを取得  (データ構造上速い)

# <class 'awkward.highlevel.Array'>
# {RUNNUMBER: 254, EVENTNUMBER: 1, TIMESTAMP: 152232731237482}
# [RUNNUMBER, EVENTNUMBER, TIMESTAMP, ...]
# [254, 254, 254, 254, 254, 254, 254, 254, ... 254, 254, 254, 254, 254, 254, 254, 254]

Awkward Arrayが便利なのは、[]演算子に整数(int)を入れれば i 番目のイベントのデータが得られ、メンバー変数を指定すればその変数の配列が得られるということである。numpyのように.Tで転置して云々みたいなことは必要ない。

既存のメンバー変数の値を変更したいとき、または新しいメンバー変数を追加したいときも、比較的直感的に行える。

tree["NEWMEMBER"] = [i for i in range(len(tree))]
# tree["NEWMEMBER"] = 0 #で全てを0にできる
print(tree["NEWMEMBER"])

なお、i番目のイベントのデータだけを書き換えるということはできない。

tree["NEWMEMBER"][0] = 10

以下のようなエラーが出る

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-141-33465948660c> in <module>
----> 1 tree["NEWMEMBER"][0] = 10

~/.local/lib/python3.8/site-packages/awkward/highlevel.py in __setitem__(self, where, what)
   1062                 or (isinstance(where, tuple) and all(isinstance(x, str) for x in where))
   1063             ):
-> 1064                 raise TypeError(
   1065                     "only fields may be assigned in-place (by field name)"
   1066                     + ak._util.exception_suffix(__file__)

TypeError: only fields may be assigned in-place (by field name)

(https://github.com/scikit-hep/awkward-1.0/blob/1.7.0/src/awkward/highlevel.py#L1066)

大量のメンバーを含むROOTファイルのうち一部しか使わない場合や、カット条件によりごく一部のデータしか使わない場合、Awkward Array を書き出しておくと便利である。

pip3 install pyarrow

でpyarrowをインストールし、to_parquetでparquet形式で出力するのが良いだろう。

import awkward as ak
ak.to_parquet(tree, "tree.parquet")

つまり、ROOTファイルをparquet形式に変換するサンプルコードは以下の通り。

import uproot
import awkward as ak
for run in [113,120]:
    tree = uproot.concatenate(f"rootfiles/run{run:04d}.root:tree",
                       filter_name=["TIMESTAMP", "F5X", "F5A"],
                       library ="ak")
    ak.to_parquet(tree, f"tree{run:04d}.parquet")
    del tree

詳細を確認していないが、stackoverflowによると浮動小数点数を含む場合はuse_dictionary=Falseを使ったほうがいいらしい。

ak.to_parquet(tree, "tree.parquet", use_dictionary=False)

圧縮をしたくない場合は compression="NONE"を使う。

ak.to_parquet(tree, "tree.parquet", compression="NONE")

読み込むときはfrom_parquetを使う。

tree = ak.from_parquet("tree.parquet", lazy=True)

参照

stackoverflow.com

Awkwardからの出力方法

https://awkward-array.readthedocs.io/en/latest/_auto/ak.to_parquet.htmlawkward-array.readthedocs.io

Parquetの全オプションについて

arrow.apache.org

ある条件で抽出したい場合(=フィルタリングする場合)、numpyと同じように扱うことができる

new_tree = tree[tree["Member1"]!=-9999]

複数の条件が絡む場合は各条件を()で囲むべし

new_tree1 = tree[(tree["Member1"]!=-9999) & (tree["Member2"]!=-9999)]
new_tree2 = tree[(tree["Member1"]!=-9999) | (tree["Member2"]!=-9999)]

カット条件

cut1 = (tree["Member1"]!=-9999)
cut2 = (tree["Member2"]!=-9999)
new_tree1 = tree[cut1 & cut2] #AND
new_tree2 = tree[cut1 | cut2] #OR
new_tree3 = tree[~(cut1 ! cut2)] #否定

ak.any, ak.all, ak.sum の使い方

array = ak.Array([[1,1,3],[2,2],[3,1],[float("nan"),2],[float("nan"),1]])
print(ak.any(np.isnan(array)))
# True どれかなので
print(ak.any(np.isnan(array),axis=0)) #axisは縦横
# [True, False, False] 調べる方向が縦なので
print(ak.any(np.isnan(array),axis=1)) #axisは縦横
# [False, False, False, True, True] 調べる方向が横なので
print(ak.all(np.isnan(array)))
# False 全てnanではないので 
print(ak.sum(np.isnan(array)))
# 2 nanは2個なので
print((np.isnan(array)))
# [[False, False, False], [False, False], [False, False], [True, False], [True, False]] 全ての要素を判定

2つ以上のするときは、awkward配列を結合するときは、ak.concatenateを使うとよい。

tree3 = ak.concatenate([tree1,tree2])

値が配列になってるとき、配列のn番目だけを抜き出すのはawkwardだけではできない(っぽい)のでnumpyの力を借りる。

ch = 0
tree["ADC"].to_numpy().T[ch]

値が配列になっているとき、配列の平均値を得る。awkwardだけではできない(っぽい)のでnumpyの力を借りる。

tree["ADCMEAN"] = np.mean(tree["ADC"].to_numpy(),axis=1)

配列の一部だけの平均値を得る方法。numpy (numpy.ndarray)の力を借りる。例えば、ADC値の3番目から8番目までの6チャンネル分の平均値を計算するときは次のようにする。

tree["ADCMEAN39"] = np.mean(tree["ADC"].to_numpy()[:,3:9],axis=1)

Awkward ArrayのTree同士の中身が一致しているかどうか確認する

def is_equal_ak(tree1,tree2):
    if type(tree1)!=type(tree2):return False
    if len(tree1)!=len(tree2):return False
    if len(tree1)!=len(tree2):return False
    if list(tree1[0].to_list().keys())!=list(tree2[0].to_list().keys()):return False
    for key in list(tree1[0].to_list().keys()):
        if ak.all(tree1[key] == tree2[key])==False:return False
    return True

https://awkward-array.readthedocs.io/en/latest/_auto/ak.Array.html#filteringawkward-array.readthedocs.io

github.com