Skip to content

Instantly share code, notes, and snippets.

@annawoodard annawoodard/
Last active May 10, 2017

What would you like to do?
from collections import defaultdict
import os
import re
import numpy as np
import ROOT
from rootpy.plotting import Hist, HistStack, Legend, Canvas
from import get_style, set_style
from rootpy.plotting.utils import draw
from import CMS_label
workdir = '/afs/'
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)
pars = fit.floatParsFinal()
workspace.saveSnapshot("snapshot", ROOT.RooArgSet(pars), ROOT.kTRUE)
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)
normalization = workspace.function(name)
histo.SetBinContent(i + 1, normalization.getVal())
histo.xaxis.SetBinLabel(i + 1, binlabels[i])
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 ='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')
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)
CMS_label("", lumi=35.9)
name = '{}fit_{}_{}.pdf'.format(
'post' if postfit else 'pre',
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.