User Tools

Site Tools


root:pyroot_ttree

Getting started with ROOTUseful pyROOT snippetsHandling TTree trees

Handling TTree trees

Here is the class reference for TTree.

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 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.

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 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.

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:

tree2.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): # 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()

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 event.px

If you also need the event's index:

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

Making a histogram with TTree

See this page.

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()

Scanning the values of a TTree

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")

TChain

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”)

root/pyroot_ttree.txt · Last modified: 2018/02/28 14:03 by iwn