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