ViewVC Help
View File | Revision Log | Show Annotations | Root Listing
root/cvsroot/UserCode/VHbb/python/write_regression_systematics.py
Revision: 1.1
Committed: Wed May 30 00:28:47 2012 UTC (12 years, 11 months ago) by nmohr
Content type: text/x-python
Branch: MAIN
CVS Tags: AN-12-181-7TeV_patch1, AN-12-181-7TeV
Log Message:
Correct systematics treatment

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