Getting started with ROOT → Useful pyROOT snippets → Handling TTree trees
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 fill a TTree
in C++
, you can pass a double as a reference, called the address of the branch, from the official tree1.C example:
TTree tree("tree_name","tree title"); Float_t px; tree.Branch("px",&px,"px/F");
In python
however, you will need to use an array instead, for example the array
class. More information on the TTree::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 a individual 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): # loop over events nParticles[0] = int(gRandom.Uniform()*10) for j in xrange(nParticles[0]): # loop over particles in this event 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 event.px
If you also need the event's index:
for i, event in enumerate(tree): print i, event.px
See this page.
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()
Print the values of some branches in columns with TTree::Scan
:
tree.Scan("run:lumi:event") tree.Scan("run:lumi:event","pt1>30 && pt2>30") # with some selections tree.Scan("run:lumi:event","pt1>30 && pt2>30","colsize=8") # with some option
The documentation of TTreePlayer::Scan() has more information on print options (column size, decimal precision, …) and formatting.
To write the output directly to a file (which gets overwritten every time):
tree.GetPlayer().SetScanRedirect(True) tree.GetPlayer().SetScanFileName("runnumber.txt") tree.Scan("run:lum:evt","pt1>30 && pt2>30")
If you have a samples split into many files, each containing the same tree, you can either add the files into one big one in the command line with hadd
or load them in a TChain
in you analysis code.
With hadd
in the command line:
hadd sample.root sample_1.root sample_2.root sample_3.root
With TChain
:
chain = TChain("tree_name") chain.Add("sample_1.root") chain.Add("sample_2.root") chain.Add("sample_3.root") for event in chain: print event.px
Looping over the events in a chain is similar as for trees. Note it's also possible to use a glob wildcard: chain.Add(“sample_*.root”)