This shows you the differences between two versions of the page.
Both sides previous revisionPrevious revisionNext revision | Previous revisionLast revisionBoth sides next revision | ||
root:pyroot [2017/06/13 08:59] – [Adding and saving a branch to a TTree in a TFile] iwn | root:pyroot [2017/09/05 15:05] – [Useful pyROOT snippets] iwn | ||
---|---|---|---|
Line 1: | Line 1: | ||
+ | [[root: | ||
+ | |||
====== Useful pyROOT snippets ====== | ====== Useful pyROOT snippets ====== | ||
- | This page lists some loose examples of '' | + | Here are some loose examples of '' |
- | + | | |
- | ===== Handling TTree trees ===== | + | |
- | + | | |
- | [[root:example_TTree|Here]] is one file containing multiple examples. | + | |
- | + | * [[https://wiki.physik.uzh.ch/lhcb/root:root|more by the LHCb group]]. | |
- | ==== 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 '' | + | |
- | + | ||
- | <file python tree.py> | + | |
- | from ROOT import TFile, TTree, gRandom | + | |
- | from array import array | + | |
- | + | ||
- | # make new root file with new tree | + | |
- | file = TFile(" | + | |
- | tree = TTree(" | + | |
- | + | ||
- | # 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(' | + | |
- | phi = array(' | + | |
- | + | ||
- | # create the branches | + | |
- | tree.Branch(" | + | |
- | tree.Branch(" | + | |
- | + | ||
- | # create some random numbers, assign them into the fill variables and call Fill() | + | |
- | for i in xrange(10000): | + | |
- | px[0] = gRandom.Gaus(20, | + | |
- | 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 '' | + | |
- | <code python> | + | |
- | import numpy | + | |
- | ... | + | |
- | px = numpy.zeros(1, | + | |
- | phi = numpy.zeros(1, | + | |
- | ... | + | |
- | </ | + | |
- | + | ||
- | Note that '' | + | |
- | <code python> | + | |
- | tree.Write() | + | |
- | </ | + | |
- | Optional arguments are a string for a new name and some writing options, e.g. to overwrite with '' | + | |
- | + | ||
- | ==== 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 '' | + | |
- | + | ||
- | <file python tree.py> | + | |
- | from ROOT import TFile, TTree, gRandom | + | |
- | from array import array | + | |
- | + | ||
- | file = TFile(" | + | |
- | tree = TTree(" | + | |
- | + | ||
- | Nmax = 10 | + | |
- | nParticles = array( ' | + | |
- | pt = array( ' | + | |
- | tree.Branch( ' | + | |
- | tree.Branch( ' | + | |
- | + | ||
- | for i in xrange(1000): | + | |
- | nParticles[0] = int(gRandom.Uniform()*10) | + | |
- | for j in xrange(nParticles[0]): | + | |
- | pt[j] = gRandom.Gaus(20, | + | |
- | tree.Fill() | + | |
- | + | ||
- | file.Write("", | + | |
- | file.Close() | + | |
- | </file> | + | |
- | ==== Looping over a TTree ==== | + | |
- | + | ||
- | The classic way of looping through a tree is: | + | |
- | <code python> | + | |
- | N = tree.GetEntries() | + | |
- | for i in xrange(N): | + | |
- | tree.GetEntry(i) | + | |
- | print tree.px | + | |
- | </code> | + | |
- | where you can access the values in branch '' | + | |
- | <code python> | + | |
- | for event in tree: | + | |
- | print tree.px | + | |
- | </code> | + | |
- | If you also need the index: | + | |
- | <code python> | + | |
- | for i, event in enumerate(tree): | + | |
- | print i, tree.px | + | |
- | </code> | + | |
- | + | ||
- | + | ||
- | ==== 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. | + | |
- | + | ||
- | <file python addBranch.py> | + | |
- | from ROOT import TFile, TTree, gRandom | + | |
- | + | ||
- | # reopen tree file | + | |
- | file = TFile(" | + | |
- | tree = file.Get(" | + | |
- | + | ||
- | # make new variable and add it as a branch to the tree | + | |
- | py = array(' | + | |
- | branch = tree.Branch(" | + | |
- | + | ||
- | # fill the branch | + | |
- | N = tree.GetEntries() | + | |
- | for i in xrange(N): | + | |
- | tree.GetEntry(i) | + | |
- | py[0] = gRandom.Gaus(23, | + | |
- | branch.Fill() | + | |
- | + | ||
- | # overwrite the tree in the output file and close the file | + | |
- | file.Write("", | + | |
- | file.Close() | + | |
- | </ | + | |
- | + | ||
- | More interestingly, | + | |
- | <code python> | + | |
- | import sqrt from math | + | |
- | ... | + | |
- | pt = array(' | + | |
- | branch = tree.Branch(" | + | |
- | + | ||
- | 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: | + | |
- | + | ||
- | <code python> | + | |
- | hist = TH1D(" | + | |
- | for i in xrange(10000): | + | |
- | hist.Fill(gRandom.Gauss()) | + | |
- | </ | + | |
- | + | ||
- | With some weight: | + | |
- | <code python> | + | |
- | weight = gRandom.Gauss(1, | + | |
- | hist.Fill(gRandom.Gauss(), | + | |
- | </ | + | |
- | ==== Filling a histogram from a tree ==== | + | |
- | + | ||
- | The classic way would be | + | |
- | <code python> | + | |
- | hist = TH1D(" | + | |
- | 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 [[https://root.cern.ch/doc/master/ | + | |
- | <code python> | + | |
- | hist = TH1D(" | + | |
- | tree.Draw(" | + | |
- | </ | + | |
- | Note that the histogram has been saved into '' | + | |
- | + | ||
- | You can also use mathematical expressions: | + | |
- | <code python> | + | |
- | tree.Draw(" | + | |
- | </ | + | |
- | + | ||
- | You can apply some selections using variable, available in the tree's branches: | + | |
- | <code python> | + | |
- | tree.Draw(" | + | |
- | </code> | + | |
- | + | ||
- | You can apply some selections and/or a weight, if a weight variable is available in the tree: | + | |
- | <code python> | + | |
- | tree.Draw(" | + | |
- | tree.Draw(" | + | |
- | </code> | + | |
- | + | ||
- | An equivalent method is creating the histogram in the '' | + | |
- | <code python> | + | |
- | from ROOT import gDirectory | + | |
- | ... | + | |
- | tree.Draw(" | + | |
- | hist = gDirectory.Get(" | + | |
- | </ | + | |
- | ===== 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 '' | + | |
- | + | ||
- | <file python addDirectory.py> | + | |
- | from ROOT import TFile | + | |
- | + | ||
- | file = TFile(" | + | |
- | + | ||
- | # make directory if it does not exist | + | |
- | dir = file.GetDirectory(" | + | |
- | if not dir: | + | |
- | print ">>> | + | |
- | dir = file.mkdir(" | + | |
- | + | ||
- | # make this directory the current one: everything you write will be saved here | + | |
- | dir.cd() | + | |
- | + | ||
- | # make tree, histogram, ... | + | |
- | + | ||
- | file.Write("", | + | |
- | file.Close() | + | |
- | </ | + | |
- | + | ||
- | + | ||
- | ==== gDirectory ==== | + | |
- | + | ||
- | Note that '' | + | |
- | + | ||
- | <file python gDirectory.py> | + | |
- | from ROOT import TFile, TTree, TH1F, gDirectory | + | |
- | + | ||
- | def printGDirectory(): | + | |
- | print ">>> | + | |
- | print ">>> | + | |
- | gDirectory.pwd() | + | |
- | print ">>> | + | |
- | gDirectory.ls() | + | |
- | + | ||
- | + | ||
- | print " | + | |
- | printGDirectory() | + | |
- | + | ||
- | print ">>> | + | |
- | file = TFile(" | + | |
- | tree = TTree(" | + | |
- | hist = TH1F(" | + | |
- | dir1 = file.mkdir(" | + | |
- | printGDirectory() | + | |
- | + | ||
- | print ">>> | + | |
- | gDirectory.Delete(" | + | |
- | printGDirectory() | + | |
- | + | ||
- | print ">>> | + | |
- | dir1.cd() | + | |
- | printGDirectory() | + | |
- | print ">>> | + | (More snippets to follow.) |
- | file.Close() | + | |
- | printGDirectory() | + | |
- | </ | + |