PyROOT——tree文件的读写

学习PyROOT代码可以参考

ROOT: PyRoot tutorials

root里面常用的TGraph TH1 TH2 Fit都有python的写法

import numpy as np
import ROOT as root
from ROOT import TCanvas, TGraph, TF1

#import root file
runnumber = sys.argv[1]
print(runnumber)
file1 = "/home/xxx/"
file2 = "_tree.root"
filename = "%s%s%s"%(file1,runnumber,file2)
print(filename)
root_file = root.TFile(filename)
#import tree
tree = root_file.Get("tree")
n_entries = tree.GetEntries()

#define file_out
file_out="_file_out.root"
filen_outname = "%s%s"%(runnumber,file_out)
new_root_file = root.TFile(filen_outname, "RECREATE")
new_tree = root.TTree("tree", "Modified Tree")

#define branch
E= np.zeros(1, dtype=np.int64)
T = np.zeros(1, dtype=np.int64)
length= np.zeros(1, dtype=np.int16)
data = np.zeros(500, dtype=np.int64)#长度500的数组
aa = np.zeros(1, dtype=np.int64)
bb = np.zeros(1, dtype=np.int64)

tree.SetBranchAddress('E', E)
tree.SetBranchAddress('T', T)
tree.SetBranchAddress('length', length)
tree.SetBranchAddress('data', data)

#数据类型根据root文件MakeClass生成的.h文件里面的branch写
new_tree.Branch("E", E, "E/s")
new_tree.Branch("T", T, "T/L")
new_tree.Branch("length", length, "length/s")
new_tree.Branch("data", data, "data[length]/s")
#aa bb是新写入的branch
new_tree.Branch("aa", aa, "aa/s")
new_tree.Branch("bb", bb, "bb/s")

for i in range(n_entries):#n_entries = tree.GetEntries()
    tree.GetEntry(i)
    if i%1000==0:
        #print('\r i:{}'.format(i),end = '')
        print(i)
    #aa=
    #bb=
    new_tree.Fill()


new_root_file.Write()
root_file.Close()
new_root_file.Close()