UZH

Physik-Institut

CMS Wiki Pages

User Tools

Site Tools

Trace:
root:pyroot

Differences

This shows you the differences between two versions of the page.

Link to this comparison view

Both sides previous revisionPrevious revision
Next revision		Previous revision
root:pyroot [2017/06/13 09:23] – [Adding and saving a branch to a TTree in a TFile] iwnroot:pyroot [2017/09/05 15:05] (current) – [Useful pyROOT snippets] iwn
Line 1: Line 1:
 +[[root:root|Getting started with ROOT]] -> [[root:pyroot|Useful pyROOT snippets]]
 +
 ====== Useful pyROOT snippets ====== ====== Useful pyROOT snippets ======
  
-This page lists some loose examples of ''ROOT'' in ''python'' that may be useful. +Here are some loose examples of ''ROOT'' in ''python'' that may be useful: 
- +  * [[root:pyroot_TTree|handling TTree trees]], 
-===== Handling TTree trees ===== +  * [[root:pyroot_TH1|handling TH1 histograms]], 
- +  * [[root:pyroot_TFile|handling TFile files and TDirectory directories]], 
-[[root:example_TTree|Here]] is one file containing multiple examples. +  * [[https://github.com/clelange/roofit|RooFit]] and [[https://github.com/clelange/roostats|RooStats]] tutorials, [[https://www.nikhef.nl/~vcroft/GettingStartedWithRooFit.html|RooFit manual]]. 
- +  * [[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 ''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 [[https://root.cern/doc/master/classTTree.html#ac1fa9466ce018d4aa739b357f981c615|Branch(name,address, leaflist) method]]. +
- +
-<file python 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() +
-</file> +
- +
-Alternatively, you could use the ''numpy'' class: +
-<code python> +
-import numpy +
-... +
-px  = numpy.zeros(1, dtype=float) +
-phi = numpy.zeros(1, dtype=float) +
-... +
-</code> +
- +
-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 [[https://root.cern.ch/doc/master/classTTree.html#af6f2d9ae4048ad85fcae5d2afa05100f|Write() method]]: +
-<code python> +
-tree.Write() +
-</code> +
-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: +
- +
-<file python tree.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): +
-    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() +
-</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 ''px'' for each event. A more elegant and pythonic way is the following: +
-<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.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() +
-</file> +
- +
-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'': +
-<code python> +
-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() +
-</code> +
- +
- +
----- +
- +
-===== Handling TH1 histograms ===== +
- +
-==== Filling a histogram ==== +
- +
-Simple enough: +
- +
-<code python> +
-hist = TH1D("hist_name","hist title",100,0,100) +
-for i in xrange(10000): +
-    hist.Fill(gRandom.Gauss()) +
-</code> +
- +
-With some weight: +
-<code python> +
-weight = gRandom.Gauss(1,0.1) +
-hist.Fill(gRandom.Gauss(),weight) +
-</code> +
-==== Filling a histogram from a tree ==== +
- +
-The classic way would be +
-<code python> +
-hist = TH1D("hist_name","hist title",100,0,100) +
-for event in tree: +
-    hist.Fill(tree.px) +
-</code> +
- +
-But the simplest and most powerful way to draw some variable into a histogram, is with the [[https://root.cern.ch/doc/master/classTTree.html#a73450649dc6e54b5b94516c468523e45|Draw]] function: +
-<code python> +
-hist = TH1D("hist_name","hist title",100,0,100) +
-tree.Draw("px >> hist_name") +
-</code> +
-Note that the histogram has been saved into ''ROOT'''s working memory and is therefore accessible via its name in the ''Draw'' function, ''"hist_name"'' in this example. +
- +
-You can also use mathematical expressions: +
-<code python> +
-tree.Draw("2*sqrt(px*px+py*py) >> hist_name") +
-tree.Draw("abs(eta) >> hist_name") +
-</code> +
- +
-You can apply some selections using variable, available in the tree's branches: +
-<code python> +
-tree.Draw("px >> hist_name","pt>20 && E>20") +
-</code> +
- +
-You can apply some selections and/or a weight, if a weight variable is available in the tree: +
-<code python> +
-tree.Draw("px >> hist_name","(pt>20 && E>20)*weight") +
-tree.Draw("px >> hist_name","weight") +
-</code> +
- +
-An equivalent method is creating the histogram in the ''Draw'' method. If you need the ''TH1F'' object, you can still get it with ''gDirectory'' (see [[root:pyroot#gDirectory|below]]): +
-<code python> +
-from ROOT import gDirectory +
-... +
-tree.Draw("pt >> h(100,0,100)") +
-hist = gDirectory.Get("h") # will be a TH1F object +
-</code> +
- +
-It's also possible to pass some option you would normally use in "TH1::Draw()". It is passed as the third string. +
-<code python> +
-tree.Draw("px >> h1(100,0,100)","","E2") +
-tree.Draw("py >> h2(100,0,100)","","E2 SAME") +
- +
-tree.Draw("px >> h1(100,0,100)","abs(eta)>2","E2") +
-tree.Draw("px >> h2(100,0,100)","abs(eta)<2","E2 SAME") +
-</code> +
- +
-__Protip__: This is very useful for quick drawing and comparing, for example in the ''ROOT'' command line. +
-<code C++> +
-$ root -l tree.root +
-[1] tree->Draw("px >> h1(100,0,100)","","E2") +
-[2] tree->Draw("py >> h2(100,0,100)","","E2 SAME") +
-</code> +
-===== 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 ''dir.cd()''. Pay attention to the order of creating an object, changing directory (''cd()'') and writing the object! +
- +
-<file python addDirectory.py> +
-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() +
-</file> +
- +
- +
-==== gDirectory ==== +
- +
-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: +
- +
-<file python gDirectory.py> +
-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()" +(More snippets to follow.)
-file.Close() +
-printGDirectory() +
-</file>+
root/pyroot.1497338589.txt.gz · Last modified: 2017/06/13 09:23 by iwn

Page Tools

Except where otherwise noted, content on this wiki is licensed under the following license: CC Attribution-Share Alike 4.0 International
CC Attribution-Share Alike 4.0 International Donate Powered by PHP Valid HTML5 Valid CSS Driven by DokuWiki