Below is an example on the use of ''TTree'' in ''pyROOT''. Run in the command line with
$ python trees.py
and check output with
$ root -l tree.root
root [1] TBrowser a
#! /usr/bin/env python
import time
start = time.time()
from ROOT import gRandom, TFile, TTree, TH1D
from array import array
import numpy # alternative to array
# Add branch to existing tree:
# https://root.cern.ch/phpBB3/viewtopic.php?t=9661
#
# Write directories into root files:
# https://root.cern.ch/root/html/tutorials/io/dirs.C.html
def writeTree():
print ">>> Writing a tree."
# make new root file with new tree
f = TFile("tree.root", 'recreate')
t = TTree("tree_name", "tree title")
# create 1 dimensional float arrays as fill variables, in this way the float array serves
# as a pointer which can be passed to the branch
n = array('d',[0])
u = array('d',[0])
# using numpy instead of array class (note python's float datatype corresponds to C++ doubles):
#n = numpy.zeros(1, dtype=float)
#u = numpy.zeros(1, dtype=float)
# create the branches and assign the fill-variables to them as doubles (D)
t.Branch("normal", n, 'normal/D')
t.Branch("uniform", u, 'uniform/D')
# create some random numbers, assign them into the fill variables and call Fill()
for i in xrange(10000):
n[0] = gRandom.Gaus()
u[0] = gRandom.Uniform()
t.Fill()
# write the tree into the output file and close the file
f.Write()
f.Close()
def addBranch():
print ">>> Add branch."
# reopen tree file
f = TFile("tree.root",'update')
t = f.Get("tree_name")
# make new variable
n = array('d',[0])
#n = numpy.zeros(1, dtype=float)
b = t.Branch("normal2", n, 'normal2/D')
# fill the branch
N = t.GetEntries()
for i in xrange(N):
t.GetEntry(i)
n[0] = gRandom.Gaus(3,2)
b.Fill()
# overwrite the tree in the output file and close the file
f.Write("",TFile.kOverwrite)
f.Close()
def addHist():
print ">>> Add histograms to a tree and one to a new branch."
# reopen tree file
f = TFile("tree.root",'update')
t = f.Get("tree_name")
# make new histogram
h1 = TH1D("normal3","Haters be histin'",50,-2,6)
# fill histogram
for i in xrange(10000):
h1.Fill(gRandom.Gaus(2,2))
f.Write("",TFile.kOverwrite)
f.Close()
def addHistToDirectory():
print ">>> Add histogram to directory."
# reopen tree file
f = TFile("tree.root",'update')
# make directory if it does not exist
dir = f.GetDirectory("dir1")
if not dir:
print ">>> created dir1"
dir = f.mkdir("dir1")
# make this directory the current one: everything you write will be saved here
dir.cd()
# declare histogram
h1 = TH1D("normal4","Hist in folder",50,-1,9)
# write histogram
for i in xrange(10000):
h1.Fill(gRandom.Gaus(4,3))
f.Write("",TFile.kOverwrite)
f.Close()
def addBranchToDirectory():
print ">>> Add branch to directory."
f = TFile("tree.root",'update')
# make directory if it does not exist
dir = f.GetDirectory("dir1")
if not dir:
print ">>> created dir1"
dir = f.mkdir("dir1")
# make this directory the current one: everything you write will be saved here
dir.cd()
# make new tree with a branch
t = TTree("tree2", "tree2")
n = array('d',[0])
#n = numpy.zeros(1, dtype=float)
t.Branch("normal5", n, 'normal5/D')
# fill tree
for i in xrange(10000):
n[0] = gRandom.Gaus(3,2)
t.Fill()
f.Write("",TFile.kOverwrite)
f.Close()
def loopThroughTree():
print ">>> Loop through tree."
f = TFile("tree.root",'read')
t = f.Get("tree_name")
# loop over tree in the easiest way
i = 0
for event in t:
if (i%2000) == 0:
print ">>> event %5i: n = %5.2f" % (i,t.normal)
print ">>> event %5i: u = %5.2f" % (i,t.uniform)
i += 1
print ">>> event %5i: n = %5.2f" % (i,t.normal)
print ">>> event %5i: u = %5.2f" % (i,t.uniform)
# alternatively
# N = t.GetEntries()
# for i in xrange(N):
# if (i%2000) == 0:
# t.GetEntry(i)
# print ">>> event %5i: n = %5.2f" % (i,t.normal)
# print ">>> event %5i: u = %5.2f" % (i,t.uniform)
# i += 1
# print ">>> event %5i: n = %5.2f" % (i,t.normal)
# print ">>> event %5i: u = %5.2f" % (i,t.uniform)
f.Close()
def main():
print "\n>>> Hello, world!"
writeTree()
addBranch()
addBranchToDirectory()
addHist()
addHistToDirectory()
loopThroughTree()
if __name__ == '__main__':
main()
end = time.time()
print ">>>\n>>> The program lasted %.3f seconds.\n" % (end-start)