postfit.py

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')