CMS Wiki Pages

This page is read only. You can view the source, but not change it. Ask your administrator if you think this is wrong.
Below is an example on the use of ''TTree'' in ''pyROOT''. Run in the command line with
<code>
$ python trees.py
</code>
and check output with
<code>
$ root -l tree.root
root [1] TBrowser a
</code>

<code python trees.py>
#! /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)
    
</code>
Physik-Institut

CMS Wiki Pages

User Tools

Site Tools

Page Tools
