def getDM(string):
"""Help function to get DM from string."""
matches = re.findall(r"_?(DM\d+)", string)
if not matches: return ""
if len(matches) > 1:
print warning('getDM: Found more than one match for shift: %s' %
(matches))
return matches[0]
def getShiftTitle(string):
"""Help function to format title, e.g. '_TES0p970' -> '-3% TES'."""
matches = re.findall(r"([a-zA-Z]+)(\d+[p\.]\d+)", string)
if not matches: return ""
if len(matches) > 1:
print warning(
'getShiftTitle: Found more than one match for shift: %s' %
(matches))
param, shift = matches[0]
shift = float(shift.replace('p', '.')) - 1.
if not shift: return "" #re.sub(r"_?[a-zA-Z]+\d+[p\.]\d+","",string)
title = " %s%% %s" % (("%+.2f" %
(100.0 * shift)).rstrip('0').rstrip('.'), param)
return title
def drawUpDownVariation(filename, dirname, samples, shifts, **kwargs):
"""Compare up/down variations of systematics, e.g. sample ZTT with it up and down variations for 'TES0.97'.
One can add several backgrounds with the same systematics into one by passing e.g. 'ZTT+ZJ+ZL=DY'."""
print '>>>\n>>> drawUpDownVariation("%s","%s")' % (filename, dirname)
file = TFile(filename, 'READ')
dir = file.Get(dirname)
tag = kwargs.get('tag', "")
var0 = kwargs.get('var', None)
xmin = kwargs.get('xmin', None)
xmax = kwargs.get('xmax', None)
outdir = kwargs.get('outdir', "shapes")
channel = kwargs.get('channel', "mutau")
text = kwargs.get('text', "")
ensureDirectory(outdir)
if not isList(samples): samples = [samples]
if not dir:
print warning('drawUpDownVariation: did not find dir "%s"' % (dirname))
for sample in samples:
for shift in shifts:
if not shift: continue
if '$' in shift:
shift = shift.replace('$CAT',
dirname).replace('$CHANNEL', channel)
print '>>> sample "%s" for shift "%s"' % (sample, shift)
skip = False
matches = re.findall(r"(.+\+.+)=(.*)",
sample) # e.g. "ZTT+ZJ+ZL=DY"
if len(matches) == 0:
sampleUp = "%s_%sUp" % (sample, shift)
sampleDn = "%s_%sDown" % (sample, shift)
hist = dir.Get(sample)
histUp = dir.Get(sampleUp)
histDn = dir.Get(sampleDn)
if hist and not histUp and not histDn:
print warning(
'drawUpDownVariation: did not find the variations "%s" or "%s" in directory "%s"'
% (sampleUp, sampleDn, dir.GetName()),
pre=" ")
skip = True
continue
names = [sampleUp, sample, sampleDn]
hists = [histUp, hist, histDn]
for hist1, name in zip(hists, names):
if not hist1:
print warning(
'drawUpDownVariation: did not find "%s" in directory "%s"'
% (name, dir.GetName()),
pre=" ")
skip = True
break
else: # add samples
hist, histUp, histDn = None, None, None
sample = matches[0][1]
for subsample in matches[0][0].split('+'):
print '>>> adding subsample "%s"' % (subsample)
subsampleUp = "%s_%sUp" % (subsample, shift)
subsampleDn = "%s_%sDown" % (subsample, shift)
subhist = dir.Get(subsample)
subhistUp = dir.Get(subsampleUp)
subhistDn = dir.Get(subsampleDn)
subhists = [subhistUp, subhist, subhistDn]
subnames = [subsampleUp, subsample, subsampleDn]
for hist1, name in zip(subhists, subnames):
if not hist1:
print warning(
'drawUpDownVariation: did not find "%s" in directory "%s"'
% (name, dir.GetName()),
pre=" ")
skip = True
break
if skip: break
if hist: hist.Add(subhist)
else: hist = subhist
if histUp: histUp.Add(subhistUp)
else: histUp = subhistUp
if histDn: histDn.Add(subhistDn)
else: histDn = subhistDn
hist.SetTitle(sample)
histUp.SetTitle("%s_%sUp" % (sample, shift))
histDn.SetTitle("%s_%sDown" % (sample, shift))
hists = [histUp, hist, histDn]
if skip: continue
var = var0 or hist.GetXaxis().GetTitle()
varname = formatVariable(var)
vartitle = variable_dict.get(var, var)
tshift = formatFilename(shift)
nshift = ('_' + tshift) if tshift else ""
title = formatCategory(sample, tshift)
hist.SetTitle("nominal")
histUp.SetTitle("up variation")
histDn.SetTitle("down variation")
canvasname = "%s/%s_%s_%s%s%s.png" % (
outdir, sample.replace('_TES1.000', '').replace(
'.', 'p'), varname, dirname, tag, nshift)
exts = ['pdf', 'png'] if args.pdf else ['png']
plot = Plot(hists)
plot.plot(vartitle,
title=title,
ratio=2,
linestyle=False,
xmin=xmin,
xmax=xmax,
errorbars=True,
text=text)
plot.saveAs(canvasname, ext=exts)
plot.close()
file.Close()
def drawVariations(filename, dirname, samples, variations, **kwargs):
"""Compare variations, e.g. sample ZTT with variations ['TES0.97','TES1.03']."""
print '>>>\n>>> drawVariations("%s","%s")' % (filename, dirname)
file = TFile(filename, 'READ')
dir = file.Get(dirname)
tag = kwargs.get('tag', "")
var0 = kwargs.get('var', None)
xmin = kwargs.get('xmin', None)
xmax = kwargs.get('xmax', None)
rmin = kwargs.get('rmin', None)
rmax = kwargs.get('rmax', None)
outdir = kwargs.get('outdir', "shapes")
position = kwargs.get('position', 'right')
text = kwargs.get('text', "")
if not isList(samples): samples = [samples]
if not dir:
print warning('drawVariations: did not find dir "%s"' % (dirname))
ensureDirectory(outdir)
ratio = True
for sample in samples:
hists = []
ratio = int(len(variations) / 2)
for i, variation in enumerate(variations):
#print '>>> sample "%s" for shift "%s"'%(sample, shift)
name = "%s_%s" % (sample, variation.lstrip('_'))
hist = dir.Get(name)
if not hist:
print warning(
'drawVariations: did not find "%s" in directory "%s"' %
(name, dir.GetName()),
pre=" ")
continue
sampletitle = sample
if variation in variation_dict:
sampletitle += variation_dict[variation]
else:
sampletitle += getShiftTitle(variation)
hist.SetTitle(sampletitle)
if sample == sampletitle or "nom" in sampletitle.lower():
ratio = i + 1
hists.append(hist)
if len(hists) != len(variations):
print warning(
'drawVariations: number of hists (%d) != number variations (%d)'
% (len(hists), len(variations)),
pre=" ")
var = var0 or hists[0].GetXaxis().GetTitle()
varname = formatVariable(var)
vartitle = variable_dict.get(var, var)
shift0 = formatFilename(variations[0])
shift1 = formatFilename(variations[-1])
shift = ""
i = 0
for i, letter in enumerate(shift0):
if i >= len(shift1): break
if letter == shift1[i]: shift += shift0[i]
else: break
nshift = "_%s-%s" % (shift0, shift1.replace(shift, ""))
title = formatCategory(dirname, shift)
canvasname = "%s/%s_%s_%s%s%s.png" % (outdir, sample, varname, dirname,
tag, nshift.replace('.', 'p'))
exts = ['pdf', 'png'] if args.pdf else ['png']
plot = Plot(hists)
plot.plot(vartitle,
title=title,
text=text,
ratio=ratio,
linestyle=False,
xmin=xmin,
xmax=xmax,
rmin=rmin,
rmax=rmax,
position=position)
plot.saveAs(canvasname, ext=exts)
plot.close()
file.Close()
def drawPrePostFit(filename, dirname, samples, **kwargs):
"""Create pre- and post-fit plots from histograms and data."""
print '>>>\n>>> drawPrePostFit("%s","%s")' % (filename, dirname)
usertitle = kwargs.get('title', None)
ratio = kwargs.get('ratio', True)
tag = kwargs.get('tag', "")
var0 = kwargs.get('var', None)
xmin = kwargs.get('xmin', None)
xmax = kwargs.get('xmax', None)
ymargin = kwargs.get('ymargin', 1.22)
dirnametag = kwargs.get('dirnametag', dirname)
outdir = kwargs.get('outdir', "plots")
group = kwargs.get('group', [])
signals = kwargs.get('signal', [])
signaltitles = kwargs.get('signaltitle', [])
signalheader = kwargs.get('signalheader', "")
signalfilenames = kwargs.get('signalfile', []) # files to take signal from
signalstrengths = kwargs.get('r', [])
upscale = kwargs.get('upscale', 1.0)
userposition = kwargs.get('position', '')
text = kwargs.get('text', "")
append_dict = kwargs.get('apptitle', {})
signalPostfit = kwargs.get('signalPostfit', "postfit")
signalposition = kwargs.get('signalposition', 'x=0.08')
ymax = None
#signalPostScale = 1.0
ensureDirectory(outdir)
signals = ensureList(signals)
signaltitles = ensureList(signaltitles)
signalstrengths = ensureList(signalstrengths)
fits = ['prefit', 'postfit']
file = ensureTFile(filename, 'READ')
if not file:
print error('drawPrePostFit: did not find file "%s"' % (filename),
pre=" ")
exit(1)
# ARRANGE signal sample files, (histogram) names and titles
signalfiles = []
signallist = []
if len(signals) > 0 and len(signaltitles) == 0:
signaltitles = [
sample_dict[s] + append_dict.get(s, "") for s in signals
]
elif len(signals) == 1 and len(signals) < len(signaltitles):
signals = signals * len(signaltitles)
if len(signals) > 0 and len(signalstrengths) == 0:
signalstrengths = [1.0] * len(signals)
if len(signals) != len(signaltitles):
print error(
'drawPrePostFit: len(signals) = %d != %d = len(signaltitles)' %
(len(signals), len(signaltitles)),
pre=" ")
exit(1)
if len(signalfilenames) + 1 == len(signals):
signalfilenames += [filename] + signalfilenames
elif len(signalfilenames) != len(signals):
print error(
'drawPrePostFit: len(signalfilenames) = %d != %d = len(signals)' %
(len(signalfilenames), len(signals)),
pre=" ")
exit(1)
if len(signalfilenames) != len(signals):
print error(
'drawPrePostFit: len(signalfilenames) = %d != %d = len(signals)' %
(len(signalfilenames), len(signals)),
pre=" ")
exit(1)
for signalfilename in signalfilenames:
signalfile = file if signalfilename == filename else ensureTFile(
signalfilename, 'READ')
signalfiles.append(signalfile)
signallist = zip(signalfiles, signals, signaltitles, signalstrengths)
# DRAW PRE-/POST-FIT
for fit in fits:
fitdirname = "%s_%s" % (dirname, fit)
dir = file.Get(fitdirname)
if not dir:
print warning('drawPrePostFit: did not find dir "%s"' %
(fitdirname),
pre=" ")
return
histsD = []
histsB = []
histsS = []
# DATA & BACKGROUNDS
for sample in reversed(samples):
histname = "%s/%s" % (fitdirname, sample)
stitle = sample_dict[sample] + append_dict.get(sample, "")
hist = file.Get(histname)
if not hist:
print warning(
'drawPrePostFit: could not find "%s" template in directory "%s_%s"'
% (sample, dirname, fit),
pre=" ")
continue
if 'data_obs' in sample:
histsD.append(hist)
hist.SetLineColor(1)
ymax = hist.GetMaximum() * ymargin
else:
histsB.append(hist)
hist.SetTitle(stitle)
# SIGNALS
for signalfile, sample, stitle, signalr in signallist:
histname = "%s/%s" % (fitdirname, sample)
hist = signalfile.Get(histname)
if not hist:
print warning(
'drawPrePostFit: could not find "%s" template in directory "%s_%s"'
% (sample, dirname, fit),
pre=" ")
continue
scale = upscale / signalr if 'post' in fit else upscale
hist.Scale(scale)
histsS.append(hist)
hist.SetTitle(stitle)
if len(histsB) == 0:
print warning(
'drawPrePostFit: could not find any templates in directory "%s"'
% (dirname),
pre=" ")
continue
if len(histsD) == 0:
print warning(
'drawPrePostFit: could not find a data template in directory "%s"'
% (dirname),
pre=" ")
continue
for groupargs in group:
histsB = groupHistsInList(histsB, *groupargs)
var = var0 or histsB[0].GetXaxis().GetTitle()
varname = formatVariable(var)
vartitle = variable_dict.get(var, var)
position = userposition if userposition else "x=0.67" if 'post' in fit else "x=0.66"
title = formatCategory(dirname) if usertitle == None else usertitle
errortitle = "pre-fit stat. + syst. unc." if 'pre' in fit else "post-fit unc."
canvasname = "%s/%s_%s%s_%s.png" % (outdir, varname, dirnametag, tag,
fit)
exts = ['pdf', 'png'] if args.pdf else ['png']
plot = Plot(histsD, histsB, histsS, stack=True)
plot.plot(vartitle,
xmin=xmin,
xmax=xmax,
ymax=ymax,
ratio=ratio,
staterror=True,
colors=colors,
linewidth=3,
title=title,
text=text,
errortitle=errortitle,
position=position,
signallegend=signalheader,
signaltitle=signalheader,
signalposition=signalposition)
plot.saveAs(canvasname, ext=exts)
plot.close()
def drawStacks(filename, dirname, samples, shifts=[""], **kwargs):
"""Compare shapes."""
print '>>>\n>>> drawStacks("%s","%s")' % (filename, dirname)
file = TFile(filename, 'READ')
dir = file.Get(dirname)
channel = kwargs.get('channel', "")
tag = kwargs.get('tag', "")
var0 = kwargs.get('var', None)
xmin = kwargs.get('xmin', None)
xmax = kwargs.get('xmax', None)
outdir = kwargs.get('outdir', "control")
blind = kwargs.get('blind', [])
group = kwargs.get('group', [])
ensureDirectory(outdir)
if not isList(samples): samples = [samples]
if not dir:
print warning('drawStacks: did not find dir "%s"' % (dirname),
pre=" ")
if channel: channel += '-'
ensureDirectory(outdir)
for shift in shifts:
if '$' in shift:
shift = shift.replace('$CAT', dirname).replace('$CHANNEL', channel)
variations = ['Up', 'Down'] if shift else [""]
for variation in variations:
histsB = []
histsD = []
samplesS = []
nShifted = 0
for sample in reversed(samples):
hist = None
if shift: # SHIFTED
sample1 = "%s_%s%s" % (sample, shift, variation)
hist = dir.Get(sample1)
nShifted += 1
if not hist:
hist = dir.Get(sample)
else: # NOMINAL
hist = dir.Get(sample)
if not hist:
print warning(
'drawStacks: could not find "%s" template in directory "%s"'
% (sample, dir.GetName()),
pre=" ")
continue
hist.SetName(hist.GetName().replace('_ttbar',
'').replace('_ztt', ''))
hist.SetTitle(sample_dict[sample])
if 'data_obs' in sample:
histsD.append(hist)
hist.SetLineColor(1)
if blind:
for bin in range(hist.FindBin(blind[0]),
hist.FindBin(blind[1] + 2)):
hist.SetBinContent(bin, 0)
else:
histsB.append(hist)
if len(histsB) == 0 or (shift and nShifted == 0):
print warning(
'drawStacks: could not find any "%s" templates in directory "%s"'
% (shift, dir.GetName()),
pre=" ")
continue
for groupargs in group:
histsB = groupHistsInList(histsB, *groupargs)
var = var0 or histsB[0].GetXaxis().GetTitle()
varname = formatVariable(var)
vartitle = variable_dict.get(var, var)
tshift = formatFilename(shift)
nshift = ('_' + tshift + variation) if tshift else ""
title = formatCategory(dirname,
tshift,
variation,
channel=channel.rstrip('-'))
canvasname = "%s/%s_%s%s%s%s.png" % (outdir, var, channel, dirname,
tag, nshift)
exts = ['pdf', 'png'] if args.pdf else ['png']
plot = Plot(histsD, histsB, stack=True)
plot.plot(vartitle,
title=title,
ratio=True,
staterror=True,
xmin=xmin,
xmax=xmax)
plot.saveAs(canvasname, ext=exts)
plot.close()
file.Close()