htt/tree-jackknife 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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
#!/usr/bin/env python
#
from rootpy.io import root_open
from rootpy.tree import Cut
import numpy as np
from numpy.lib import recfunctions
from root_numpy import array2tree
# local imports
from mva.cmd import get_parser
from mva.analysis import get_analysis
from mva.categories import CATEGORIES
from mva import log
args = get_parser(actions=False).parse_args()
analysis = get_analysis(args)

CB = CATEGORIES['cuts']
MVA = CATEGORIES['mva_all']

CB_cuts = Cut()
for cat in CB:
    CB_cuts |= cat.cuts

MVA_cuts = Cut()
for cat in MVA:
    MVA_cuts |= cat.cuts

fields = ['EventNumber', 'mmc1_mass', 'dEta_tau1_tau2']
names = ('EvtNum', 'MMC', 'DETA')
target_region = args.target_region
data = analysis.data

def get_record_overlap(mva_cat, cba_cat, clf=None):
    rec = data.merged_records(mva_cat, target_region,
                              fields=fields, include_weight=False, cuts=cb_cat.cuts)
    rec.dtype.names = names
    if clf:
        scores, _ = data.scores(clf, mva_cat, target_region,
                                systematics=False, cuts=cb_cat.cuts)
    else:
        scores = -999*np.ones(len(rec))
    rec = recfunctions.rec_append_fields(rec, data=scores, names='BDT', dtypes='f4')
    jks_mva = [mva_cat.jk_number for i in range(0, len(rec))]
    jks_cba = [cba_cat.jk_number for i in range(0, len(rec))]
    rec = recfunctions.rec_append_fields(rec, data=jks_mva, names='catMVA', dtypes='i8')
    rec = recfunctions.rec_append_fields(rec, data=jks_cba, names='catCBA', dtypes='i8')
    ismva = np.ones(len(rec))
    iscba = np.ones(len(rec))
    rec = recfunctions.rec_append_fields(rec, data=ismva, names='IsMVA', dtypes='i8')
    rec = recfunctions.rec_append_fields(rec, data=iscba, names='IsCBA', dtypes='i8')
    return rec

def get_record_disjonction(cat, exclusion_cut, clf=None, cat_type='MVA', excl_type='CBA'):
    rec = data.merged_records(cat, target_region,
                              fields=fields, include_weight=False, cuts=exclusion_cut)

    rec.dtype.names = names
    if clf:
        scores, _ = data.scores(clf, cat, target_region,
                                systematics=False, cuts=exclusion_cut)
    else:
        scores = -999*np.ones(len(rec))
    rec = recfunctions.rec_append_fields(rec, data=scores, names='BDT', dtypes='f4')
    jks = [cat.jk_number for i in range(0, len(rec))]
    rec = recfunctions.rec_append_fields(rec, data=jks, names='cat{0}'.format(cat_type), dtypes='i8')
    rec = recfunctions.rec_append_fields(rec, data=-999*np.ones(len(rec)),
                                         names='cat{0}'.format(excl_type), dtypes='i8')
    rec = recfunctions.rec_append_fields(rec, data=np.ones(len(rec)),
                                         names='Is{0}'.format(cat_type), dtypes='i8')
    rec = recfunctions.rec_append_fields(rec, data=np.zeros(len(rec)),
                                         names='Is{0}'.format(excl_type), dtypes='i8')
    return rec

recs = []
# Iterate over MVA categories
for category in analysis.iter_categories('mva_all'):
    if not category.analysis_control:
        clf = analysis.get_clf(category, load=True, mass=125, transform=True)
    else:
        clf = None

    for cb_cat in CB:
        rec = get_record_overlap(category, cb_cat, clf=clf)
        recs.append(rec)
    rec = get_record_disjonction(category, -CB_cuts, clf=clf)
    recs.append(rec)

# Iterate over CBA categories
for category in analysis.iter_categories('cuts'):
    clf = None
    rec = get_record_disjonction(category, -MVA_cuts, cat_type='CBA', excl_type='MVA')
    if len(rec)>0:
        recs.append(rec)

big_rec = np.concatenate(recs)
out = root_open('hh_jackknife.root', 'recreate')
out.cd()
output_name = 'datatree'
outtree = array2tree(big_rec, name='datatree')
outtree.Write()
out.Close()