def MultiDraw(self, Formulae, CommonWeight="1"):
"""Draws many histograms in one loop over a tree.
Instead of:
MyTree.Draw( "nlcts >> a(100, -1, 1)", "weightA" )
MyTree.Draw( "nlcts >> b(100, -1, 1)", "weightB" )
Do:
MyTree.MultiDraw( ( "nlcts >> a(100, -1, 1)", "weightA" ),
( "nlcts >> b(100, -1, 1)", "weightB" ) )
This is significantly faster when there are many histograms to be drawn.
The first parameter, CommonWeight, decides a weight given to all
histograms.
An arbitrary number of additional histograms may be specified. They can
either be specified with just a string containing the formula to be
drawn, the histogram name and bin configuration.
Alternatively it can be a tuple, with said string, and an additional
string specifying the weight to be applied to that histogram only.
"""
if type(CommonWeight) == tuple:
Formulae = (CommonWeight, ) + Formulae
CommonWeight = "1"
results, formulae, weights = [], [], []
lastFormula, lastWeight = None, None
# A weight common to everything being drawn
CommonWeightFormula = TTreeFormula("CommonWeight", CommonWeight, self)
CommonWeightFormula.SetQuickLoad(True)
if not CommonWeightFormula.GetTree():
raise RuntimeError("TTreeFormula didn't compile: " + CommonWeight)
hists = {}
for i, origFormula in enumerate(Formulae):
print "Have an origFormula", origFormula
# Expand out origFormula and weight, otherwise just use weight of 1.
if type(origFormula) == tuple:
origFormula, weight = origFormula
else:
origFormula, weight = origFormula, "1"
# print origFormula, weight
# Pluck out histogram name and arguments
match = re.match(r"^(.*?)\s*>>\s*(.*?)\s*\(\s*(.*?)\s*\)$", origFormula)
if match:
formula, name, arguments = match.groups()
arguments = re.split(",\s*", arguments)
bins, minX, maxX = arguments
bins, minX, maxX = int(bins), float(minX), float(maxX)
# Create histogram with name and arguments
hist = TH1D(name, name, bins, minX, maxX)
hist.Sumw2()
else:
# without arguments
match = re.match(r"^(.*?)\s*>>\s*(.*?)\s*$", origFormula)
if not match:
raise RuntimeError("MultiDraw: Couldn't parse formula: '%s'" % origFormula)
formula, name = match.groups()
# print formula, name
if name.startswith("+") and name[1:] in hists:
# Drawing additionally into a histogram
hist = hists[name[1:]]
else:
# name = name[1:] # JAN: ???
hist = gDirectory.Get(name)
if not hist:
raise RuntimeError("MultiDraw: Couldn't find histogram to fill '%s' in current directory." % name)
if name not in hists:
hists[name] = hist
results.append(hist)
# The following two 'if' clauses check that the next formula is different
# to the previous one. If it is not, we add an ordinary TObject.
# Then, the dynamic cast in MultiDraw.cxx fails, giving 'NULL', and
# The previous value is used. This saves the recomputing of identical values
if formula != lastFormula:
f = TTreeFormula("formula%i" % i, formula, self)
if not f.GetTree():
raise RuntimeError("TTreeFormula didn't compile: " + formula)
f.SetQuickLoad(True)
formulae.append(f)
else:
formulae.append(TObject())
if weight != lastWeight:
f = TTreeFormula("weight%i" % i, weight, self)
if not f.GetTree():
raise RuntimeError("TTreeFormula didn't compile: " + formula)
f.SetQuickLoad(True)
weights.append(f)
else:
weights.append(TObject())
lastFormula, lastWeight = formula, weight
# Only compile MultiDraw once
try:
from ROOT import MultiDraw as _MultiDraw
except ImportError:
# gROOT.ProcessLine(".L %sMultiDraw.cxx+O" % "./")
if "/sMultiDraw_cc.so" not in gSystem.GetLibraries():
gROOT.ProcessLine(".L %s/../SFrameAnalysis_emu/datacard/MultiDraw.cc+" % os.environ['CMSSW_BASE']);
from ROOT import MultiDraw as _MultiDraw
from time import time
start = time()
# Ensure that formulae are told when tree changes
fManager = TTreeFormulaManager()
for formula in formulae + weights + [CommonWeightFormula, ]:
if type(formula) == TTreeFormula:
fManager.Add(formula)
fManager.Sync()
self.SetNotify(fManager)
# Draw everything!
_MultiDraw(self, CommonWeightFormula,
MakeTObjArray(formulae),
MakeTObjArray(weights),
MakeTObjArray(results),
len(Formulae))
print "Took %.2fs" % (time() - start), " "*20
return results
def MultiDraw(self, varexps, selection='1', drawoption="", **kwargs):
"""Draws multiple histograms in one loop over a tree (self).
Instead of:
tree.Draw( "pt_1 >> a(100, 0, 100)", "weightA" )
tree.Draw( "pt_2 >> b(100, 0, 100)", "weightB" )
Do:
tree.MultiDraw( ( "pt_1 >> a(100, 0, 100)", "weightA" ),
( "pt_2 >> b(100, 0, 100)", "weightB" ) )
This is significantly faster when there are many histograms to be drawn.
The first parameter, commonWeight, decides a weight given to all histograms.
An arbitrary number of additional histograms may be specified. They can
either be specified with just a string containing the formula to be
drawn, the histogram name and bin configuration.
Alternatively it can be a tuple, with said string, and an additional
string specifying the weight to be applied to that histogram only."""
selection = kwargs.get('cut', selection) # selections cuts
verbosity = kwargs.get('verbosity', 0) # verbosity
poisson = kwargs.get('poisson', False) # kPoisson errors for data
sumw2 = kwargs.get('sumw2', False) # sumw2 for MC
histlist = kwargs.get('hists',
[]) # to not rely on gDirectory.Get(histname)
hists = {}
results, xformulae, yformulae, weights = [], [], [], []
lastXVar, lastYVar, lastWeight = None, None, None
# A weight common to everything being drawn
commonFormula = TTreeFormula("commonFormula", selection, self)
commonFormula.SetQuickLoad(True)
if not commonFormula.GetTree():
raise error(
"MultiDraw: TTreeFormula 'selection' did not compile:\n selection: %r\n varexps: %s"
% (selection, varexps))
for i, varexp in enumerate(varexps):
#print ' Variable expression: %s'%(varexp,)
yvar = None
# EXPAND varexp
weight = None
if isinstance(varexp, (tuple, list)) and len(varexp) == 2:
varexp, weight = varexp
elif not isinstance(varexp, str):
raise IOError(
"MultiDraw: given varexp is not a string or tuple of length 2! Got varexp=%s (%s)"
% (varexp, type(varexp)))
if not varexp: varexp = '1'
if not weight: weight = '1'
# PREPARE histogram
match = varregex.match(varexp)
if match: # create new histogram: varexp = "x >> h(100,0,100)" or "y:x >> h(100,0,100,100,0,100)"
xvar, name, binning = match.group(1), match.group(2), match.group(
3)
# CREATE HISTOGRAM
vmatch = varregex2D.match(xvar)
if not vmatch or xvar.replace('::', '').count(':') == xvar.count(
'?'): # 1D, allow "(x>100 ? 1 : 0) >> h(2,0,2)"
bmatch = binregex.match(binning)
if not bmatch:
raise error(
"MultiDraw: Could not parse formula for %r: %r" %
(name, varexp))
nxbins, xmin, xmax = int(bmatch.group(1)), float(
bmatch.group(2)), float(bmatch.group(3))
hist = TH1D(name, name, nxbins, xmin, xmax)
elif vmatch: # 2D histogram
yvar, xvar = vmatch.group(1), vmatch.group(2)
bmatch = binregex2D.match(binning)
if not bmatch:
raise error(
'MultiDraw: Could not parse formula for %r to pattern %r: "%s"'
% (name, binregex2D.pattern, varexp))
nxbins, xmin, xmax = int(bmatch.group(1)), float(
bmatch.group(2)), float(bmatch.group(3))
nybins, ymin, ymax = int(bmatch.group(4)), float(
bmatch.group(5)), float(bmatch.group(6))
hist = TH2D(name, name, nxbins, xmin, xmax, nybins, ymin, ymax)
else: # impossible
raise error(
'MultiDraw: Could not parse variable %r for %r to pattern %r: %r'
% (xvar, name, varregex2D.pattern, varexp))
else: # get existing histogram: varexp = "x >> h" or "y:x >> h"
match = varregex2.match(varexp)
if not match:
raise error(
'MultiDraw: Could not parse formula to pattern %r: %r' %
(varregex2.pattern, varexp))
xvar, name = match.groups()
if name.startswith("+") and name[1:] in hists:
hist = hists[name[1:]] # add content to existing histogram
else:
if i < len(histlist):
hist = histlist[i]
if hist.GetName() != histlist[i].GetName():
raise error(
"MultiDraw: Hisogram mismatch: looking for %r, but found %r."
% (hist.GetName(), histlist[i].GetName()))
else:
hist = gDirectory.Get(name)
if not hist:
raise error(
"MultiDraw: Could not find histogram to fill %r in current directory (varexp %r)."
% (name, varexp))
# SANITY CHECKS
vmatch = varregex2D.match(xvar)
if not vmatch or xvar.replace('::', '').count(':') == xvar.count(
'?'): # 1D, allow "(x>100 ? 1 : 0) >> h(2,0,2)"
pass
elif vmatch: # 2D histogram
yvar, xvar = vmatch.group(1), vmatch.group(2)
if not isinstance(hist, TH2):
raise error(
"MultiDraw: Existing histogram with name %r is not 2D! Found xvar=%r, yvar=%r..."
% (name, xvar, yvar))
else: # impossible
raise error(
'MultiDraw: Could not parse variable %r for %r to pattern %r: "%s"'
% (xvar, name, varregex2D.pattern, varexp))
if sumw2:
hist.Sumw2()
elif poisson:
hist.SetBinErrorOption(TH1D.kPoisson)
if drawoption:
hist.SetDrawOption(drawoption)
if name not in hists:
hists[name] = hist
results.append(hist)
# CHECK that the next formula is different to the previous one.
# If it is not, we add an ordinary TObject. In this way, the
# dynamic cast in MultiDraw.cxx fails, giving 'NULL', and the previous
# value is used. This saves the recomputing of identical values
if xvar != lastXVar:
formula = TTreeFormula("formula%i" % i, xvar, self)
if not formula.GetTree():
raise error(
"MultiDraw: TTreeFormula 'xvar' did not compile for %r:\n xvar: %r\n varexp: %r"
% (name, xvar, varexp))
formula.SetQuickLoad(True)
xformulae.append(formula)
else:
xformulae.append(TObject())
if yvar != None:
if yvar != lastYVar:
formula = TTreeFormula("formula%i" % i, yvar, self)
if not formula.GetTree():
raise error(
"MultiDraw: TTreeFormula 'yvar' did not compile for %r:\n yvar: %r\n varexp: %r"
% (name, yvar, varexp))
formula.SetQuickLoad(True)
yformulae.append(formula)
else:
yformulae.append(TObject())
if weight != lastWeight:
formula = TTreeFormula("weight%i" % i, weight, self)
if not formula.GetTree():
raise error(
"MultiDraw: TTreeFormula 'weight' did not compile for %r:\n weight: %r\n varexp: %r"
% (name, weight, varexp))
formula.SetQuickLoad(True)
weights.append(formula)
else:
weights.append(TObject())
lastXVar, lastYVar, lastWeight = xvar, yvar, weight
# CHECK that formulae are told when tree changes
manager = TTreeFormulaManager()
for formula in xformulae + yformulae + weights + [commonFormula]:
if isinstance(formula, TTreeFormula):
manager.Add(formula)
manager.Sync()
self.SetNotify(manager)
# DRAW
if verbosity >= 2:
print ">>> MultiDraw: xformulae=%s, yformulae=%s" % (
[x.GetTitle()
for x in xformulae], [y.GetTitle() for y in yformulae])
print ">>> MultiDraw: weights=%s, results=%s" % (
[w.GetTitle() for w in weights], results)
if len(yformulae) == 0:
_MultiDraw(self, commonFormula, makeTObjArray(xformulae),
makeTObjArray(weights), makeTObjArray(results),
len(xformulae))
elif len(xformulae) == len(yformulae):
_MultiDraw2D(self, commonFormula, makeTObjArray(xformulae),
makeTObjArray(yformulae), makeTObjArray(weights),
makeTObjArray(results), len(xformulae))
else:
raise error(
"MultiDraw: Given a mix of arguments for 1D (%d) and 2D (%d) histograms!"
% (len(xformulae), len(yformulae)))
return results
