htt/yields-postfit at master · qbuat/htt · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
#!/usr/bin/env python
# python imports
import os
import pickle
# root/rootpy imports
from rootpy.extern.ordereddict import OrderedDict
# local imports
from statstools.ufloat import ufloat
from mva import log
# -------------------------------------
def get_table_template():
    latex_lines = OrderedDict()
    latex_lines['cat_name'] = '&'
    latex_lines['sep_1']    = '\\hline'
    latex_lines['Signal_Z'] = 'ZH($m_H$ = 125 GeV) &'
    latex_lines['Signal_W'] = 'WH($m_H$ = 125 GeV) &'
    latex_lines['Signal_gg'] = 'ggH($m_H$ = 125 GeV) &'
    latex_lines['Signal_VBF'] = 'VBF H($m_H$ = 125 GeV) &'
    latex_lines['Higgs']    = 'H($m_H$ = 125 GeV) &'
    latex_lines['sep_2']    = '\\hline'
    latex_lines['Ztautau']  = 'Z$\\rightarrow\\tau\\tau$ &'
    latex_lines['QCD']      = 'Multi-jets &'
    latex_lines['Others']   = 'Others &'
    latex_lines['sep_3']    = '\\hline'
    latex_lines['TotalBkg'] = 'Total Bkg. &'
    latex_lines['sep_4']    = '\\hline'
    latex_lines['Data']     = 'Data &'
    return latex_lines

# --------------------------------------------
def get_sample_yield(yields_postfit, cat, sample_string):
    sample_yield = ufloat(0, 0)
    for key, yields in yields_postfit[cat].items():
        if sample_string in key:
            sample_yield = ufloat(yields[0], yields[1])
    return str(sample_yield)


if __name__ == '__main__':
    from rootpy.extern.argparse import ArgumentParser
    parser = ArgumentParser()
    parser.add_argument('file')
    args = parser.parse_args()

    #     output = os.path.splitext(args.file)[0] + '_postfit.pickle'
    with open(args.file) as fpickle:
        yields_postfit = pickle.load(fpickle)

        latex_lines = get_table_template()
        for cat, yields in yields_postfit.items():
            if '_11' in cat:
                continue
            log.info(cat)
            latex_lines['cat_name'] += cat + '&'
            latex_lines['Signal_Z']  += get_sample_yield(yields_postfit, cat, 'Signal_Z') + '&'
            latex_lines['Signal_W']  += get_sample_yield(yields_postfit, cat, 'Signal_W') + '&'
            latex_lines['Signal_gg'] += get_sample_yield(yields_postfit, cat, 'Signal_gg') + '&'
            latex_lines['Signal_VBF']+= get_sample_yield(yields_postfit, cat, 'Signal_VBF') + '&'
            latex_lines['Higgs']     += get_sample_yield(yields_postfit, cat, 'sum_sig') + '&'
            latex_lines['Ztautau']   += get_sample_yield(yields_postfit, cat, 'Ztautau') + '&'
            latex_lines['QCD']       += get_sample_yield(yields_postfit, cat, 'Fakes') + '&'
            latex_lines['Others']    += get_sample_yield(yields_postfit, cat, 'Others') + '&'
            latex_lines['TotalBkg']  += get_sample_yield(yields_postfit, cat, 'sum_bkg') + '&'
            latex_lines['Data']      += str(int(yields_postfit[cat]['Data'][0])) + '&'

        for _, line in latex_lines.items():
            if not 'hline' in line:
                line += '\\\\'

        for _, line in latex_lines.items():
            log.info(line)