PyROOTでルートファイルのtreeのbranchを読ませるのに苦労した話

PyROOTに関して

あるときケムコワは思いました。
「ROOTは使いにくい！！！」
と。
ポインタだのnewだのよくわからないし独自型を返す関数が独自型の引数をとるのでそういう部分に囚われてにっちもさっちも解析が進みません。
すぐSegmentationViolaionとか言われるし…

そこでPyROOTという存在に目をつけました。
ちょうど研究でPythonを使っていたので良い機会でした。

研究には当然TTreeオブジェクトのブランチを読んでゴニョゴニョするわけですが、それがPyROOTでやるのに苦労したというのが今回の記事の内容になります。
ただしある程度ROOTの基礎知識を要求するのは許してください。

PyROOTの導入に関して

ちょっとPyROOTの導入に関して述べておきます
まずなんらかの手段でROOTとPython3を導入してください。
以下Linuxを想定します
rootのインストールディレクトリに行きます。
rootにパスが通っていれば、

$ which root

として出てくる場所がインストールディレクトリです。
そこにthisroot.shが多分あるので、

$ bash thisroot.sh

とすれば何か知らないけど設定完了するはずです。(自分はこれだけでOKでした)

確認するには、
pythonの対話実行モードを立ち上げてROOTライブラリを読み込んで何も表示されなければOKです。
なお、ライブラリそのものが重いので読み込みには結構時間がかかります。

$ python3
>>> import ROOT #何も表示されずに次の入力ラインが表示されればOK
>>>

PyROOTでのtreeのブランチの読ませ方

.

結論

先に結論を述べておきます。

import ROOT
# TTree オブジェクトの取得
rootfile = ROOT.TFile(rootfilename)
tree = rootfile.Get("tree") # 右辺の引数はルートファイル上のTTreeオブジェクトの名前

# ほしいイベントに移動
tree.GetEntry(0)

#「ふつう」のブランチの読み方 ---------------------------------------

# ブランチがスカラーである場合
# ブランチ名をevenumとする
evenum = tree.evenum

# ブランチが配列である場合
# ブランチ名をmomentumとする
momentum = tree.momentum

# C++組み込み型のメンバ変数をもつ構造体の場合 -------------------------

# 例として次のような構造体を仮定
# event
#  ├── evenum : int型
#  └──  momentum : float型1次元配列

struct_branch = tree.GetBranch("struct_branch_name")
# ここでは実際には struct_branch = tree.GetBranch("event") とする
evenum = struct_branch.GetLeaf("evenum").Getvalue()
momentum = []
for i in range(struct_branch.GetLeaf("momentum").GetLen()):
    momentum.append(struct_branch.GetLeaf("momentum").Getvalue(i))

経緯

初めは、伝統的な方法であるSetBranchAddressを用いてtreeを読むことを試みました。
実際自分がこれまでROOTで解析を行う際にはそれを用いていました。
例：

double momentum[3];
tree->SetBranchAddress("momentum",&momentum)

また、TTreeReaderによって読むことも試みました。
しかし両者ともにうまく行きません。

最後にたどり着いた、Kamonowikiに書いてある情報では、

variable = tree.brach_name

として簡潔にブランチの値を取得できるようでした。これである程度うまくいったのですが、やはり、ある問題を抱えていました。

生じていた問題

問題となっていたのは、C++で定義された構造体のブランチがあったことです。その構造体そのものはPython上で定義されているわけではないので、

evenum = tree.event.evenum

と言った処理はできませんでした。というか、eventという名前のブランチが存在するにも関わらず、エラーメッセージは、
「eventという名前のメンバ変数は持っていません」
といったようなことを言ってきます。

ここで試しに

evenum = tree.evenum

としたところ、値を取得することができました。
どうやら、tree.branch_nameの形でアクセスできる変数は、

$root file.root
root[0] tree->Show()

というようにしたときに表示される変数名一覧を参照できるようでした。
自分の実際に解析しているデータにおける表示例

さらなる問題

しかしこれでも問題があります。それは、実は、全く同じ名前のメンバ変数をもつ構造体ブランチを2つルートファイルが持っていたことでした。
pla1, pla2という2つのブランチがあり、それぞれthickness, tdcという2つのメンバ変数を持っていたとします。

pla1
  ├── thickness : double
  └── tdc : float型1次元配列
pla２
  ├── thickness : double
  └── tdc : float型1次元配列

ROOT上でtree->Print()としてブランチの一覧を確認したとき、pla1のほうが先にある(ブランチナンバーが先)とします。
このとき、PyROOT上でtree.thicknessやtree.tdcとして値を取得できるのはpla1の方のみです。もう一度tree.thicknessなどとしてもpla1の値しか取得できません。

解決策

これでは困るので、ChatGPTに聞きながら色々とやってみました。その結果が上記の方法です。
TBranchオブジェクトを取得し、そこからTLeafオブジェクトを取得し、GetValue()によって値を取得するという方法になります。

上に挙げた例をもう少し詳しく説明するとこのようになります。

# 例として次のような構造体を仮定
# event
#  ├── evenum : int型
#  └── momentum : float型1次元配列

# TBranchオブジェクトを取得
struct_branch = tree.GetBranch("event")
# evenumに対応するTLeafオブジェクトを取得
evenum_leaf = struct_branch.GetLeaf("evenum")
# 値を取得
evenum = evenum_leaf.GetValue()

# momentumに対応するTLeafオブジェクトを取得
momentum_leaf =  struct_branch.GetLeaf("momentum")
# momentum_leafの長さを取得
length = momentum_leaf.GetLen()
# 各配列要素を取得して配列に詰める
momentum = []
for i in range(length):
    momentum.append(momentum_leaf.GetValue(i))

たとえメンバ変数が同じであっても

pla1_branch = tree.GetBranch("pla1")
pla2_branch = tree.GetBranch("pla2")

としてブランチを区別できるので、解決できたというわけです。

ちなみに、この方法で構造体でないブランチの値を取得することも可能です。
runという(たとえば整数型の)ブランチを持っているのであれば、

run_branch = tree.GetBranch("run")
run_leaf = tree.GetLeaf("run")
run = run_leaf.GetValue()

として値を取得できます。

ROOTのそもそもの話

rootファイルは高度に構造化されており、

TFile オブジェクト
    ├── ヒストグラムなどのオブジェクト
    └── TTreeオブジェクト
        └── TBranchオブジェクト
            └── TLeafオブジェクト
                └── 値

というさながら一本の樹のような構造をしているわけです(他にもメンバ変数やメンバ関数などがたくさんありますが)。それを踏まえれば今回のような方法は極めて素朴な方法であると言えます。
なぜこんな方法に気づかなかったのかというと、調べて出てくる情報はすべてSetBranchAddressとかTTreeReaderによる方法であったからです。
なぜなんでしょうかね…？

未解決の問題

実はまだ未解決の(おそらく自分が解決方法を知らないだけの)問題があります。

• ブランチが多次元配列の場合
• ブランチが構造体であったとして、そのメンバ変数がさらに構造体であった場合
• ブランチがstd::vectorであった場合

とりあえず現状で解析を進める上で障害にはならないので放置しています。というのは、そういうブランチを持つrootファイルが手元にありません。
まぁないなら作れば良いのですが面倒なのでまた今度試します。
学会も近いし
1つめに関しては、たぶん簡単に解決できて、

leaf.GetValue(i,j)

で各要素を取得して値を詰め直せばOKです。

2つめに関してはTLeafオブジェクトがさらにTLeafオブジェクトを持つのかな…と想像してます(想像だけ)

3つ目は配列と同じように取得できる…?（知らないけど…)

さいごに

というわけでPyROOTのtreeのbranchを読む方法を述べました。
たぶんほとんど同様の方法でROOTでブランチを読むことができると思います。(というかこの方法がprimitiveな気がしないでもない)

なにか質問等ありましたらどうぞお気軽にお願いします。