This is an old revision of the document!
This page lists some loose examples of ROOT
in python
that may be useful.
Here is one file containing multiple examples.
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.
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
.
Here is an example of filling a tree with a variable number of particles, each having a pt
variable:
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()
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
Here is simple example of adding a completely new branch to an existing tree.
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()
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)
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 can be accessed via its name, “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, and getting the TH1F
object with gDirectory
(see below):
from ROOT import gDirectory ... tree.Draw("pt >> h(100,0,100)") hist = gDirectory.Get("h") # will be a TH1F object
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!
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()
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:
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()