-
Notifications
You must be signed in to change notification settings - Fork 1
/
createTurnOn_ditau_pnet.py
211 lines (192 loc) · 10 KB
/
createTurnOn_ditau_pnet.py
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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
#!/usr/bin/env python
import argparse
from array import array
import math
import numpy as np
import os
import re
import sys
import ROOT
parser = argparse.ArgumentParser(description='Create turn on curves.')
parser.add_argument('--input', required=True, type=str, nargs='+', help="input files")
parser.add_argument('--output', required=True, type=str, help="output file prefix")
parser.add_argument('--channels', required=False, type=str, default='mutau', help="channels to process")
parser.add_argument('--decay-modes', required=False, type=str, default='all', help="decay modes to process")
parser.add_argument('--working-points', required=False, type=str,
default='Tight',
help="working points to process")
args = parser.parse_args()
path_prefix = '' if 'TAU-Trigger-NANO' in os.getcwd() else 'TAU-Trigger-NANO/'
sys.path.insert(0, path_prefix + 'Common/python')
from AnalysisTypes import *
from AnalysisTools import *
import RootPlotting
ROOT.ROOT.EnableImplicitMT(4)
ROOT.gROOT.SetBatch(True)
ROOT.TH1.SetDefaultSumw2()
RootPlotting.ApplyDefaultGlobalStyle()
def CreateBins(max_pt, for_fitting):
#bins = np.arange(20, 100, step=10)
bins = np.arange(20, 40, step=4)
bins = np.append(bins, np.arange(40, 60, step=5))
bins = np.append(bins, np.arange(60, 100, step=10))
high_pt_bins = [ 100, 150, 200, 300, 400, 500, 650, 800, 1000 ]
# bins = np.arange(20, 100, step=8)
# high_pt_bins = [ 100, 200, 300, 400, 500, 650, 800, 1000 ]
n = 0
while n < len(high_pt_bins) and high_pt_bins[n] < max_pt:
n += 1
use_logx = max_pt > 200
return np.append(bins, high_pt_bins[0:n+1]), use_logx
class TurnOnData:
def __init__(self):
self.hist_total = None
self.hist_passed = None
self.eff = None
def CreateHistograms(input_file, channels, decay_modes, discr_name, working_points, hist_models, label, var,
output_file):
df = ROOT.RDataFrame('Events', input_file)
turnOn_data = {}
dm_labels = {}
for dm in decay_modes:
if dm == 'all':
dm_labels[dm] = ''
df_dm = df
elif dm == '1011':
dm_labels[dm] = '_dm{}'.format(dm)
df_dm = df.Filter('tau_decayMode == 10 | tau_decayMode == 11')
else:
dm_labels[dm] = '_dm{}'.format(dm)
df_dm = df.Filter('tau_decayMode == {}'.format(dm))
turnOn_data[dm] = {}
for wp in working_points:
wp_bit = ParseEnum(DiscriminatorWP, wp)
# df_wp = df_dm.Filter('({} & (1 << {})) != 0'.format(discr_name, wp_bit))
df_wp = df_dm.Filter('{0} >= {1}'.format(discr_name, wp_bit))
turnOn_data[dm][wp] = {}
for channel in channels:
turnOn_data[dm][wp][channel] = {}
year = input_file[-9:-5]
df_ch = df_wp.Filter('pass_{0}_deeptau_{1} > 0.5'.format(channel, year))
for model_name, hist_model in hist_models.items():
turn_on = TurnOnData()
turn_on.hist_total = df_wp.Histo1D(hist_model, var, 'weight')
turn_on.hist_passed = df_ch.Histo1D(hist_model, var, 'weight')
turnOn_data[dm][wp][channel][model_name] = turn_on
for dm in decay_modes:
for wp in working_points:
for channel in channels:
for model_name, hist_model in hist_models.items():
turn_on = turnOn_data[dm][wp][channel][model_name]
name_pattern = '{}_{}_{}{}_{}_{{}}'.format(label, channel, wp, dm_labels[dm], model_name)
turn_on.name_pattern = name_pattern
if 'fit' in model_name:
passed, total, eff = AutoRebinAndEfficiency(turn_on.hist_passed.GetPtr(),
turn_on.hist_total.GetPtr(), bin_scan_pairs)
else:
passed, total = turn_on.hist_passed.GetPtr(), turn_on.hist_total.GetPtr()
# # new add by botao
# if (passed.Integral() < 0) or (total.Integral() < 0):
# continue
# # end add
try:
FixEfficiencyBins(passed, total)
except:
continue
turn_on.eff = ROOT.TEfficiency(passed, total)
eff = turn_on.eff
output_file.WriteTObject(total, name_pattern.format('total'), 'Overwrite')
output_file.WriteTObject(passed, name_pattern.format('passed'), 'Overwrite')
output_file.WriteTObject(eff, name_pattern.format('eff'), 'Overwrite')
# print(name_pattern)
# print('hist_total {}'.format(turn_on.hist_total.GetPtr().GetNbinsX()))
# for n in range(turn_on.hist_total.GetPtr().GetNbinsX() + 1):
# print('\t{} {} +/- {} {} +/- {}'.format(turn_on.hist_total.GetPtr().GetBinLowEdge(n+1), turn_on.hist_total.GetPtr().GetBinContent(n+1), turn_on.hist_total.GetPtr().GetBinError(n+1), turn_on.hist_passed.GetPtr().GetBinContent(n+1), turn_on.hist_passed.GetPtr().GetBinError(n+1)))
# print('hist_passed {}'.format(turn_on.hist_passed.GetPtr().GetNbinsX()))
# for n in range(turn_on.hist_passed.GetPtr().GetNbinsX() + 1):
# print('\t{} {}'.format(turn_on.hist_passed.GetPtr().GetBinLowEdge(n+1), ))
#if 'fit' not in model_name:
# turn_on.eff = ROOT.TEfficiency(turn_on.hist_passed.GetPtr(), turn_on.hist_total.GetPtr())
# output_file.WriteTObject(turn_on.eff, name_pattern.format('passed'), 'Overwrite')
return turnOn_data
output_file = ROOT.TFile(args.output + '.root', 'RECREATE')
input_files = args.input
print(input_files)
n_inputs = len(input_files)
labels = []
for i in range(n_inputs):
input_name = os.path.basename(input_files[i])
labels.append(os.path.splitext(input_name)[0][-4:])
var = 'leading_tau_pt'
title, x_title = '#tau p_{T}', '#tau p_{T} (GeV)'
decay_modes = args.decay_modes.split(',')
channels = args.channels.split(',')
working_points = args.working_points.split(',')
bins, use_logx = CreateBins(200, False)
hist_models = {
'plot': ROOT.RDF.TH1DModel(var, var, len(bins) - 1, array('d', bins)),
}
turnOn_data = [None] * n_inputs
for input_id in range(n_inputs):
print("Creating {} histograms...".format(labels[input_id]))
turnOn_data[input_id] = CreateHistograms(input_files[input_id], channels, decay_modes, 'leading_tau_idDeepTauVSjet', # tau_idDeepTau2017v2p1VSjet,
working_points, hist_models, labels[input_id], var, output_file)
colors_list = [ ROOT.kBlack, ROOT.kRed, ROOT.kGreen, ROOT.kViolet]
colors = colors_list[0:n_inputs]
canvas = RootPlotting.CreateCanvas()
n_plots = len(decay_modes) * len(channels) * len(working_points)
plot_id = 0
for channel in channels:
for wp in working_points:
for dm in decay_modes:
if dm == 'all':
dm_label = ''
dm_plain_label = ''
else:
dm_label = ' DM={}'.format(dm)
dm_plain_label = '_dm{}'.format(dm)
ratio_graph = None
ref_hist = hist_models['plot'].GetHistogram()
ratio_ref_hist = ref_hist.Clone()
turnOns = [None] * n_inputs
curves = [None] * n_inputs
for input_id in range(n_inputs):
turnOns[input_id] = turnOn_data[input_id][dm][wp][channel]['plot']
curves[input_id] = turnOns[input_id].eff
y_min, y_max = (0, 1)
y_title = 'Efficiency'
title = '{} {}{}'.format(channel, wp, dm_label)
plain_title = '{}_{}{}'.format(channel, wp, dm_plain_label)
main_pad, ratio_pad, title_controls = RootPlotting.CreateTwoPadLayout(canvas, ref_hist, ratio_ref_hist,
log_x=use_logx, title=title)
RootPlotting.ApplyAxisSetup(ref_hist, ratio_ref_hist, x_title=x_title, y_title=y_title,
ratio_y_title='Ratio', y_range=(y_min, y_max * 1.1), max_ratio=1.5)
legend = RootPlotting.CreateLegend(pos=(0.68, 0.28), size=(0.2, 0.15))
for input_id in range(n_inputs):
curve = curves[input_id]
try:
curve.Draw('SAME')
except:
continue
RootPlotting.ApplyDefaultLineStyle(curve, colors[input_id])
legend.AddEntry(curve, labels[input_id], 'PLE')
# if input_id < n_inputs - 1:
if input_id > 0:
ratio_graph = RootPlotting.CreateEfficiencyRatioGraph(turnOns[input_id].hist_passed,
turnOns[input_id].hist_total,
turnOns[0].hist_passed,
turnOns[0].hist_total)
if ratio_graph:
output_file.WriteTObject(ratio_graph, 'ratio_{}'.format(plain_title), 'Overwrite')
ratio_pad.cd()
ratio_color = colors[input_id] if n_inputs > 2 else ROOT.kBlack
RootPlotting.ApplyDefaultLineStyle(ratio_graph, ratio_color)
ratio_graph.Draw("0PE SAME")
main_pad.cd()
legend.Draw()
canvas.Update()
output_file.WriteTObject(canvas, 'canvas_{}'.format(plain_title), 'Overwrite')
RootPlotting.PrintAndClear(canvas, args.output + '.pdf', plain_title, plot_id, n_plots,
[ main_pad, ratio_pad ])
plot_id += 1
output_file.Close()