【高能物理.ROOT】ROOT文件IO(下)

ROOT文件IO实战

ROOT文件(TTree)读出

读取TTree最重要的是搞清楚TTree的结构。
当碰到每个事例的其中一类数据(Branch)是二维数组甚至更高维情况就很难办,这种时候不能直接.Draw或者.Histo
一般读入数据后以什么方式存储是个问题。

1. 保存为root本身自带格式,后续一些处理分析需要查相应的代码,例如RooDataSet
2. 转存为numpy数组,后续处理分析的相应资料更多,像一些sklean
   保存为什么格式根据不同问题而定，我更喜欢对于涉及事例间问题转为numpy,而每个事例独立更喜欢用RDataFram格式

直接读取TTree

有一种直接读TTree的方法¹,好用,但是我不常用

import ROOT
import array
# 创建TTree
tree=ROOT.TTree('testtree','title')
a=array.array("I",[9])
b=array.array("d",[1,2,3,4,5,6,7,8,9])
tree.Branch('a',a,"a/i")
tree.Branch('b',b,"b[a]/D")
tree.Fill()
a[0]=8
b.pop()
b[0]=1.1
tree.Fill()
# 读取TTree
for i,Evt in enumerate(tree):
    print('Event:',i)
    print('a=',Evt.a)
    print('b=',end='')
    for j in range(Evt.a):
        print(Evt.b[j],end=' ')
    print()

运行结果

Event: 0
a= 9
b=1.0 2.0 3.0 4.0 5.0 6.0 7.0 8.0 9.0 
Event: 1
a= 8
b=1.1 2.0 3.0 4.0 5.0 6.0 7.0 8.0 

有时候会用这种方法读,因为后一种方法占有内存多.

RDataFrame读取TTree

这种读入方式可以批量筛选数据(CUT)，还可以转numpy,还可以接入ROOT自带的一些分析工具。只不过一些函数可能需要自己写(C++).

import ROOT
#假设有一个函数想要批量应用到每个事例,在PyROOT中要如此定义此函数
ROOT.gInterpreter.Declare(
"""
ROOT::RVecD GeteDepDingdian(ROOT::RVecD &eDep,ROOT::RVecD &eDepx,ROOT::RVecD &eDepy,ROOT::RVecD &eDepz)
{
    ROOT::RVecD r={0,0,0};
    double hitall=0.0;
    for(int i =0;i<eDep.size();i++){
        hitall+=eDep[i];
        r[0]+=eDep[i]*eDepx[i];
        r[1]+=eDep[i]*eDepy[i];
        r[2]+=eDep[i]*eDepz[i];
    }
    r[0]=r[0]/hitall;r[1]=r[1]/hitall;r[2]=r[2]/hitall;
    return r;//mm
}
""")
path='thedata.root'
对象名字='tree'
file1=ROOT.TFile.Open(path)#打开文件
RDF1=ROOT.RDataFrame(file1.Get(对象名字))#转为RDataFrame
RDF1=RDF1.Define("R","GeteDepDingdian(eDep,eDepx,eDepy,eDepz)")#定义R
RDF1=RDF1.Define("Z","R[2]").Define("X","R[0]")#定义其他的
RDF1=RDF1.Filter("Z>0")#筛选条件
RDF1.Snapshot('thedata_CUT.root','tree',{'X','Z'})#保存为root文件
DATA_py1=RDF1.AsNumpy(list(['X','Z']))#或者转为numpy数组
TH1=RDF1.Histo1D("Z","Z",100,-10,10)#或者也可以直接画直方图
print(DATA_py1)
print(TH1)
print(RDF1.Count().GetValue())#统计事例数
print(RDF1.Sum("Z").GetValue())#统计总Z
print(RDF1)

一些关于RDataFrame的用法可以参考官方文档
学习ROOT一般从官方文档,tutorial,碰到奇怪问题有论坛

---

1. https://root.cern/doc/v630/group__tutorial__tree.html ↩︎