-
Notifications
You must be signed in to change notification settings - Fork 0
/
reweight_sstop_files.py
186 lines (147 loc) · 7.05 KB
/
reweight_sstop_files.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
import numpy as np
import os, sys
from optparse import OptionParser
from cross_sections_DM import *
from copy import copy
import ROOT
gcd = ROOT.gROOT.cd
processes = ['tt_exclusive', 'ttu_onshellV', 'ttu_offshellV']
# Nominal model parameters
mediator_masses = [1000, 1500, 2000, 2500, 3000]
mDM = 1.
def make_folder_if_not_exists(folder):
if not os.path.exists(folder):
try:
os.makedirs(folder)
except:
print "Could not create a folder ", folder
def reweight_root_file(mV, process, data):
global output_folder, systematics, BR_run
root_file_prefix = 'sstops_'
if systematics:
root_file_prefix = 'SystVar_sstops_'
# slice in columns
a_r = data[:,2]
g = data[:,3]
BR = data[:,0]
if 'tt_excl' in process:
xsection = data[:,9]
MC_cross_section = cross_sections_tt_excl[str(mV)]
elif 'onshell' in process:
xsection = data[:,8]
MC_cross_section = cross_sections_onshellV[str(mV)]
elif 'offshell' in process:
xsection = data[:,7]
MC_cross_section = cross_sections_offshellV[str(mV)]
else:
print 'Unknown process: ', process
sys.exit(1)
weights = xsection/MC_cross_section
input_root_file = ROOT.TFile(input_folder + '/' + root_file_prefix + process + '_mV' + str(mV) + '.root', "read")
for i in range(len(weights)):
if BR_run:
weight_suffix = '_a_r%.2f_BR%.2f' % (a_r[i], BR[i])
else:
weight_suffix = '_a_r%.2f_g%.2f' % (a_r[i], g[i])
# create the new root file with reweighted histograms
reweighted_root_file = ROOT.TFile(output_folder + '/' + root_file_prefix + process + '_mV' + str(mV) + weight_suffix + '.root', "recreate")
# get the histograms from the input root file
for instance in input_root_file.GetListOfKeys():
histogram = input_root_file.Get(instance.GetName())
histogram_weighted = histogram.Clone()
histogram_weighted.Scale(weights[i])
histogram_weighted.Write(instance.GetName())
reweighted_root_file.Close()
input_root_file.Close()
def sum_files(mV, data):
global output_folder, systematics, BR_run
root_file_prefix = 'sstops_'
if systematics:
root_file_prefix = 'SystVar_sstops_'
# slice in columns
a_r = data[:,2]
g = data[:,3]
BR = data[:,0]
if BR_run:
first_file_name = output_folder + '/' + root_file_prefix + processes[0] + '_mV' + str(mV) + '_a_r{0:.2f}'.format(a_r[0]) + '_BR{0:.2f}'.format(BR[0]) + '.root'
else:
first_file_name = output_folder + '/' + root_file_prefix + processes[0] + '_mV' + str(mV) + '_a_r{0:.2f}'.format(a_r[0]) + '_g{0:.2f}'.format(g[0]) + '.root'
first_file = ROOT.TFile(first_file_name, "read")
list_of_histograms = [histogram.GetName() for histogram in first_file.GetListOfKeys()]
first_file.Close()
for i in range(len(a_r)):
if BR_run:
weight_suffix = '_a_r%.2f_BR%.2f' % (a_r[i], BR[i])
else:
weight_suffix = '_a_r%.2f_g%.2f' % (a_r[i], g[i])
histograms_dict = {}
for process in processes:
reweighted_root_file = ROOT.TFile(output_folder + '/' + root_file_prefix + process + '_mV' + str(mV) + weight_suffix + '.root', "read")
# print 'Reading from file ', output_folder + '/' + root_file_prefix + process + '_mV' + str(mV) + weight_suffix + '.root'
gcd()
# sum the histograms and place in dictionary
for histogram_name in list_of_histograms:
histogram = reweighted_root_file.Get(histogram_name)
if histogram_name in histograms_dict:
histograms_dict[histogram_name] += copy(histogram.Clone())
else:
histograms_dict[histogram_name] = copy(histogram.Clone())
reweighted_root_file.Close()
# write histograms
summed_root_file = ROOT.TFile(output_folder + '/' + root_file_prefix + 'mV' + str(mV) + weight_suffix + '.root', "recreate")
# print 'Creating file ', output_folder + '/' + root_file_prefix + 'mV' + str(mV) + weight_suffix + '.root'
for histogram_name in list_of_histograms:
histograms_dict[histogram_name].Write(histogram_name)
summed_root_file.Close()
def clean_up_files():
global output_folder, systematics
root_file_prefix = 'sstops_'
if systematics:
root_file_prefix = 'SystVar_sstops_'
for process in processes:
os.system("rm " + output_folder + '/' + root_file_prefix + process + '*.root')
if __name__ == '__main__':
parser = OptionParser()
parser.add_option( "-o", "--output_folder", dest = "output_folder", default = 'reweighted_sstop_files/',
help = "set path to save weighted root files" )
parser.add_option( "-i", "--input_folder", dest= "input_folder", default = '../',
help = "set path with input root files to reweight" )
parser.add_option( "-t", "--input_table", dest= "input_table", default = 'big_table.csv',
help = "set a filename with the parameterization table (make_results_table.py output)" )
parser.add_option( "-s", "--syst", action = "store_true", dest = "systematics",
help = "Reweight systematic variations (SystVar_*.root)" )
parser.add_option( "-k", "--keep_subprocesses", action = "store_true", dest = "keep_subprocesses",
help = "Keep reweighted root files for all subprocesses" )
parser.add_option( "-c", "--clean_up_subprocesses", action = "store_true", dest = "clean_up_subprocesses",
help = "Only clean up reweighted root files for all subprocesses" )
parser.add_option( "-m", "--mediator_mass", dest = "mass", default = '',
help = "Choose the mediator mass in GeV (1000/1500/2000/2500/3000" )
parser.add_option( "-B", "--BR_parameters", action="store_true", dest="BR_run",
help="Use BR parameterisation instead of the nominal one")
( options, args ) = parser.parse_args()
output_folder = options.output_folder
make_folder_if_not_exists(output_folder)
input_folder = options.input_folder
input_table = options.input_table
systematics = options.systematics
BR_run = options.BR_run
if options.clean_up_subprocesses:
print 'Cleaning up reweighted subrpocess files...'
clean_up_files()
print 'Done.'
sys.exit(0)
whole_data = np.genfromtxt(input_table, delimiter=',')
#delete first row (variable names)
whole_data = np.delete(whole_data, (0), axis=0)
if options.mass:
mediator_masses = [ int(options.mass) ]
options.keep_subprocesses = True
# work with a single mV:
for mV in mediator_masses:
print 'Working with mV = ', mV
single_mV_data = whole_data[whole_data[:,5]==mV]
for process in processes:
reweight_root_file(mV, process, single_mV_data)
sum_files(mV, single_mV_data)
if not options.keep_subprocesses:
clean_up_files()