VHbb/python/write_systematics.py

#!/usr/bin/env python
from samplesclass import sample
from printcolor import printc
import pickle
import sys
import os
import ROOT 
import shutil
from ROOT import TFile
import ROOT
from array import array
import warnings
warnings.filterwarnings( action='ignore', category=RuntimeWarning, message='creating converter.*' )


#usage: ./write_systematic.py path

path=sys.argv[1]

#load info
infofile = open(path+'/samples.info','r')
info = pickle.load(infofile)
infofile.close()
#os.mkdir(path+'/sys')

for job in info:
    if eval(job.active):
        if job.type != 'DATA':
            print '\t - %s' %(job.name)
            input = TFile.Open(job.getpath(),'read')
            output = TFile.Open(job.path+'/sys/'+job.prefix+job.identifier+'.root','recreate')

            input.cd()
            obj = ROOT.TObject
            for key in ROOT.gDirectory.GetListOfKeys():
                input.cd()
                obj = key.ReadObj()
                #print obj.GetName()
                if obj.GetName() == job.tree:
                    continue
                output.cd()
                #print key.GetName()
                obj.Write(key.GetName())

            tree = input.Get(job.tree)
            nEntries = tree.GetEntries()
            
            job.addpath('/sys')
            job.SYS = ['Nominal','JER_up','JER_down','JES_up','JES_down','beff_up','beff_down','bmis_up','bmis_down']

            output.cd()
            newtree = tree.CloneTree(0)
            
            hJ0 = ROOT.TLorentzVector()
            hJ1 = ROOT.TLorentzVector()
              
            #JER branches
            hJet_pt_JER_up = array('f',[0]*2)
            newtree.Branch('hJet_pt_JER_up',hJet_pt_JER_up,'hJet_pt_JER_up[2]/F')
            hJet_pt_JER_down = array('f',[0]*2)
            newtree.Branch('hJet_pt_JER_down',hJet_pt_JER_down,'hJet_pt_JER_down[2]/F')
            hJet_e_JER_up = array('f',[0]*2)
            newtree.Branch('hJet_e_JER_up',hJet_e_JER_up,'hJet_e_JER_up[2]/F')
            hJet_e_JER_down = array('f',[0]*2)
            newtree.Branch('hJet_e_JER_down',hJet_e_JER_down,'hJet_e_JER_down[2]/F')
            H_JER = array('f',[0]*4)
            newtree.Branch('H_JER',H_JER,'mass_up:mass_down:pt_up:pt_down/F')
            
            #JES branches
            hJet_pt_JES_up = array('f',[0]*2)
            newtree.Branch('hJet_pt_JES_up',hJet_pt_JES_up,'hJet_pt_JES_up[2]/F')
            hJet_pt_JES_down = array('f',[0]*2)
            newtree.Branch('hJet_pt_JES_down',hJet_pt_JES_down,'hJet_pt_JES_down[2]/F')
            hJet_e_JES_up = array('f',[0]*2)
            newtree.Branch('hJet_e_JES_up',hJet_e_JES_up,'hJet_e_JES_up[2]/F')
            hJet_e_JES_down = array('f',[0]*2)
            newtree.Branch('hJet_e_JES_down',hJet_e_JES_down,'hJet_e_JES_down[2]/F')
            H_JES = array('f',[0]*4)
            newtree.Branch('H_JES',H_JES,'mass_up:mass_down:pt_up:pt_down/F')
            
            #Add training Flag
            EventForTraining = array('f',[0])
            newtree.Branch('EventForTraining',EventForTraining,'EventForTraining/F')
            
            #iter=0
            
            TFlag=ROOT.TTreeFormula("EventForTraining","EVENT.event%2",tree)
            
            for entry in range(0,nEntries):
                tree.GetEntry(entry)

                #fill training flag 
                #iter+=1
                #if (iter%2==0):
                #    EventForTraining[0]=1
                #else:
                #    EventForTraining[0]=0
                #iter+=1
                
                EventForTraining[0]=int(not TFlag.EvalInstance())

                #get
                hJet_pt0 = tree.hJet_pt[0]
                hJet_pt1 = tree.hJet_pt[1]
                hJet_eta0 = tree.hJet_eta[0]
                hJet_eta1 = tree.hJet_eta[1]
                hJet_genPt0 = tree.hJet_genPt[0]
                hJet_genPt1 = tree.hJet_genPt[1]
                hJet_e0 = tree.hJet_e[0]
                hJet_e1 = tree.hJet_e[1]
                hJet_phi0 = tree.hJet_phi[0]
                hJet_phi1 = tree.hJet_phi[1]
                hJet_JECUnc0 = tree.hJet_JECUnc[0]
                hJet_JECUnc1 = tree.hJet_JECUnc[1]


                for updown in ['up','down']:
                    #JER
                    if updown == 'up':
                        inner = 0.06
                        outer = 0.1
                    if updown == 'down':
                        inner = -0.06
                        outer = -0.1
                    #Calculate
                    if abs(hJet_eta0)<1.1: res0 = inner
                    else: res0 = outer
                    if abs(hJet_eta1)<1.1: res1 = inner
                    else: res1 = outer
                    rPt0 = hJet_pt0 + (hJet_pt0-hJet_genPt0)*res0
                    rPt1 = hJet_pt1 + (hJet_pt1-hJet_genPt1)*res1
                    rE0 = hJet_e0*rPt0/hJet_pt0
                    rE1 = hJet_e1*rPt1/hJet_pt1
                    hJ0.SetPtEtaPhiE(rPt0,hJet_eta0,hJet_phi0,rE0)
                    hJ1.SetPtEtaPhiE(rPt1,hJet_eta1,hJet_phi1,rE1)
                    #Set
                    if updown == 'up':
                        hJet_pt_JER_up[0]=rPt0
                        hJet_pt_JER_up[1]=rPt1
                        hJet_e_JER_up[0]=rE0
                        hJet_e_JER_up[1]=rE1
                        H_JER[0]=(hJ0+hJ1).M()
                        H_JER[2]=(hJ0+hJ1).Pt()
                    if updown == 'down':
                        hJet_pt_JER_down[0]=rPt0
                        hJet_pt_JER_down[1]=rPt1
                        hJet_e_JER_down[0]=rE0
                        hJet_e_JER_down[1]=rE1
                        H_JER[1]=(hJ0+hJ1).M()
                        H_JER[3]=(hJ0+hJ1).Pt()
                    
                    #JES
                    if updown == 'up':
                        variation=1
                    if updown == 'down':
                        variation=-1
                    #calculate
                    rPt0 = hJet_pt0*(1+variation*hJet_JECUnc0)
                    rPt1 = hJet_pt1*(1+variation*hJet_JECUnc1)
                    rE0 = hJet_e0*(1+variation*hJet_JECUnc0)
                    rE1 = hJet_e1*(1+variation*hJet_JECUnc1)
                    hJ0.SetPtEtaPhiE(rPt0,hJet_eta0,hJet_phi0,rE0)
                    hJ1.SetPtEtaPhiE(rPt1,hJet_eta1,hJet_phi1,rE1)
                    #Fill
                    if updown == 'up':
                        hJet_pt_JES_up[0]=rPt0
                        hJet_pt_JES_up[1]=rPt1
                        hJet_e_JES_up[0]=rE0
                        hJet_e_JES_up[1]=rE1
                        H_JES[0]=(hJ0+hJ1).M()
                        H_JES[2]=(hJ0+hJ1).Pt()
                    if updown == 'down':
                        hJet_pt_JES_down[0]=rPt0
                        hJet_pt_JES_down[1]=rPt1
                        hJet_e_JES_down[0]=rE0
                        hJet_e_JES_down[1]=rE1
                        H_JES[1]=(hJ0+hJ1).M()
                        H_JES[3]=(hJ0+hJ1).Pt()
                
                newtree.Fill()
                       
            newtree.AutoSave()
            output.Close()
            
        else: #(is data)
        
            print '\t - %s' %(job.name)
            input = TFile.Open(job.getpath(),'read')
            output = TFile.Open(job.path+'/sys/'+job.prefix+job.identifier+'.root','recreate')
            #input.cd()
            obj = ROOT.TObject
            for key in ROOT.gDirectory.GetListOfKeys():
                input.cd()
                obj = key.ReadObj()
                if obj.GetName() == job.tree:
                    continue
                output.cd()
                obj.Write(key.GetName())
            #output.cd()
            tree = input.Get(job.tree)
            nEntries = tree.GetEntries()
            job.addpath('/sys')
            output.cd()
            newtree = tree.CloneTree(0)
            #Add training Flag
            EventForTraining = array('f',[0])
            newtree.Branch('EventForTraining',EventForTraining,'EventForTraining/F')
            for entry in range(0,nEntries):
                tree.GetEntry(entry)
                EventForTraining[0]=0
                newtree.Fill()
            newtree.AutoSave()
            output.Close()


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