ViewVC Help
View File | Revision Log | Show Annotations | Root Listing
root/cvsroot/UserCode/VHbb/python/write_regression_systematics.py
Revision: 1.3
Committed: Mon Jun 4 14:55:36 2012 UTC (12 years, 11 months ago) by nmohr
Content type: text/x-python
Branch: MAIN
Changes since 1.2: +1 -0 lines
Log Message:
Fix

File Contents

# User Rev Content
1 nmohr 1.1 #!/usr/bin/env python
2     from samplesclass import sample
3     from printcolor import printc
4     import pickle
5     import sys
6     import os
7     import ROOT
8     import math
9     import shutil
10     from ROOT import TFile
11     import ROOT
12     from array import array
13     import warnings
14     warnings.filterwarnings( action='ignore', category=RuntimeWarning, message='creating converter.*' )
15    
16    
17     #usage: ./write_regression_systematic.py path
18    
19     path=sys.argv[1]
20    
21     #load info
22     infofile = open(path+'/samples.info','r')
23     info = pickle.load(infofile)
24     infofile.close()
25     #os.mkdir(path+'/sys')
26    
27     def deltaPhi(phi1, phi2):
28     result = phi1 - phi2
29     while (result > math.pi): result -= 2*math.pi
30     while (result <= -math.pi): result += 2*math.pi
31     return result
32    
33     def corrCSV(btag, csv, flav):
34     if(csv < 0.): return csv
35     if(csv > 1.): return csv;
36     if(flav == 0): return csv;
37     if(math.fabs(flav) == 5): return btag.ib.Eval(csv)
38     if(math.fabs(flav) == 4): return btag.ic.Eval(csv)
39     if(math.fabs(flav) != 4 and math.fabs(flav) != 5): return btag.il.Eval(csv)
40     return -10000
41    
42    
43     for job in info:
44     #print job.name
45     #if job.name != 'ZH125': continue
46     ROOT.gROOT.ProcessLine(
47     "struct H {\
48     int HiggsFlag;\
49     float mass;\
50     float pt;\
51     float eta;\
52     float phi;\
53     float dR;\
54     float dPhi;\
55     float dEta;\
56     } ;"
57     )
58     ROOT.gROOT.LoadMacro('../interface/btagshape.h+')
59     from ROOT import BTagShape
60     btagNom = BTagShape("../data/csvdiscr.root")
61     btagNom.computeFunctions()
62     btagUp = BTagShape("../data/csvdiscr.root")
63     btagUp.computeFunctions(+1.,0.)
64     btagDown = BTagShape("../data/csvdiscr.root")
65     btagDown.computeFunctions(-1.,0.)
66     btagFUp = BTagShape("../data/csvdiscr.root")
67     btagFUp.computeFunctions(0.,+1.)
68     btagFDown = BTagShape("../data/csvdiscr.root")
69     btagFDown.computeFunctions(0.,-1.)
70    
71     print '\t - %s' %(job.name)
72     input = TFile.Open(job.getpath(),'read')
73     output = TFile.Open(job.path+'/sys/'+job.prefix+job.identifier+'.root','recreate')
74    
75     input.cd()
76     obj = ROOT.TObject
77     for key in ROOT.gDirectory.GetListOfKeys():
78     input.cd()
79     obj = key.ReadObj()
80     #print obj.GetName()
81     if obj.GetName() == job.tree:
82     continue
83     output.cd()
84     #print key.GetName()
85     obj.Write(key.GetName())
86    
87 nmohr 1.3 input.cd()
88 nmohr 1.1 tree = input.Get(job.tree)
89     nEntries = tree.GetEntries()
90    
91     job.addpath('/sys')
92     if job.type != 'DATA':
93     job.SYS = ['Nominal','JER_up','JER_down','JES_up','JES_down','beff_up','beff_down','bmis_up','bmis_down']
94    
95     H = ROOT.H()
96     HNoReg = ROOT.H()
97     tree.SetBranchStatus('H',0)
98     output.cd()
99     newtree = tree.CloneTree(0)
100    
101     hJ0 = ROOT.TLorentzVector()
102     hJ1 = ROOT.TLorentzVector()
103    
104     regWeight = "../data/MVA_BDT_REG_May23.weights.xml"
105     regDict = {"Jet_pt": "hJet_pt", "Jet_eta": "hJet_eta", "Jet_e": "hJet_e", "Jet_JECUnc": "hJet_JECUnc", "Jet_chf": "hJet_chf","Jet_nconstituents": "hJet_nconstituents", "Jet_vtxPt": "hJet_vtxPt", "Jet_vtx3dL": "hJet_vtx3dL", "Jet_vtx3deL": "hJet_vtx3deL"}
106     regVars = ["Jet_pt","Jet_eta","Jet_e","Jet_JECUnc", "Jet_chf","Jet_nconstituents", "Jet_vtxPt", "Jet_vtx3dL", "Jet_vtx3deL"]
107    
108    
109     #Regression branches
110     applyRegression = True
111     hJet_pt = array('f',[0]*2)
112     hJet_e = array('f',[0]*2)
113     newtree.Branch( 'H', H , 'HiggsFlag/I:mass/F:pt/F:eta:phi/F:dR/F:dPhi/F:dEta/F' )
114     newtree.Branch( 'HNoReg', HNoReg , 'HiggsFlag/I:mass/F:pt/F:eta:phi/F:dR/F:dPhi/F:dEta/F' )
115     Event = array('f',[0])
116     METet = array('f',[0])
117     rho25 = array('f',[0])
118     METphi = array('f',[0])
119     fRho25 = ROOT.TTreeFormula("rho25",'rho25',tree)
120     fEvent = ROOT.TTreeFormula("Event",'EVENT.event',tree)
121     fMETet = ROOT.TTreeFormula("METet",'METnoPU.et',tree)
122     fMETphi = ROOT.TTreeFormula("METphi",'METnoPU.phi',tree)
123     hJet_MET_dPhi = array('f',[0]*2)
124     hJet_regWeight = array('f',[0]*2)
125     hJet_MET_dPhiArray = [array('f',[0]),array('f',[0])]
126     newtree.Branch('hJet_MET_dPhi',hJet_MET_dPhi,'hJet_MET_dPhi[2]/F')
127     newtree.Branch('hJet_regWeight',hJet_regWeight,'hJet_regWeight[2]/F')
128     readerJet0 = ROOT.TMVA.Reader("!Color:!Silent" )
129     readerJet1 = ROOT.TMVA.Reader("!Color:!Silent" )
130    
131     theForms = {}
132     theVars0 = {}
133     for var in regVars:
134     theVars0[var] = array( 'f', [ 0 ] )
135     readerJet0.AddVariable(var,theVars0[var])
136     theForms['form_reg_%s_0'%(regDict[var])] = ROOT.TTreeFormula("form_reg_%s_0"%(regDict[var]),'%s[0]' %(regDict[var]),tree)
137     readerJet0.AddVariable( "Jet_MET_dPhi", hJet_MET_dPhiArray[0] )
138     readerJet0.AddVariable( "METet", METet )
139     readerJet0.AddVariable( "rho25", rho25 )
140    
141     theVars1 = {}
142     for var in regVars:
143     theVars1[var] = array( 'f', [ 0 ] )
144     readerJet1.AddVariable(var,theVars1[var])
145     theForms['form_reg_%s_1'%(regDict[var])] = ROOT.TTreeFormula("form_reg_%s_1"%(regDict[var]),'%s[1]' %(regDict[var]),tree)
146     readerJet1.AddVariable( "Jet_MET_dPhi", hJet_MET_dPhiArray[1] )
147     readerJet1.AddVariable( "METet", METet )
148     readerJet1.AddVariable( "rho25", rho25 )
149     readerJet0.BookMVA( "jet0Regression", regWeight );
150     readerJet1.BookMVA( "jet1Regression", regWeight );
151    
152    
153     if job.type != 'DATA':
154     #CSV branches
155     hJet_flavour = array('f',[0]*2)
156     hJet_csv = array('f',[0]*2)
157     hJet_csvOld = array('f',[0]*2)
158     hJet_csvUp = array('f',[0]*2)
159     hJet_csvDown = array('f',[0]*2)
160     hJet_csvFUp = array('f',[0]*2)
161     hJet_csvFDown = array('f',[0]*2)
162     newtree.Branch('hJet_csvOld',hJet_csvOld,'hJet_csvOld[2]/F')
163     newtree.Branch('hJet_csvUp',hJet_csvUp,'hJet_csvUp[2]/F')
164     newtree.Branch('hJet_csvDown',hJet_csvDown,'hJet_csvDown[2]/F')
165     newtree.Branch('hJet_csvFUp',hJet_csvFUp,'hJet_csvFUp[2]/F')
166     newtree.Branch('hJet_csvFDown',hJet_csvFDown,'hJet_csvFDown[2]/F')
167    
168     #JER branches
169     hJet_pt_JER_up = array('f',[0]*2)
170     newtree.Branch('hJet_pt_JER_up',hJet_pt_JER_up,'hJet_pt_JER_up[2]/F')
171     hJet_pt_JER_down = array('f',[0]*2)
172     newtree.Branch('hJet_pt_JER_down',hJet_pt_JER_down,'hJet_pt_JER_down[2]/F')
173     hJet_e_JER_up = array('f',[0]*2)
174     newtree.Branch('hJet_e_JER_up',hJet_e_JER_up,'hJet_e_JER_up[2]/F')
175     hJet_e_JER_down = array('f',[0]*2)
176     newtree.Branch('hJet_e_JER_down',hJet_e_JER_down,'hJet_e_JER_down[2]/F')
177     H_JER = array('f',[0]*4)
178     newtree.Branch('H_JER',H_JER,'mass_up:mass_down:pt_up:pt_down/F')
179    
180     #JES branches
181     hJet_pt_JES_up = array('f',[0]*2)
182     newtree.Branch('hJet_pt_JES_up',hJet_pt_JES_up,'hJet_pt_JES_up[2]/F')
183     hJet_pt_JES_down = array('f',[0]*2)
184     newtree.Branch('hJet_pt_JES_down',hJet_pt_JES_down,'hJet_pt_JES_down[2]/F')
185     hJet_e_JES_up = array('f',[0]*2)
186     newtree.Branch('hJet_e_JES_up',hJet_e_JES_up,'hJet_e_JES_up[2]/F')
187     hJet_e_JES_down = array('f',[0]*2)
188     newtree.Branch('hJet_e_JES_down',hJet_e_JES_down,'hJet_e_JES_down[2]/F')
189     H_JES = array('f',[0]*4)
190     newtree.Branch('H_JES',H_JES,'mass_up:mass_down:pt_up:pt_down/F')
191    
192     #Add training Flag
193     EventForTraining = array('f',[0])
194     newtree.Branch('EventForTraining',EventForTraining,'EventForTraining/F')
195    
196    
197     #iter=0
198    
199     TFlag=ROOT.TTreeFormula("EventForTraining","EVENT.event%2",tree)
200    
201     for entry in range(0,nEntries):
202     tree.GetEntry(entry)
203    
204     #fill training flag
205     #iter+=1
206     #if (iter%2==0):
207     # EventForTraining[0]=1
208     #else:
209     # EventForTraining[0]=0
210     #iter+=1
211    
212     if job.type != 'DATA':
213     EventForTraining[0]=int(not TFlag.EvalInstance())
214     else:
215     EventForTraining[0]=0
216    
217     #get
218     hJet_pt = tree.hJet_pt
219     hJet_e = tree.hJet_e
220     hJet_pt0 = tree.hJet_pt[0]
221     hJet_pt1 = tree.hJet_pt[1]
222     hJet_eta0 = tree.hJet_eta[0]
223     hJet_eta1 = tree.hJet_eta[1]
224     hJet_genPt0 = tree.hJet_genPt[0]
225     hJet_genPt1 = tree.hJet_genPt[1]
226     hJet_e0 = tree.hJet_e[0]
227     hJet_e1 = tree.hJet_e[1]
228     hJet_phi0 = tree.hJet_phi[0]
229     hJet_phi1 = tree.hJet_phi[1]
230     hJet_JECUnc0 = tree.hJet_JECUnc[0]
231     hJet_JECUnc1 = tree.hJet_JECUnc[1]
232    
233     Event[0]=fEvent.EvalInstance()
234     METet[0]=fMETet.EvalInstance()
235     rho25[0]=fRho25.EvalInstance()
236     METphi[0]=fMETphi.EvalInstance()
237     for key, value in regDict.items():
238     theVars0[key][0] = theForms["form_reg_%s_0" %(value)].EvalInstance()
239     theVars1[key][0] = theForms["form_reg_%s_1" %(value)].EvalInstance()
240     for i in range(2):
241     hJet_MET_dPhi[i] = deltaPhi(METphi[0],tree.hJet_phi[i])
242     hJet_MET_dPhiArray[i][0] = deltaPhi(METphi[0],tree.hJet_phi[i])
243    
244     if applyRegression:
245     hJ0.SetPtEtaPhiE(hJet_pt0,hJet_eta0,hJet_phi0,hJet_e0)
246     hJ1.SetPtEtaPhiE(hJet_pt1,hJet_eta1,hJet_phi1,hJet_e1)
247     HNoReg.HiggsFlag = 1
248     HNoReg.mass = (hJ0+hJ1).M()
249     HNoReg.pt = (hJ0+hJ1).Pt()
250     HNoReg.eta = (hJ0+hJ1).Eta()
251     HNoReg.phi = (hJ0+hJ1).Phi()
252     HNoReg.dR = hJ0.DeltaR(hJ1)
253     HNoReg.dPhi = hJ0.DeltaPhi(hJ1)
254     HNoReg.dEta = abs(hJ0.Eta()-hJ1.Eta())
255     rPt0 = readerJet0.EvaluateRegression( "jet0Regression" )[0]
256     rPt1 = readerJet1.EvaluateRegression( "jet1Regression" )[0]
257     hJet_regWeight[0] = rPt0/hJet_pt0
258     hJet_regWeight[1] = rPt1/hJet_pt1
259     rE0 = hJet_e0*hJet_regWeight[0]
260     rE1 = hJet_e1*hJet_regWeight[1]
261     hJ0.SetPtEtaPhiE(rPt0,hJet_eta0,hJet_phi0,rE0)
262     hJ1.SetPtEtaPhiE(rPt1,hJet_eta1,hJet_phi1,rE1)
263     tree.hJet_pt[0] = rPt0
264     tree.hJet_pt[1] = rPt1
265     tree.hJet_e[0] = rE0
266     tree.hJet_e[1] = rE1
267     H.HiggsFlag = 1
268     H.mass = (hJ0+hJ1).M()
269     H.pt = (hJ0+hJ1).Pt()
270     H.eta = (hJ0+hJ1).Eta()
271     H.phi = (hJ0+hJ1).Phi()
272     H.dR = hJ0.DeltaR(hJ1)
273     H.dPhi = hJ0.DeltaPhi(hJ1)
274     H.dEta = abs(hJ0.Eta()-hJ1.Eta())
275     if hJet_regWeight[0] > 5. or hJet_regWeight[1] > 5.:
276     print 'MET %.2f' %(METet[0])
277     print 'rho25 %.2f' %(rho25[0])
278     for key, value in regDict.items():
279     print '%s 0: %.2f'%(key, theVars0[key][0])
280     print '%s 0: %.2f'%(key, theVars1[key][0])
281     for i in range(2):
282     print 'dPhi %.0f %.2f' %(i,hJet_MET_dPhiArray[i][0])
283     print 'corr 0 %.2f' %(hJet_regWeight[0])
284     print 'corr 1 %.2f' %(hJet_regWeight[1])
285     print 'Event %.0f' %(Event[0])
286     print 'rPt0 %.2f' %(rPt0)
287     print 'rPt1 %.2f' %(rPt1)
288     print 'rE0 %.2f' %(rE0)
289     print 'rE1 %.2f' %(rE1)
290     print 'Mass %.2f' %(H.mass)
291    
292     if job.type == 'DATA':
293     newtree.Fill()
294     continue
295    
296     for i in range(2):
297     flavour = tree.hJet_flavour[i]
298     csv = tree.hJet_csv[i]
299     hJet_csvOld[i] = csv
300     tree.hJet_csv[i] = corrCSV(btagNom,csv,flavour)
301     hJet_csvDown[i] = corrCSV(btagDown,csv,flavour)
302     hJet_csvUp[i] = corrCSV(btagUp,csv,flavour)
303     hJet_csvFDown[i] = corrCSV(btagFDown,csv,flavour)
304     hJet_csvFUp[i] = corrCSV(btagFUp,csv,flavour)
305    
306     for updown in ['up','down']:
307     #JER
308     if updown == 'up':
309     inner = 0.06
310     outer = 0.1
311     if updown == 'down':
312     inner = -0.06
313     outer = -0.1
314     #Calculate
315     if abs(hJet_eta0)<1.1: res0 = inner
316     else: res0 = outer
317     if abs(hJet_eta1)<1.1: res1 = inner
318     else: res1 = outer
319     rPt0 = hJet_pt0 + (hJet_pt0-hJet_genPt0)*res0
320     rPt1 = hJet_pt1 + (hJet_pt1-hJet_genPt1)*res1
321     rE0 = hJet_e0*rPt0/hJet_pt0
322     rE1 = hJet_e1*rPt1/hJet_pt1
323     if applyRegression:
324     theVars0['Jet_pt'][0] = rPt0
325     theVars1['Jet_pt'][0] = rPt1
326     theVars0['Jet_e'][0] = rE0
327     theVars1['Jet_e'][0] = rE1
328     rPt0 = readerJet0.EvaluateRegression( "jet0Regression" )[0]
329     rPt1 = readerJet1.EvaluateRegression( "jet1Regression" )[0]
330     rE0 = hJet_e0*rPt0/hJet_pt0
331     rE1 = hJet_e1*rPt1/hJet_pt1
332     hJ0.SetPtEtaPhiE(rPt0,hJet_eta0,hJet_phi0,rE0)
333     hJ1.SetPtEtaPhiE(rPt1,hJet_eta1,hJet_phi1,rE1)
334     #Set
335     if updown == 'up':
336     hJet_pt_JER_up[0]=rPt0
337     hJet_pt_JER_up[1]=rPt1
338     hJet_e_JER_up[0]=rE0
339     hJet_e_JER_up[1]=rE1
340     H_JER[0]=(hJ0+hJ1).M()
341     H_JER[2]=(hJ0+hJ1).Pt()
342     if updown == 'down':
343     hJet_pt_JER_down[0]=rPt0
344     hJet_pt_JER_down[1]=rPt1
345     hJet_e_JER_down[0]=rE0
346     hJet_e_JER_down[1]=rE1
347     H_JER[1]=(hJ0+hJ1).M()
348     H_JER[3]=(hJ0+hJ1).Pt()
349    
350     #JES
351     if updown == 'up':
352     variation=1
353     if updown == 'down':
354     variation=-1
355     #calculate
356     rPt0 = hJet_pt0*(1+variation*hJet_JECUnc0)
357     rPt1 = hJet_pt1*(1+variation*hJet_JECUnc1)
358     rE0 = hJet_e0*(1+variation*hJet_JECUnc0)
359     rE1 = hJet_e1*(1+variation*hJet_JECUnc1)
360     if applyRegression:
361     theVars0['Jet_pt'][0] = rPt0
362     theVars1['Jet_pt'][0] = rPt1
363     theVars0['Jet_e'][0] = rE0
364     theVars1['Jet_e'][0] = rE1
365     rPt0 = readerJet0.EvaluateRegression( "jet0Regression" )[0]
366     rPt1 = readerJet1.EvaluateRegression( "jet1Regression" )[0]
367     rE0 = hJet_e0*rPt0/hJet_pt0
368     rE1 = hJet_e1*rPt1/hJet_pt1
369     hJ0.SetPtEtaPhiE(rPt0,hJet_eta0,hJet_phi0,rE0)
370     hJ1.SetPtEtaPhiE(rPt1,hJet_eta1,hJet_phi1,rE1)
371     #Fill
372     if updown == 'up':
373     hJet_pt_JES_up[0]=rPt0
374     hJet_pt_JES_up[1]=rPt1
375     hJet_e_JES_up[0]=rE0
376     hJet_e_JES_up[1]=rE1
377     H_JES[0]=(hJ0+hJ1).M()
378     H_JES[2]=(hJ0+hJ1).Pt()
379     if updown == 'down':
380     hJet_pt_JES_down[0]=rPt0
381     hJet_pt_JES_down[1]=rPt1
382     hJet_e_JES_down[0]=rE0
383     hJet_e_JES_down[1]=rE1
384     H_JES[1]=(hJ0+hJ1).M()
385     H_JES[3]=(hJ0+hJ1).Pt()
386    
387     newtree.Fill()
388    
389     newtree.AutoSave()
390     output.Close()
391    
392     #dump info
393     infofile = open(path+'/sys'+'/samples.info','w')
394     pickle.dump(info,infofile)
395     infofile.close()