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

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