User Tools

Site Tools


root:pyroot

This is an old revision of the document!


Useful pyROOT snippets

This page lists some loose examples of ROOT in python that may be useful.

Handling TTree trees

Here is one file containing multiple examples.

Creating a TTree with branches and writing it to a TFile

Below is an example of creating and saving a tree with two branches (variables) each containing a randomly distributed variable for each event.

If you would fill a TTree in C++, you can pass a double as a reference, called the address of the branch. In python however, you will need to use an array instead, for example the array class. More information on the Branch(name,address, leaflist) method.

tree.py
from ROOT import TFile, TTree, gRandom
from array import array
 
# make new root file with new tree
file = TFile("tree.root", 'recreate')
tree = 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
px  = array('d',[0])
phi = array('d',[0])
 
# create the branches and assign the fill-variables to them as doubles (D)
tree.Branch("px",  px,  'normal/D')
tree.Branch("phi", phi, 'uniform/D')
 
# create some random numbers, assign them into the fill variables and call Fill()
for i in xrange(10000):
    px[0]  = gRandom.Gaus(20,2)
    phi[0] = gRandom.Uniform(2*3.1416)
    tree.Fill()
 
# write the tree into the output file and close the file
file.Write()
file.Close()

Alternatively, you could use the numpy class:

import numpy
...
px  = numpy.zeros(1, dtype=float)
phi = numpy.zeros(1, dtype=float)
...

Note that file.Write() will write all objects (trees, histograms, graphs, …) that are still opened. Instead, you can write an object, in this case a TTree, with the Write() method:

tree.Write()

Optional arguments are a string for a new name and some writing options, e.g. to overwrite with TFile.kOverwrite.

Filling a branch with multiple values per event

Here is an example of filling a tree with a variable number of particles, each having a pt variable:

tree.py
from ROOT import TFile, TTree, gRandom
from array import array
 
file = TFile("tree.root", 'recreate')
tree = TTree("tree_name", "tree title")
 
Nmax = 10
nParticles = array( 'i', [ 0 ] )
pt = array( 'd', Nmax*[ 0. ] )
tree.Branch( 'nParticles', nParticles, 'nParticles/I' )
tree.Branch( 'pt', pt, 'pt[nParticles]/D' )
 
for i in xrange(1000):
    nParticles[0] = int(gRandom.Uniform()*10)
    for j in xrange(nParticles[0]):
       pt[j] = gRandom.Gaus(20,2)
    tree.Fill()
 
file.Write("",TFile.kOverwrite)
file.Close()

Looping over a TTree

The classic way of looping through a tree is:

N = tree.GetEntries()
for i in xrange(N):
    tree.GetEntry(i)
    print tree.px

where you can access the values in branch px for each event. A more elegant and pythonic way is the following:

for event in tree:
    print tree.px

If you also need the index:

for i, event in enumerate(tree):
    print i, tree.px

Adding and saving a branch to a TTree in a TFile

Here is simple example of adding a completely new branch to an existing tree.

addBranch.py
from ROOT import TFile, TTree, gRandom
 
# reopen tree file
file = TFile("tree.root",'update')
tree = file.Get("tree_name")
 
# make new variable and add it as a branch to the tree
py = array('d',[0])
branch = tree.Branch("py", py, 'py/D')
 
# fill the branch
N = tree.GetEntries()
for i in xrange(N):
    tree.GetEntry(i)
    py[0] = gRandom.Gaus(23,4)
    branch.Fill()
 
# overwrite the tree in the output file and close the file
file.Write("",TFile.kOverwrite)
file.Close()

More interestingly, you can access existing branches in the tree and use those values to create new variables. Here is an example of making a new branch containing transverse momentum for each event, assuming the tree has branches px and py:

import sqrt from math
...
pt = array('d',[0])
branch = tree.Branch("pt", pt, 'pt/D')
 
N = tree.GetEntries()
for i in xrange(N):
    tree.GetEntry(i)
    pt[0] = sqrt(tree.px**2 + tree.py**2)
    branch.Fill()

Handling TH1 histograms

Filling a histogram

Simple enough:

hist = TH1D("hist_name","hist title",100,0,100)
for i in xrange(10000):
    hist.Fill(gRandom.Gauss())

With some weight:

weight = gRandom.Gauss(1,0.1)
hist.Fill(gRandom.Gauss(),weight)

Filling a histogram from a tree

The classic way would be

hist = TH1D("hist_name","hist title",100,0,100)
for event in tree:
    hist.Fill(tree.px)

But the simplest and most powerful way to draw some variable into a histogram, is with the Draw function:

hist = TH1D("hist_name","hist title",100,0,100)
tree.Draw("px >> hist_name")

Note that the histogram has been saved into ROOT's working memory and is therefore accessible via its name in the Draw function, “hist_name” in this example.

You can also use mathematical expressions:

tree.Draw("2*sqrt(px*px+py*py) >> hist_name")

You can apply some selections using variable, available in the tree's branches:

tree.Draw("px >> hist_name","pt>20 && E>20")

You can apply some selections and/or a weight, if a weight variable is available in the tree:

tree.Draw("px >> hist_name","(pt>20 && E>20)*weight")
tree.Draw("px >> hist_name","weight")

An equivalent method is creating the histogram in the Draw method. If you need the TH1F object, you can still get it with gDirectory (see below):

from ROOT import gDirectory
...
tree.Draw("pt >> h(100,0,100)")
hist = gDirectory.Get("h") # will be a TH1F object

It's also possible to pass some option you would normally use in “TH1::Draw()”. It is passed as the third string.

tree.Draw("px >> h1(100,0,100)","","E2")
tree.Draw("py >> h2(100,0,100)","","E2 SAME")

Protip: This is very useful for quick drawing and comparing, for example in the ROOT command line.

$ root -l tree.root
[1] tree->Draw("px >> h1(100,0,100)","","E2")
[2] tree->Draw("py >> h2(100,0,100)","","E2 SAME")

Handling TFiles

Adding a directory to a TFile

Here is a basic example creating a file with on directory. If you want to create and save an object in this directory, e.g. a tree or histogram, use dir.cd(). Pay attention to the order of creating an object, changing directory (cd()) and writing the object!

addDirectory.py
from ROOT import TFile
 
file = TFile("tree.root",'update')
 
# make directory if it does not exist
dir = file.GetDirectory("dir")
if not dir:
    print ">>> created dir"
    dir = file.mkdir("dir")
 
# make this directory the current one: everything you write will be saved here
dir.cd()
 
# make tree, histogram, ...
 
file.Write("",TFile.kOverwrite)
file.Close()

gDirectory

Note that gDirectory is the current directory containing all newly created saved objects, or those saved to this directory. If you just opened a file, this will be your current directory until you close it or use cd(). Test the behavior:

gDirectory.py
from ROOT import TFile, TTree, TH1F, gDirectory
 
def printGDirectory():
    print ">>> gDirectory.GetName()\n%s" % gDirectory.GetName()
    print ">>> gDirectory.pwd()"
    gDirectory.pwd()
    print ">>> gDirectory.ls()"
    gDirectory.ls()
    print
 
print "\ndefault"
printGDirectory()
 
print ">>> creating some file with some contents"
file = TFile("test.root","recreate")
tree = TTree("tree","tree")
hist = TH1F("hist","hist",100,0,100)
dir1 = file.mkdir("dir1")
printGDirectory()
 
print ">>> gDirectory.Delete(\"hist\")"
gDirectory.Delete("hist")
printGDirectory()
 
print ">>> dir1.cd()"
dir1.cd()
printGDirectory()
 
print ">>> file.Close()"
file.Close()
printGDirectory()
root/pyroot.1497338367.txt.gz · Last modified: 2017/06/13 09:19 by iwn