rootpy is a more pythonic interface for the ROOT libraries on top of the PyROOT bindings.
See the docs here: rootpy.org
Get the latest code on GitHub: github.com/rootpy
git clone https://github.com/rootpy/rootpy.git
cd rootpy; python setup.py install --user
Aside from making ROOT better in Python, rootpy allows you to integrate ROOT with other popular scientific packages in the Python universe.
The root_numpy package provides an interface between ROOT and NumPy. See the docs here: rootpy.github.io/root_numpy
Let's see how rootpy can make your life easier with some examples...
# histograms are unified!
from rootpy.plotting import Hist
hist = Hist(10, -3, 3) # <== notice that the name or title is not needed
hist.name # the name is automatically a unique ID
'7982e742cdd94b34b741818a8113c833'
# Hist is really a subclass of the appropriate ROOT TH1X class
hist.__class__.__bases__
(rootpy.plotting.hist._Hist, ROOT.TH1F)
# the base class is determined dynamically at runtime
hist = Hist(10, -3, 3, type='D', title="asd", linecolor="#ccddee")
hist.__class__.__bases__
(rootpy.plotting.hist._Hist, ROOT.TH1D)
print hist.GetLineColor()
# rootpy has ipython notebook support so things can be drawn inline
hist.FillRandom('gaus')
hist
#ccddee
The arguments that specify the binning are consumed from left to right as either three values for the number of bins, left bound, and right bound, or a list for bins of variable width. Fixed and variable width bins may be mixed for 2D and 3D histograms:
from rootpy.plotting import Hist2D, Hist3D
# variable-width bins
h1d = Hist([-10, -3.2, 5.2, 6])
# 2D histogram with fixed-width bins; 10 along x and 5 along y
h2d = Hist2D(10, 0, 1, 5, -2, 4, drawstyle='E')
# 3 variable-width bins along x and 4 fixed-width bins along y
h2d = Hist2D([10, 30, 100, 1000], 4, 10, 33.5)
# 3D histogram with fixed-width bins
h3d = Hist3D(100, 0, 1, 20, 5.8, 7.2, 1000, -10, 1)
# fixed-width bins along x and z and variable-width bins along y
h3d = Hist3D(5, 0, 1, [-20, 5, 50, 1e3], 10, 0, 1)
Style settings are much friendlier in rootpy (no need to remember the ROOT colour codes!). rootpy also has the official styles for ATLAS, CMS and LHCb built-in.
from rootpy.plotting import set_style
set_style('ATLAS')
hist = Hist(50, -3, 3)
hist.FillRandom('gaus')
hist.xaxis.title = 'Variable [Units]'
hist.yaxis.title = 'Entries'
hist.linewidth = 3
hist.linestyle = 'dashed'
hist.linecolor = 'forestgreen' # SVG
hist.markercolor = '#334455' # hex
hist.fillstyle = '/'
hist.fillcolor = 'red'
hist.drawstyle = 'HIST'
hist
INFO:rootpy.plotting.style:using ROOT style 'ATLAS'
# access bin contents and errors
hist[1] = (3, 1.2)
hist[1].value
3.0
hist[1].error
1.2
# reverse bin contents (and errors)
hist[:] = hist.Clone(shallow=True)[::-1]
hist
# rebin by 2 along x
hist.rebinned(5, axis=0)
Hist(hist[::10]) # you can construct histograms from "views" of other histograms. The step value in the slice creates a rebinning.
# merge specific bins
hist.merge_bins([(0, 5), (40, -1)])
# same functionality in 2D and 3D...
from rootpy.plotting import F2
h2 = Hist2D(20, 0, 1, 20, 0, 1)
h2.FillRandom(F2('x*y'))
h2.drawstyle = 'LEGO2'
h2
h2_merged = h2.merge_bins([(6, 12)], axis=1)
h2_merged.drawstyle = 'LEGO2'
h2_merged
h2_merged = h2_merged.merge_bins([(6, 12)], axis=0)
h2_merged.drawstyle = 'LEGO2'
h2_merged
Hist2D(h2[5:10,::5], drawstyle='LEGO2') # crop along x and rebin along y
h2[:,2:8] = (0, 0) # set specific ranges of bins (value, error)
h2
# flatten a histogram by rebinning according to the quantiles
h = Hist(1000, -3, 3)
h.FillRandom('gaus', 100000)
print h.quantiles(5, strict=True)
h_flat = h.rebinned(h.quantiles(5, strict=True))
h_flat.SetMinimum(0)
h_flat
[-3. -0.666 0. 0.672 2.994]
from rootpy.io import root_open
from rootpy.testdata import get_filepath
# objects are automatically cast to the rootpy subclasses if they exist
# support for the with-statement
with root_open(get_filepath()) as f:
print f
for path, dirs, objects in f.walk():
print path, objects
File('/home/endw/.local/lib/python2.7/site-packages/rootpy-dev-py2.7.egg/rootpy/testdata/test_file.root') [] dimensions ['hist2d', 'hist3d'] scales ['hist1', 'hist3', 'hist2', 'hist4'] means ['hist1', 'hist3', 'hist2', 'hist4'] graphs ['tgrapherrors', 'tgraph2d', 'tgraphasymmerrors', 'tgraph'] gaps ['hist1', 'hist3', 'hist2', 'hist4'] efficiencies ['hist1', 'hist3', 'hist2', 'hist4', 'eff3v1', 'eff2v1', 'eff4v1']
f = root_open(get_filepath())
list(f.find('hist1'))
[('/scales/hist1', <_sre.SRE_Match at 0x54d6850>), ('/means/hist1', <_sre.SRE_Match at 0x54d6780>), ('/gaps/hist1', <_sre.SRE_Match at 0x54d68b8>), ('/efficiencies/hist1', <_sre.SRE_Match at 0x54d6920>)]
# only loop over objects of a given class
for path, dirs, objects in f.walk(class_pattern='TH1F'):
print objects
[] [] ['hist1', 'hist3', 'hist2', 'hist4'] ['hist1', 'hist3', 'hist2', 'hist4'] [] ['hist1', 'hist3', 'hist2', 'hist4'] ['hist1', 'hist3', 'hist2', 'hist4']
f.scales.hist1 # easy acess
# does something exist in a file?
'scales' in f
True
'blah' in f
False
'scales/hist1' in f
True
# create a new file
g = root_open('test.root', 'recreate')
hist = Hist(10, -3, 3)
hist.FillRandom('gaus')
# write to this file
g.myhist = hist
list(g)
[Hist('d02e8d09a8214f68be2c52bd3ae2e1e6')]
g.Get('myhist')
# ROOT errors become real Python exceptions
root_open('does_not_exist.root')
ERROR:ROOT.TFile.TFile:file does_not_exist.root does not exist
--------------------------------------------------------------------------- ROOTError Traceback (most recent call last) <ipython-input-29-24d63a85d9b6> in <module>() 1 # ROOT errors become real Python exceptions ----> 2 root_open('does_not_exist.root') /home/endw/.local/lib/python2.7/site-packages/rootpy-dev-py2.7.egg/rootpy/io/file.pyc in root_open(filename, mode) 125 filename = expand_path(filename) 126 prev_dir = ROOT.gDirectory.func() --> 127 root_file = ROOT.R.TFile.Open(filename, mode) 128 if not root_file: 129 raise IOError("could not open file: '{0}'".format(filename)) /home/endw/.local/lib/python2.7/site-packages/rootpy-dev-py2.7.egg/rootpy/io/file.pyc in root_open(filename, mode) 125 filename = expand_path(filename) 126 prev_dir = ROOT.gDirectory.func() --> 127 root_file = ROOT.R.TFile.Open(filename, mode) 128 if not root_file: 129 raise IOError("could not open file: '{0}'".format(filename)) /home/endw/.local/lib/python2.7/site-packages/rootpy-dev-py2.7.egg/rootpy/logger/roothandler.pyc in python_logging_error_handler(level, root_says_abort, location, msg) 100 # Hence the need for dark magic, we re-raise it within a trace. 101 from .. import ROOTError --> 102 raise ROOTError(level, location, msg) 103 except RuntimeError: 104 _, exc, traceback = sys.exc_info() ROOTError: level=3000, loc='TFile::TFile', msg='file does_not_exist.root does not exist'
g['asasd']
--------------------------------------------------------------------------- DoesNotExist Traceback (most recent call last) <ipython-input-30-8c57f6b0c853> in <module>() ----> 1 g['asasd'] /home/endw/.local/lib/python2.7/site-packages/rootpy-dev-py2.7.egg/rootpy/io/file.pyc in __getitem__(self, name) 195 def __getitem__(self, name): 196 --> 197 return self.Get(name) 198 199 def __setitem__(self, name, thing): /home/endw/.local/lib/python2.7/site-packages/rootpy-dev-py2.7.egg/rootpy/io/file.pyc in get(self, name, *args, **kwargs) 92 raise DoesNotExist( 93 "requested path '{0}' does not exist in {1}".format( ---> 94 name, self._path)) 95 return get 96 DoesNotExist: requested path 'asasd' does not exist in test.root
from rootpy.tree import Tree, TreeModel, IntCol, FloatCol
from rootpy.vector import LorentzVector
from random import gauss, uniform, expovariate
from math import pi
class Event(TreeModel):
number = IntCol()
MET_phi = FloatCol()
MET = FloatCol(default=-1)
jet = LorentzVector
output = root_open('test.root', 'recreate')
tree = Tree('data', model=Event)
for i in xrange(1000):
tree.number = i
tree.MET_phi = uniform(-pi, pi)
if i % 2 == 1:
tree.MET = expovariate(10)
tree.jet.SetPtEtaPhiM(expovariate(10), gauss(0, 2), uniform(-pi, pi), 0)
tree.Fill(reset=True)
len(tree)
1000
tree.Draw('MET_phi')
tree.Draw('jet.Pt()', drawstyle='hist', linewidth=3, linestyle='dotted')
# a more complicated model...
class Particle(TreeModel):
charge = IntCol()
fourvector = LorentzVector
class Jet(TreeModel):
num_tracks = IntCol()
vertex_fraction = FloatCol()
fourvector = LorentzVector
class Event(Particle.prefix('mu1_'), Particle.prefix('mu2_'),
Jet.prefix('jet1_'), Jet.prefix('jet2_')):
number = IntCol()
MET_phi = FloatCol()
MET = FloatCol(default=-1)
tree2 = Tree('events', model=Event)
tree2._buffer
jet1_fourvector -> LorentzVector(px=0.000000, py=0.000000, pz=0.000000, E=0.000000) jet2_fourvector -> LorentzVector(px=0.000000, py=0.000000, pz=0.000000, E=0.000000) mu1_fourvector -> LorentzVector(px=0.000000, py=0.000000, pz=0.000000, E=0.000000) mu2_fourvector -> LorentzVector(px=0.000000, py=0.000000, pz=0.000000, E=0.000000) mu1_charge -> Int(0) at 0x69cb8f0 mu2_charge -> Int(0) at 0x69cb890 jet1_num_tracks -> Int(0) at 0x69cb950 jet2_num_tracks -> Int(0) at 0x69cb9b0 jet1_vertex_fraction -> Float(0.0) at 0x69cba10 jet2_vertex_fraction -> Float(0.0) at 0x69cba70 number -> Int(0) at 0x69cbad0 MET_phi -> Float(0.0) at 0x69cbb30 MET -> Float(-1.0) at 0x69cbb90
from rootpy.tree import Cut
a = Cut()
a.__class__.__bases__
(ROOT.TCut,)
cut1 = Cut('a < 10')
cut2 = Cut('b % 2 == 0')
cut1 & cut2
'(a<10)&&(b%2==0)'
# expansion of ternary conditions
cut3 = Cut('10 < a < 20')
cut3
'(10<a)&&(a<20)'
# easily combine cuts arbitrarily
((cut1 & cut2) | - cut3)
'((a<10)&&(b%2==0))||(!((10<a)&&(a<20)))'