VHbb/python/write_systematics.py

#!/usr/bin/env python
from samplesinfo 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


#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 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')
        output.cd()
        newtree = tree.CloneTree(0)


        '''
        input = TFile.Open(job.getpath(),'read')
        Count = input.Get("Count")
        CountWithPU = input.Get("CountWithPU")
        CountWithPU2011B = input.Get("CountWithPU2011B")
        tree = input.Get(job.tree)
        nEntries = tree.GetEntries()
        
        job.addpath('/sys')
        output = ROOT.TFile(job.getpath(), 'RECREATE')
        newtree = tree.CloneTree(0)
        '''
        job.SYS = ['Nominal','JER_up','JER_down','JES_up','JES_down','beff_up','beff_down','bmis_up','bmis_down']
        
        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
        
        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

            #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()
        #newtree.Write()            
        #Count.Write()
        #CountWithPU.Write()
        #CountWithPU2011B.Write()
        output.Close()

    else: #(is data)
    
        shutil.copy(job.getpath(),path+'/sys')
        job.addpath('/sys')

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