| 1 |
#!/usr/bin/env python
|
| 2 |
from samplesinfo import sample
|
| 3 |
from printcolor import printc
|
| 4 |
import pickle
|
| 5 |
import sys
|
| 6 |
import os
|
| 7 |
import ROOT
|
| 8 |
import shutil
|
| 9 |
from ROOT import TFile
|
| 10 |
import ROOT
|
| 11 |
from array import array
|
| 12 |
import warnings
|
| 13 |
warnings.filterwarnings( action='ignore', category=RuntimeWarning, message='creating converter.*' )
|
| 14 |
|
| 15 |
|
| 16 |
#usage: ./write_systematic.py path
|
| 17 |
|
| 18 |
path=sys.argv[1]
|
| 19 |
|
| 20 |
#load info
|
| 21 |
infofile = open(path+'/samples.info','r')
|
| 22 |
info = pickle.load(infofile)
|
| 23 |
infofile.close()
|
| 24 |
#os.mkdir(path+'/sys')
|
| 25 |
|
| 26 |
for job in info:
|
| 27 |
if job.type != 'DATA':
|
| 28 |
print '\t - %s' %(job.name)
|
| 29 |
input = TFile.Open(job.getpath(),'read')
|
| 30 |
output = TFile.Open(job.path+'/sys/'+job.prefix+job.identifier+'.root','recreate')
|
| 31 |
|
| 32 |
input.cd()
|
| 33 |
obj = ROOT.TObject
|
| 34 |
for key in ROOT.gDirectory.GetListOfKeys():
|
| 35 |
input.cd()
|
| 36 |
obj = key.ReadObj()
|
| 37 |
#print obj.GetName()
|
| 38 |
if obj.GetName() == job.tree:
|
| 39 |
continue
|
| 40 |
output.cd()
|
| 41 |
#print key.GetName()
|
| 42 |
obj.Write(key.GetName())
|
| 43 |
|
| 44 |
tree = input.Get(job.tree)
|
| 45 |
nEntries = tree.GetEntries()
|
| 46 |
|
| 47 |
job.addpath('/sys')
|
| 48 |
job.SYS = ['Nominal','JER_up','JER_down','JES_up','JES_down','beff_up','beff_down','bmis_up','bmis_down']
|
| 49 |
|
| 50 |
output.cd()
|
| 51 |
newtree = tree.CloneTree(0)
|
| 52 |
|
| 53 |
hJ0 = ROOT.TLorentzVector()
|
| 54 |
hJ1 = ROOT.TLorentzVector()
|
| 55 |
|
| 56 |
#JER branches
|
| 57 |
hJet_pt_JER_up = array('f',[0]*2)
|
| 58 |
newtree.Branch('hJet_pt_JER_up',hJet_pt_JER_up,'hJet_pt_JER_up[2]/F')
|
| 59 |
hJet_pt_JER_down = array('f',[0]*2)
|
| 60 |
newtree.Branch('hJet_pt_JER_down',hJet_pt_JER_down,'hJet_pt_JER_down[2]/F')
|
| 61 |
hJet_e_JER_up = array('f',[0]*2)
|
| 62 |
newtree.Branch('hJet_e_JER_up',hJet_e_JER_up,'hJet_e_JER_up[2]/F')
|
| 63 |
hJet_e_JER_down = array('f',[0]*2)
|
| 64 |
newtree.Branch('hJet_e_JER_down',hJet_e_JER_down,'hJet_e_JER_down[2]/F')
|
| 65 |
H_JER = array('f',[0]*4)
|
| 66 |
newtree.Branch('H_JER',H_JER,'mass_up:mass_down:pt_up:pt_down/F')
|
| 67 |
|
| 68 |
#JES branches
|
| 69 |
hJet_pt_JES_up = array('f',[0]*2)
|
| 70 |
newtree.Branch('hJet_pt_JES_up',hJet_pt_JES_up,'hJet_pt_JES_up[2]/F')
|
| 71 |
hJet_pt_JES_down = array('f',[0]*2)
|
| 72 |
newtree.Branch('hJet_pt_JES_down',hJet_pt_JES_down,'hJet_pt_JES_down[2]/F')
|
| 73 |
hJet_e_JES_up = array('f',[0]*2)
|
| 74 |
newtree.Branch('hJet_e_JES_up',hJet_e_JES_up,'hJet_e_JES_up[2]/F')
|
| 75 |
hJet_e_JES_down = array('f',[0]*2)
|
| 76 |
newtree.Branch('hJet_e_JES_down',hJet_e_JES_down,'hJet_e_JES_down[2]/F')
|
| 77 |
H_JES = array('f',[0]*4)
|
| 78 |
newtree.Branch('H_JES',H_JES,'mass_up:mass_down:pt_up:pt_down/F')
|
| 79 |
|
| 80 |
#Add training Flag
|
| 81 |
EventForTraining = array('f',[0])
|
| 82 |
newtree.Branch('EventForTraining',EventForTraining,'EventForTraining/F')
|
| 83 |
|
| 84 |
#iter=0
|
| 85 |
|
| 86 |
TFlag=ROOT.TTreeFormula("EventForTraining","EVENT.event%2",tree)
|
| 87 |
|
| 88 |
for entry in range(0,nEntries):
|
| 89 |
tree.GetEntry(entry)
|
| 90 |
|
| 91 |
#fill training flag
|
| 92 |
#iter+=1
|
| 93 |
#if (iter%2==0):
|
| 94 |
# EventForTraining[0]=1
|
| 95 |
#else:
|
| 96 |
# EventForTraining[0]=0
|
| 97 |
#iter+=1
|
| 98 |
|
| 99 |
EventForTraining[0]=int(not TFlag.EvalInstance())
|
| 100 |
|
| 101 |
#get
|
| 102 |
hJet_pt0 = tree.hJet_pt[0]
|
| 103 |
hJet_pt1 = tree.hJet_pt[1]
|
| 104 |
hJet_eta0 = tree.hJet_eta[0]
|
| 105 |
hJet_eta1 = tree.hJet_eta[1]
|
| 106 |
hJet_genPt0 = tree.hJet_genPt[0]
|
| 107 |
hJet_genPt1 = tree.hJet_genPt[1]
|
| 108 |
hJet_e0 = tree.hJet_e[0]
|
| 109 |
hJet_e1 = tree.hJet_e[1]
|
| 110 |
hJet_phi0 = tree.hJet_phi[0]
|
| 111 |
hJet_phi1 = tree.hJet_phi[1]
|
| 112 |
hJet_JECUnc0 = tree.hJet_JECUnc[0]
|
| 113 |
hJet_JECUnc1 = tree.hJet_JECUnc[1]
|
| 114 |
|
| 115 |
|
| 116 |
for updown in ['up','down']:
|
| 117 |
#JER
|
| 118 |
if updown == 'up':
|
| 119 |
inner = 0.06
|
| 120 |
outer = 0.1
|
| 121 |
if updown == 'down':
|
| 122 |
inner = -0.06
|
| 123 |
outer = -0.1
|
| 124 |
#Calculate
|
| 125 |
if abs(hJet_eta0)<1.1: res0 = inner
|
| 126 |
else: res0 = outer
|
| 127 |
if abs(hJet_eta1)<1.1: res1 = inner
|
| 128 |
else: res1 = outer
|
| 129 |
rPt0 = hJet_pt0 + (hJet_pt0-hJet_genPt0)*res0
|
| 130 |
rPt1 = hJet_pt1 + (hJet_pt1-hJet_genPt1)*res1
|
| 131 |
rE0 = hJet_e0*rPt0/hJet_pt0
|
| 132 |
rE1 = hJet_e1*rPt1/hJet_pt1
|
| 133 |
hJ0.SetPtEtaPhiE(rPt0,hJet_eta0,hJet_phi0,rE0)
|
| 134 |
hJ1.SetPtEtaPhiE(rPt1,hJet_eta1,hJet_phi1,rE1)
|
| 135 |
#Set
|
| 136 |
if updown == 'up':
|
| 137 |
hJet_pt_JER_up[0]=rPt0
|
| 138 |
hJet_pt_JER_up[1]=rPt1
|
| 139 |
hJet_e_JER_up[0]=rE0
|
| 140 |
hJet_e_JER_up[1]=rE1
|
| 141 |
H_JER[0]=(hJ0+hJ1).M()
|
| 142 |
H_JER[2]=(hJ0+hJ1).Pt()
|
| 143 |
if updown == 'down':
|
| 144 |
hJet_pt_JER_down[0]=rPt0
|
| 145 |
hJet_pt_JER_down[1]=rPt1
|
| 146 |
hJet_e_JER_down[0]=rE0
|
| 147 |
hJet_e_JER_down[1]=rE1
|
| 148 |
H_JER[1]=(hJ0+hJ1).M()
|
| 149 |
H_JER[3]=(hJ0+hJ1).Pt()
|
| 150 |
|
| 151 |
#JES
|
| 152 |
if updown == 'up':
|
| 153 |
variation=1
|
| 154 |
if updown == 'down':
|
| 155 |
variation=-1
|
| 156 |
#calculate
|
| 157 |
rPt0 = hJet_pt0*(1+variation*hJet_JECUnc0)
|
| 158 |
rPt1 = hJet_pt1*(1+variation*hJet_JECUnc1)
|
| 159 |
rE0 = hJet_e0*(1+variation*hJet_JECUnc0)
|
| 160 |
rE1 = hJet_e1*(1+variation*hJet_JECUnc1)
|
| 161 |
hJ0.SetPtEtaPhiE(rPt0,hJet_eta0,hJet_phi0,rE0)
|
| 162 |
hJ1.SetPtEtaPhiE(rPt1,hJet_eta1,hJet_phi1,rE1)
|
| 163 |
#Fill
|
| 164 |
if updown == 'up':
|
| 165 |
hJet_pt_JES_up[0]=rPt0
|
| 166 |
hJet_pt_JES_up[1]=rPt1
|
| 167 |
hJet_e_JES_up[0]=rE0
|
| 168 |
hJet_e_JES_up[1]=rE1
|
| 169 |
H_JES[0]=(hJ0+hJ1).M()
|
| 170 |
H_JES[2]=(hJ0+hJ1).Pt()
|
| 171 |
if updown == 'down':
|
| 172 |
hJet_pt_JES_down[0]=rPt0
|
| 173 |
hJet_pt_JES_down[1]=rPt1
|
| 174 |
hJet_e_JES_down[0]=rE0
|
| 175 |
hJet_e_JES_down[1]=rE1
|
| 176 |
H_JES[1]=(hJ0+hJ1).M()
|
| 177 |
H_JES[3]=(hJ0+hJ1).Pt()
|
| 178 |
|
| 179 |
newtree.Fill()
|
| 180 |
|
| 181 |
newtree.AutoSave()
|
| 182 |
output.Close()
|
| 183 |
|
| 184 |
else: #(is data)
|
| 185 |
|
| 186 |
print '\t - %s' %(job.name)
|
| 187 |
input = TFile.Open(job.getpath(),'read')
|
| 188 |
output = TFile.Open(job.path+'/sys/'+job.prefix+job.identifier+'.root','recreate')
|
| 189 |
#input.cd()
|
| 190 |
obj = ROOT.TObject
|
| 191 |
for key in ROOT.gDirectory.GetListOfKeys():
|
| 192 |
input.cd()
|
| 193 |
obj = key.ReadObj()
|
| 194 |
if obj.GetName() == job.tree:
|
| 195 |
continue
|
| 196 |
output.cd()
|
| 197 |
obj.Write(key.GetName())
|
| 198 |
#output.cd()
|
| 199 |
tree = input.Get(job.tree)
|
| 200 |
nEntries = tree.GetEntries()
|
| 201 |
job.addpath('/sys')
|
| 202 |
output.cd()
|
| 203 |
newtree = tree.CloneTree(0)
|
| 204 |
#Add training Flag
|
| 205 |
EventForTraining = array('f',[0])
|
| 206 |
newtree.Branch('EventForTraining',EventForTraining,'EventForTraining/F')
|
| 207 |
for entry in range(0,nEntries):
|
| 208 |
tree.GetEntry(entry)
|
| 209 |
EventForTraining[0]=0
|
| 210 |
newtree.Fill()
|
| 211 |
newtree.AutoSave()
|
| 212 |
output.Close()
|
| 213 |
|
| 214 |
|
| 215 |
#dump info
|
| 216 |
infofile = open(path+'/sys'+'/samples.info','w')
|
| 217 |
pickle.dump(info,infofile)
|
| 218 |
infofile.close() |
