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
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");
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',) phi = array('d',) # 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 = gRandom.Gaus(20,2) phi = 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
import numpy ... px = numpy.zeros(1, dtype=float) phi = numpy.zeros(1, dtype=float) ...
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:
Optional arguments are a string for a new name and some writing options, e.g. to overwrite with
Here is an example of filling a tree with a variable number of particles, each having a
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 = int(gRandom.Uniform()*10) for j in xrange(nParticles): # 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',) branch = tree.Branch("py", py, 'py/D') # fill the branch N = tree.GetEntries() for i in xrange(N): tree.GetEntry(i) py = 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
import sqrt from math ... pt = array('d',) branch = tree.Branch("pt", pt, 'pt/D') N = tree.GetEntries() for i in xrange(N): tree.GetEntry(i) pt = sqrt(tree.px**2 + tree.py**2) branch.Fill()
Print the values of some branches in columns with
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.
hadd in the command line:
hadd sample.root sample_1.root sample_2.root sample_3.root
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: