ViewVC Help
View File | Revision Log | Show Annotations | Root Listing
root/cvsroot/UserCode/VHbb/python/write_systematics.py
Revision: 1.3
Committed: Thu May 10 16:16:05 2012 UTC (13 years ago) by peller
Content type: text/x-python
Branch: MAIN
Changes since 1.2: +31 -24 lines
Log Message:
update

File Contents

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