Skip to content

Instantly share code, notes, and snippets.

@annawoodard annawoodard/postfit.py
Last active May 10, 2017

Embed
What would you like to do?
from collections import defaultdict
import os
import re
import numpy as np
import ROOT
ROOT.gROOT.SetBatch(True)
from rootpy.plotting import Hist, HistStack, Legend, Canvas
from rootpy.plotting.style import get_style, set_style
from rootpy.plotting.utils import draw
from rootpy.plotting.style.cmstdr.labels import CMS_label
workdir = '/afs/crc.nd.edu/user/a/awoodard/www/ttV/34/1/24'
if not os.path.isdir(os.path.join(workdir, 'plots', 'fit')):
os.makedirs(os.path.join(workdir, 'plots', 'fit'))
labels = {
'ttZ': 't#bar{t}Z',
'WZ': 'WZ',
'ttX': 't#bar{t}X',
'rare': 'rare',
'ttW': 't#bar{t}W',
'ttH': 't#bar{t}H',
'charge': 'charge',
'ZZ': 'ZZ',
'Fake': 'nonprompt',
'fake': 'nonprompt'
}
colors = {
'ttZ': 91,
'WZ': 51,
'ttH': 'saddlebrown',
'rare': 8,
'ttW': '#006517',
'ttX': '#FF3119',
'charge': 'darkgray',
'ZZ': 'orchid',
'Fake': '#9196fa',
'fake': '#9196fa'
}
analyses = {
'2l': {
'bins': [
'ch1', 'ch2', 'ch13', 'ch17', 'ch18',
'ch19', 'ch20', 'ch21', 'ch22', 'ch23',
'ch3', 'ch4', 'ch5', 'ch6', 'ch7',
'ch8', 'ch9', 'ch10', 'ch11', 'ch12'
],
'processes': ['ttW', 'Fake', 'charge', 'ttH', 'ttX', 'rare', 'WZ'],
'labels': ['2j', '3j1bj', '3j#geq2bj', '#geq4j1bj', '#geq4j2bj'] * 4,
'prefix': 'ch2_'
},
'2l-cr': {
'bins': ['ch16', 'ch14', 'ch15'],
'processes': ['ttW', 'Fake', 'charge', 'ttH', 'ttX', 'rare', 'WZ'],
'labels': ['BDT<0, 2J', 'BDT<0, 3J', 'BDT<0, #geq4J'],
'prefix': 'ch2_'
},
'3l': {
'bins': ['ch1', 'ch2', 'ch3', 'ch4', 'ch5', 'ch6', 'ch7', 'ch8', 'ch9'],
'processes': ['Fake', 'rare', 'WZ', 'ttX', 'ttH', 'ttZ'],
'labels': ['nB0, nJ2', 'nB0, nJ3', 'nB0, nJ#geq4', 'nB1, nJ2', 'nB1, nJ3', 'nB1,nJ#geq4', 'nB#geq2,nJ2',
'nB#geq2,nJ3', 'nB#geq2,nJ#geq4'],
'prefix': 'ch1_ch1_'
},
'4l': {
'bins': ['B0', 'B1'],
'processes': ['fake', 'rare', 'ttX', 'ttH', 'ttZ', 'ZZ'],
'labels': ['0', '#geq1'],
'prefix': 'ch1_ch2_'
}
}
def make_plot(tag, analysis, postfit=True, eft=None, width=1000, **kwargs):
bins = analyses[analysis]['bins']
processes = analyses[analysis]['processes']
prefix = '' if (tag in analyses) else analyses[analysis]['prefix'] # prefix is added when combining cards
binlabels = analyses[analysis]['labels']
if postfit:
fit_name = 'fit_s' if (tag == 'ttZ' or (tag in analyses)) else 'fit'
else: # we need the SM version for pre-fit, because MultiDimFit doesn't save the prefit result
fit_name = 'nuisances_prefit_res'
file = ROOT.TFile(os.path.join(workdir, 'workspaces', '{}.root'.format(tag)))
workspace = file.Get('w')
file = ROOT.TFile(os.path.join(workdir, 'fit-result-{}.root'.format(tag)))
fit = file.Get(fit_name)
file.Close()
pars = fit.floatParsFinal()
workspace.saveSnapshot("snapshot", ROOT.RooArgSet(pars), ROOT.kTRUE)
workspace.loadSnapshot("snapshot")
histos = []
normalizations = defaultdict(list)
for process in processes:
histo = Hist(len(bins), 0, len(bins), title=labels[process], drawstyle='hist', markersize=0, legendstyle='F')
histo.fillstyle = 'solid'
histo.fillcolor = colors[process]
histo.linecolor = colors[process]
histo.linewidth = 0
for i, bin in enumerate(bins):
name = 'n_exp_bin{}_proc_{}'.format(prefix + bin, process)
normalizations[bin].append(name)
normalization = workspace.function(name)
histo.SetBinContent(i + 1, normalization.getVal())
histo.xaxis.SetBinLabel(i + 1, binlabels[i])
histos.append(histo)
histos.sort(key=lambda x: x.Integral(), reverse=True) # biggest at the bottom, smallest at the top
mc = HistStack(histos)
uncertainties = Hist(len(bins), 0, len(bins), title='uncertainties', drawstyle='e2', markerstyle=1, legendstyle='F')
uncertainties.fillstyle = '\\'
uncertainties.fillcolor = 'gray'
for i, bin in enumerate(bins):
name = 'total_mc_{}'.format(prefix + bin)
workspace.factory('sum::{}({})'.format(name, ', '.join(normalizations[bin])))
total = workspace.function(name)
uncertainties.SetBinContent(i + 1, total.getVal())
uncertainties.SetBinError(i + 1, total.getPropagatedError(fit))
data_obs = workspace.data('data_obs')
data = Hist(len(bins), 0, len(bins), title='data', drawstyle='E1 X0', legendstyle='LEP')
indices = dict(((b, v) for v, b in enumerate(bins))) # sometimes bins aren't plotted in ascending order
for i in range(data_obs.numEntries()):
observed = data_obs.sumEntries("CMS_channel == {}".format(i))
bin = data_obs.get(i).getCatLabel("CMS_channel").replace(prefix, '')
if bin in bins:
data.SetBinContent(indices[bin] + 1, observed)
data.SetBinError(indices[bin] + 1, np.sqrt(observed))
style = get_style('cmstdr')
set_style(style)
canvas = Canvas(width=width, height=600)
draw([mc, uncertainties, data], ytitle='Events', xdivisions=len(bins) + 1, **kwargs)
legend = Legend(histos + [data, uncertainties], rightmargin=0.03, leftmargin=0.09, entryheight=0.001)
legend.SetNColumns(4)
legend.Draw()
CMS_label("", lumi=35.9)
canvas.Modified()
canvas.Update()
name = '{}fit_{}_{}.pdf'.format(
'post' if postfit else 'pre',
analysis,
eft if eft else tag
)
canvas.Print(os.path.join(workdir, 'plots', 'fit', name))
make_plot('2l', '2l', ylimits=(0, 80))
make_plot('2l', '2l-cr', ylimits=(0, 80))
make_plot('ttZ', '3l', ylimits=(0.5, 6000), logy=True)
make_plot('ttZ', '4l', ylimits=(0, 40), xtitle='b-jet multiplicity')
# make_plot('2d', '2l')
# make_plot('2d', '2l-cr', ylimits=(0, 350))
# make_plot('2d', '3l', ylimits=(0.5, 6000), logy=True)
# make_plot('2d', '4l', xtitle='b-jet multiplicity')
# for coefficient in ['cu', 'cHu', 'cuW', 'cuB']:
# make_plot(coefficient, '2l', eft=coefficient)
# make_plot(coefficient, '2l-cr', eft=coefficient, ylimits=(0, 350))
# make_plot(coefficient, '3l', eft=coefficient, ylimits=(0.5, 6000), logy=True)
# make_plot(coefficient, '4l', eft=coefficient, xtitle='b-jet multiplicity')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.