PythonGitHub/systematic_yield_extraction.py at master · chrisjcc/PythonGitHub · 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
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
#!/usr/bin/env python


############################################
#                                          #
#   Write by: Christian Contreras-Campana  #
#   email: christian.contreras@desy.de     #
#   Date: 20.10.2017                       #
#                                          #
############################################


import collections
import re
import os

from optparse import OptionParser

''' CODE PURPOSE: To create a comma separated value (csv) file which contains the break down of the nominal values for systematics included below \n
    RUN EXAMPLE: python systematic_yield_extraction.py -f EventYields_selectionRoot_mvaEventA/Nominal/combined/lumiCorrected_events_step7.txt'''

# ---- Command line option parsing
parser = OptionParser()
parser.add_option("-f", "--file", dest="filename",
                  help="write report to FILE", metavar="FILE")
parser.add_option("-q", "--quiet",
                  action="store_false", dest="verbose", default=True,
                  help="don't print status messages to stdout")

(options, args) = parser.parse_args()

# ---- List of systematic uncertainties to process
systematics = [
"BTAGDISCR_BPURITY",
"BTAGDISCR_BSTAT1",
"BTAGDISCR_BSTAT2",
"BTAGDISCR_CERR1",
"BTAGDISCR_CERR2",
"BTAGDISCR_LPURITY",
"BTAGDISCR_LSTAT1",
"BTAGDISCR_LSTAT2",
"JER",
"JES",
"LEPT",
"MEFACSCALE",
"MERENSCALE",
"PU",
"TRIG",
]

syst_nominal={}

# ---- Set up files and directory paths
sub_directories = str(options.filename).split('/')
input_filename  = os.path.basename(options.filename)
output_filename = os.path.splitext(input_filename)[0]+".csv"

# ---- Set input and output file handlers
f_out = open(output_filename, 'a')
f_nom = open(options.filename,'r')

# ---- Read in file from the nominal case
for line in f_nom.xreadlines():

   nominal_values = re.split(':|    \+\-',line)

   if nominal_values[0] == 'Total background' or nominal_values[0] == 'Total MC':
         continue
   try:
      syst_nominal[nominal_values[0]].append({'Nominal': [nominal_values[1]]})
   except KeyError:
      syst_nominal[nominal_values[0]] = {'Nominal': [nominal_values[1]]}

# ----- Loop through list of systematic to extract information from text
for systematic in systematics:

   # ---- Read in file from the up and down variation case
   f_nom_up   = open(sub_directories[0]+"/"+systematic+'_UP'+"/"+sub_directories[2]+"/"+input_filename,'r')
   f_nom_down = open(sub_directories[0]+"/"+systematic+'_DOWN'+"/"+sub_directories[2]+"/"+input_filename,'r')

   for line in f_nom_up.xreadlines():
      nominal_values = re.split(':|    \+\-',line)

      if nominal_values[0] == 'Total background' or nominal_values[0] == 'Total MC':
         continue

      try:
         syst_nominal[nominal_values[0]][systematic].append([nominal_values[1]])
      except KeyError:
         syst_nominal[nominal_values[0]][systematic] = [nominal_values[1]]

   for line in f_nom_down.xreadlines():
      nominal_values = re.split(':|    \+\-',line)

      if nominal_values[0] == 'Total background' or nominal_values[0] == 'Total MC':
         continue

      syst_nominal[nominal_values[0]][systematic].append(nominal_values[1])

   f_nom.close()
   f_nom_down.close()
   f_nom_up.close()

# ---- Print to screen and store information in output file
print 'PROCESS,',
f_out.write('PROCESS,')
for systematic in syst_nominal.values()[0].keys():
   if systematic != 'Nominal':
      print systematic+"_UP,"+systematic+"_DOWN,",
      f_out.write(systematic+"_UP,"+systematic+"_DOWN,")
   else:
      print systematic+",",
      f_out.write(systematic+",")

print
f_out.write("\n")

for process in syst_nominal.keys():
   print process+",",
   f_out.write(process+",")

   for values in syst_nominal[process].values():
      print ",".join(values),',',
      f_out.write(",".join(values)+",")
   print
   f_out.write("\n")