def MultiDraw(self, Formulae, Compiled=False):
results, formulae, weights, formulaeStr, weightsStr = [], [], [], [], []
# lastFormula, lastWeight = None, None
for i, origFormula in enumerate(Formulae):
# Expand out origFormula and weight, otherwise just use weight of 1.
if type(origFormula) == tuple:
origFormula, weight = origFormula
else:
origFormula, weight = origFormula, "1"
# Our way is simpler, require each variable to end in (...) or [...] to give the binning
# and always create a new hist
split_var = origFormula.split(';')
origFormula = split_var[0]
print "Formula: ", origFormula, weight
var_binned_x = False
var_binned_y = False
var_binned_z = False
is_2d = False
is_3d = False
if origFormula[-1] == ')':
pos_open = origFormula.rfind('(')
pos_close = origFormula.rfind(')')
if origFormula[-1] == ']':
var_binned_x = True
pos_open = origFormula.rfind('[')
pos_close = origFormula.rfind(']')
if pos_open is -1 or pos_close is -1 or pos_open > pos_close:
raise RuntimeError('You bus')
bin_args_x = GetBinningArgs(origFormula[pos_open+1:pos_close], var_binned_x)
formula = origFormula[:pos_open].strip()
# Check if this is a 2D histogram with syntax
# [var_y],[var_x][binning_y],[binning_x]
if formula[-1] == ',':
is_2d = True
if formula[-2] == ')':
pos_open_y = formula.rfind('(')
pos_close_y = formula.rfind(')')
if formula[-2] == ']':
var_binned_y = True
pos_open_y = formula.rfind('[')
pos_close_y = formula.rfind(']')
if pos_open_y is -1 or pos_close_y is -1 or pos_open_y > pos_close_y:
raise RuntimeError('You bus')
bin_args_y = GetBinningArgs(formula[pos_open_y + 1:pos_close_y], var_binned_y)
formula = origFormula[:pos_open_y].strip()
# Check if this is a 3D histogram with syntax
# [var_y],[var_x],[var_z][binning_y],[binning_x],[binning_z]
if formula[-1] == ',':
is_3d = True
if formula[-2] == ')':
pos_open_z = formula.rfind('(')
pos_close_z = formula.rfind(')')
if formula[-2] == ']':
var_binned_z = True
pos_open_z = formula.rfind('[')
pos_close_z = formula.rfind(']')
if pos_open_z is -1 or pos_close_z is -1 or pos_open_z > pos_close_z:
raise RuntimeError('You bus')
bin_args_z = GetBinningArgs(formula[pos_open_z + 1:pos_close_z], var_binned_z)
formula = formula[:pos_open_z].split(',')
else:
formula = formula[:pos_open_y].split(',')
else:
formula = [formula]
ROOT.TH1.AddDirectory(False)
if not is_2d and not is_3d:
hist = ROOT.TH1D(origFormula+':'+weight, origFormula, *bin_args_x)
elif not is_3d:
hist = ROOT.TH2F(origFormula+':'+weight, origFormula, *(bin_args_x + bin_args_y))
else:
hist = ROOT.TH3F(origFormula+':'+weight, origFormula, *(bin_args_x + bin_args_y + bin_args_z))
if len(split_var) > 1:
hist.GetXaxis().SetTitle(split_var[1])
if len(split_var) > 2:
hist.GetXaxis().SetTitle(split_var[2])
hist.GetYaxis().SetTitle(split_var[1])
if len(split_var) > 2:
hist.GetXaxis().SetTitle(split_var[3])
hist.GetYaxis().SetTitle(split_var[2])
hist.GetZaxis().SetTitle(split_var[1])
if is_2d:
results.append(ROOT.TObject())
if is_3d:
results.append(ROOT.TObject())
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
for form in formula:
f = ROOT.TTreeFormula("formula%i" % i, form, self)
f.SetTitle(form)
if not f.GetTree():
raise RuntimeError("TTreeFormula didn't compile: " + form)
f.SetQuickLoad(True)
formulae.append(f)
formulaeStr.append(form)
f = ROOT.TTreeFormula("weight%i" % i, weight, self)
f.SetTitle(weight)
if not f.GetTree():
raise RuntimeError("TTreeFormula didn't compile: " + weight)
f.SetQuickLoad(True)
weights.append(f)
weightsStr.append(weight)
if Compiled:
fname = "%sSelector%s" % (self.GetName(), randomword(7))
self.MakeSelector(fname, "=legacy")
for line in fileinput.input('%s.h' % fname, inplace=1):
print line,
if line.startswith('#include <TSelector.h>'):
print '#include <TH1F.h>'
print '#include <TObjArray.h>'
if line.startswith(' virtual void Terminate();'):
print '\n TObjArray *hists;'
for line in fileinput.input('%s.C' % fname, inplace=1):
if line.startswith(' return kTRUE;'):
print ' %s::GetEntry(entry);' % fname
print ' double weight_value = 0.;'
for i, f in enumerate(formulaeStr):
print ' weight_value = %s;' % (weightsStr[i])
print ' if (weight_value) static_cast<TH1D*>(hists->UncheckedAt(%i))->Fill(%s, weight_value);' % (i, f)
# print ' std::cout << "%i %s:" << %s << "\\n";' % (i, weightsStr[i], weightsStr[i])
print ' if (entry % 50000 == 0) {'
print ' std::cout << "Done " << (double(entry) / (double(fChain->GetEntries())) * 100.0) << "% \\r";'
print ' std::cout.flush();'
print ' }'
print line,
if line.startswith('#include <TStyle.h>'):
print '#include<iostream>'
print 'using std::abs;'
ROOT.gROOT.LoadMacro('%s.C++' % fname)
selector = getattr(ROOT, fname)()
objarr = MakeTObjArray(results, takeOwnership=False)
selector.hists = objarr
self.Process(selector)
ROOT.gSystem.Unload('%s_C.so' % fname)
os.system('rm %s*' % fname)
return results
from ROOT import MultiDraw as _MultiDraw
from time import time
start = time()
# Ensure that formulae are told when tree changes
fManager = ROOT.TTreeFormulaManager()
for formula in formulae + weights:
if type(formula) == ROOT.TTreeFormula:
fManager.Add(formula)
fManager.Sync()
self.SetNotify(fManager)
# Draw everything!
_MultiDraw(self,
MakeTObjArray(formulae),
MakeTObjArray(weights),
MakeTObjArray(results, takeOwnership=False),
len(formulae))
print "Took %.2fs" % (time() - start), " " * 20
return results
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, *draw_list):
"""
Draws (projects) many histograms in one loop over a tree.
Instead of:
tree.Project("hname1", "ph_pt", "weightA")
tree.Project("hname2", "met_et", "weightB")
Do:
tree.MultiDraw( ("hname1", "ph_pt", "weightA" ),
("hname2", "met_et", "weightB" ) )
"""
histograms, variables, selections = [], [], []
last_variable, last_selection = None, None
for i, drawexp in enumerate(draw_list):
# Expand out origFormula and weight, otherwise just use weight of 1.
hname, variable, selection = drawexp
hist = ROOT.gDirectory.Get(hname)
if not hist:
raise RuntimeError(
"MultiDraw: Couldn't find histogram to fill '%s' in current directory."
% name)
histograms.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 variable != last_variable:
f = ROOT.TTreeFormula("variable%i" % i, variable, self)
if not f.GetTree():
raise RuntimeError("TTreeFormula didn't compile: " + variable)
f.SetQuickLoad(True)
variables.append(f)
else:
variables.append(ROOT.TObject())
if selection != last_selection:
f = ROOT.TTreeFormula("selection%i" % i, selection, self)
if not f.GetTree():
raise RuntimeError("TTreeFormula didn't compile: " + selection)
f.SetQuickLoad(True)
selections.append(f)
else:
selections.append(ROOT.TObject())
last_variable, last_selection = variable, selection
# Only compile MultiDraw once
try:
from ROOT import MultiDraw as _MultiDraw
except ImportError:
ROOT.gInterpreter.Declare(
open(os.environ['SUSY_ANALYSIS'] + '/lib/MultiDraw.cxx').read())
from ROOT import MultiDraw as _MultiDraw
# Ensure that formulae are told when tree changes
fManager = ROOT.TTreeFormulaManager()
for variable in variables + selections:
if type(variable) == ROOT.TTreeFormula:
fManager.Add(variable)
fManager.Sync()
self.SetNotify(fManager)
# Draw everything!
variables = MakeTObjArray(variables)
selections = MakeTObjArray(selections)
histograms = MakeTObjArray(histograms)
_MultiDraw(self, variables, selections, histograms, len(variables))
variables.Delete()
selections.Delete()
del fManager
return
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 = ROOT.TTreeFormula("CommonWeight", CommonWeight, self)
CommonWeightFormula.SetQuickLoad(True)
if not CommonWeightFormula.GetTree():
raise RuntimeError("TTreeFormula didn't compile: " + CommonWeight)
hists = {}
print Formulae
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:]
# 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
# Our way is simpler, require each variable to end in (...) or [...] to give the binning
# and always create a new hist
pos_open = origFormula.rfind('(')
pos_close = origFormula.rfind(')')
if pos_open is -1 or pos_close is -1 or pos_open > pos_close:
raise RuntimeError('You bus')
str_binning = [x.strip() for x in origFormula[pos_open+1:pos_close].split(',')]
binning = []
if len(str_binning) == 3:
binning = [int(str_binning[0]), float(str_binning[1]), float(str_binning[2])]
formula = origFormula[:pos_open].strip()
print binning
ROOT.TH1.AddDirectory(False)
hist = ROOT.TH1D(origFormula+':'+weight, origFormula, *binning)
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 = ROOT.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(ROOT.TObject())
if weight != lastWeight:
f = ROOT.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(ROOT.TObject())
lastFormula, lastWeight = formula, weight
from ROOT import MultiDraw as _MultiDraw
from time import time
start = time()
# Ensure that formulae are told when tree changes
fManager = ROOT.TTreeFormulaManager()
for formula in formulae + weights + [CommonWeightFormula, ]:
if type(formula) == ROOT.TTreeFormula:
fManager.Add(formula)
fManager.Sync()
self.SetNotify(fManager)
# Draw everything!
_MultiDraw(self, CommonWeightFormula,
MakeTObjArray(formulae),
MakeTObjArray(weights),
MakeTObjArray(results, takeOwnership=False),
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
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/TauTauResonances/plot/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
