kiesel/plotTree/splitCandidates.py

#! /usr/bin/env python2
# -*- coding: utf-8 -*-
from treeFunctions import *
import Styles
#Styles.tdrStyle()
import numpy
import ROOT
import argparse
import ConfigParser

def deltaPhi( phi1, phi2):
        import math
        result = phi1 - phi2
        while result > math.pi:
                result -= 2*math.pi
        while result <= -math.pi:
                result += 2*math.pi
        return result

def deltaR( object1, object2 ):
        import math
        return math.sqrt( (object1.eta-object2.eta)**2 + (deltaPhi(object1.phi,object2.phi))**2 )

def gammaSelectionClone( tree, newTreeName ):
        """ Clones a tree and get access to photon vector"""
        newTree = tree.CloneTree(0)
        newTree.SetName( newTreeName )
        photons = ROOT.std.vector("tree::Photon")()
        newTree.SetBranchAddress("photon", photons )

        return newTree, photons

def combineVectors( v1, v2 ):
        for element in v1:
                v2.push_back(element)
        return v2

def recElectronMatching( genElectrons, photonElectrons ):
        for recE in photonElectrons:
                minDeltaR = 50
                minPt = 8000
                minIndex = -1
                for i in range( genElectrons.size()):
                        dr = deltaR( recE, genElectrons[i] )
                        pt = ( recE.pt - genElectrons[i].pt ) / genElectrons[i].pt
                        if dr < minDeltaR:
                                minIndex = i
                                minDeltaR = dr
                                minPt = pt

                if minDeltaR < 0.02 and abs(minPt) < 0.5:
                        photonElectrons[minIndex].phi = 5
                        genElectrons[minIndex].phi = 5

        return genElectrons, photonElectrons

def genElectronMatching( photons, photonElectrons, genElectrons, hist ):
        for genE in genElectrons:
                minDeltaR = 50
                minPt = 8000
                minIndex = -1
                typeName = ""
                for i in range( photons.size()):
                        dr = deltaR( genE, photons[i] )
                        pt = ( photons[i].pt - genE.pt ) / genE.pt
                        if dr < minDeltaR:
                                minIndex = i
                                minDeltaR = dr
                                minPt = pt
                                typeName = "g"
                for i in range( photonElectrons.size()):
                        dr = deltaR( genE, photonElectrons[i] )
                        pt = ( photonElectrons[i].pt - genE.pt ) / genE.pt
                        if dr < minDeltaR:
                                minIndex = i
                                minDeltaR = dr
                                minPt = pt
                                typeName = "e"

                #if minDeltaR < 0.02 and abs(minPt) < 0.5:
                if minDeltaR<50:
                        hist.Fill( minPt, minDeltaR )
                        if typeName == "e":
                                photonElectrons[minIndex].pixelseed = -2
                        if typeName == "g":
                                photons[minIndex].pixelseed = -2
        return photons, photonElectrons, hist

def weightTree( photonTree, photonJetTree ):
        # a tree can't be modified, therefore a new tree is saved within the same file

        # this is the binning with witch the fo reweighting will be done and has to be studied
        nBins = 10
        xMin = 80
        xMax = 800
        h_photonPt = createHistoFromTree( photonTree, "photon[0].pt", "", nBins, xMin, xMax )
        h_jetPt = createHistoFromTree( photonJetTree, "photon[0].pt", "", nBins, xMin, xMax )

        # h_jetPt is now histogram with w^{-1} = h_jetPt / h_photonPt
        h_jetPt.Divide( h_photonPt )
        h_jetPt.SetTitle(";p_{T,#gamma};w^{-1}")
        #h_jetPt.Draw()
        #raw_input()

        weightTree = ROOT.TTree("QCDWeightTree", "my TFriend for weights")
        import numpy

        # a python float corresponds to a root double
        weight = numpy.zeros(1, dtype=float)
        weightError = numpy.zeros(1, dtype=float)
        weightTree.Branch( "weight", weight, "weight/D" )
        weightTree.Branch( "weightError", weightError, "weightError/D" )

        for event in photonJetTree:
                try:
                        pt = event.photon.at(0).pt
                except:
                        pt = 0
                bin = h_jetPt.FindBin( pt )
                weight[0] = h_jetPt.GetBinContent( bin )
                weightError[0] = h_jetPt.GetBinError( bin )
                weightTree.Fill()

        return weightTree

def splitCandidates( inputFileName, nExpected, isQCD, isEWK):
        outputFileName = "slim"+inputFileName

        eventHisto = readHisto( inputFileName )
        nGenerated = eventHisto.GetBinContent(1)

        tree = readTree( inputFileName )

        photonTree, photons = gammaSelectionClone( tree, "photonTree" )
        photonJetTree, photonJets = gammaSelectionClone( tree, "photonJetTree" )
        photonElectronTree, photonElectrons = gammaSelectionClone( tree, "photonElectronTree" )
        #genElectronTree = gammaSelectionClone( tree, "genElectronTree" )[0]

        weight = numpy.zeros(1, dtype=float)
        photonTree.SetBranchAddress("weight", weight)
        photonElectronTree.SetBranchAddress("weight", weight)
        photonJetTree.SetBranchAddress("weight", weight)
        #genElectronTree.SetBranchAddress("weight", weight)

        #genElectrons = ROOT.std.vector("tree::Particle")()
        #photonElectronTree.SetBranchAddress("genElectron", genElectrons )


        hist2 = ROOT.TH2F( inputFileName, "", 100, -.30, .30, 100, 0, 0.06 )
        hist2.SetTitle(";#frac{p_{T,rec}-p_{T,gen}}{p_{T,gen}};min #DeltaR")

        for event in tree:
                if not event.GetReadEntry()%100000:
                        print 100*event.GetReadEntry()/event.GetEntries(), "%"
                weight[0] = event.weight * nExpected / nGenerated
                photons.clear()
                photonElectrons.clear()
                photonJets.clear()

                for gamma in event.photon:

                        # cuts for every object to rject spikes
                        if gamma.r9 < 1 \
                        and gamma.sigmaIetaIeta > 0.001 \
                        and abs(gamma.eta) < 1.479 \
                        and gamma.pt > 20:

                                # look for gamma and electrons
                                if gamma.hadTowOverEm < 0.05 \
                                and gamma.sigmaIetaIeta < 0.012 \
                                and gamma.chargedIso < 2.6 \
                                and gamma.neutralIso < 3.5 + 0.04*gamma.pt \
                                and gamma.photonIso < 1.3 + 0.005*gamma.pt:
                                        if gamma.pixelseed:
                                                photonElectrons.push_back( gamma )
                                        else:
                                                photons.push_back(gamma)

                                # QCD fake object definition
                                if gamma.ptJet > 75 \
                                and gamma.hadTowOverEm < 0.05 \
                                and gamma.sigmaIetaIeta < 0.014 \
                                and gamma.chargedIso < 15 \
                                and gamma.neutralIso < 3.5 + 0.04*gamma.pt \
                                and gamma.photonIso < 1.3 + 0.005*gamma.pt \
                                and not gamma.pixelseed \
                                and ( gamma.sigmaIetaIeta>=0.012 or gamma.chargedIso>=2.6):
                                        photonJets.push_back( gamma )

                photons, photonElectrons, hist2 = genElectronMatching( photons, photonElectrons, event.genElectron, hist2 )

                #genElectrons = recElectronMatching( event.genElectron, photonElectrons )

                if photons.size() > 0:
                        photonTree.Fill()
                else:
                        if photonElectrons.size() > 0:
                                photonElectronTree.Fill()
                        if photonJets.size() > 0:
                                photonJetTree.Fill()
                #if (photons.size()>0 or photonElectrons.size()>0) and event.genElectron.size() > 0:
                #       genElectronTree.Fill()

        can = ROOT.TCanvas()
        can.cd()
        hist2.Draw("colz")
        can.SaveAs("pt_r_%s_new.pdf"%inName[0:5])

        if isQCD:
                qcdWeights = weightTree( photonTree, photonJetTree )


        f = ROOT.TFile( outputFileName, "recreate" )
        f.cd()
        photonTree.Write()
        photonJetTree.Write()
        photonElectronTree.Write()
        #genElectronTree.Write()
        if isQCD:
                qcdWeights.Write()
        f.Write()
        f.Close()


if __name__ == "__main__":
        # include knowledge about objects saved in the tree
        ROOT.gSystem.Load("libTreeObjects.so")

        arguments = argparse.ArgumentParser( description="Slim tree" )
        arguments.add_argument("--input", default=["WJets_V01.03_tree.root"], nargs="+" )
        opts = arguments.parse_args()

        ROOT.gROOT.SetBatch()


        datasetConf = ConfigParser.SafeConfigParser()
        datasetConf.read( "dataset.cfg" )

        integratedLumi = 19300 #pb

        for inName in opts.input:
                shortName = None
                for configName in datasetConf.sections():
                        if inName.count( configName ):
                                shortName = configName
                                crosssection = datasetConf.getfloat( configName, "crosssection" )

                qcd = False
                ewk = False
                if shortName in ["QCD_250_500", "QCD_500_1000", "QCD_1000_inf","GJets"]:
                        qcd = True
                if shortName in ["TTJets", "WJets"]:
                        ewk = True

                # N = L * sigma
                nExpected = integratedLumi * crosssection

                splitCandidates( inName, nExpected, isQCD=qcd, isEWK=ewk )

Revision:	1.2
Committed:	Fri May 3 08:56:35 2013 UTC (12 years ago) by kiesel
Content type:	text/x-python
Branch:	MAIN
Changes since 1.1:	+101 -16 lines
Log Message:	added plotting helpers (ugly now, but only for saving)
#	User	Rev	Content
1	kiesel	1.1	#! /usr/bin/env python2
2			# -- coding: utf-8 --
3			from treeFunctions import *
4	kiesel	1.2	import Styles
5			#Styles.tdrStyle()
6	kiesel	1.1	import numpy
7			import ROOT
8			import argparse
9			import ConfigParser
10
11			def deltaPhi( phi1, phi2):
12			import math
13			result = phi1 - phi2
14			while result > math.pi:
15			result -= 2*math.pi
16			while result <= -math.pi:
17			result += 2*math.pi
18			return result
19
20			def deltaR( object1, object2 ):
21			import math
22			return math.sqrt( (object1.eta-object2.eta)2 + (deltaPhi(object1.phi,object2.phi))2 )
23
24			def gammaSelectionClone( tree, newTreeName ):
25			""" Clones a tree and get access to photon vector"""
26			newTree = tree.CloneTree(0)
27			newTree.SetName( newTreeName )
28			photons = ROOT.std.vector("tree::Photon")()
29			newTree.SetBranchAddress("photon", photons )
30
31			return newTree, photons
32
33			def combineVectors( v1, v2 ):
34			for element in v1:
35			v2.push_back(element)
36			return v2
37
38	kiesel	1.2	def recElectronMatching( genElectrons, photonElectrons ):
39			for recE in photonElectrons:
40			minDeltaR = 50
41			minPt = 8000
42			minIndex = -1
43			for i in range( genElectrons.size()):
44			dr = deltaR( recE, genElectrons[i] )
45			pt = ( recE.pt - genElectrons[i].pt ) / genElectrons[i].pt
46			if dr < minDeltaR:
47			minIndex = i
48			minDeltaR = dr
49			minPt = pt
50
51			if minDeltaR < 0.02 and abs(minPt) < 0.5:
52			photonElectrons[minIndex].phi = 5
53			genElectrons[minIndex].phi = 5
54
55			return genElectrons, photonElectrons
56
57			def genElectronMatching( photons, photonElectrons, genElectrons, hist ):
58	kiesel	1.1	for genE in genElectrons:
59			minDeltaR = 50
60			minPt = 8000
61			minIndex = -1
62			typeName = ""
63			for i in range( photons.size()):
64			dr = deltaR( genE, photons[i] )
65			pt = ( photons[i].pt - genE.pt ) / genE.pt
66			if dr < minDeltaR:
67			minIndex = i
68			minDeltaR = dr
69			minPt = pt
70			typeName = "g"
71			for i in range( photonElectrons.size()):
72			dr = deltaR( genE, photonElectrons[i] )
73			pt = ( photonElectrons[i].pt - genE.pt ) / genE.pt
74			if dr < minDeltaR:
75			minIndex = i
76			minDeltaR = dr
77			minPt = pt
78			typeName = "e"
79
80	kiesel	1.2	#if minDeltaR < 0.02 and abs(minPt) < 0.5:
81			if minDeltaR<50:
82	kiesel	1.1	hist.Fill( minPt, minDeltaR )
83			if typeName == "e":
84			photonElectrons[minIndex].pixelseed = -2
85			if typeName == "g":
86			photons[minIndex].pixelseed = -2
87			return photons, photonElectrons, hist
88
89	kiesel	1.2	def weightTree( photonTree, photonJetTree ):
90			# a tree can't be modified, therefore a new tree is saved within the same file
91
92			# this is the binning with witch the fo reweighting will be done and has to be studied
93			nBins = 10
94			xMin = 80
95			xMax = 800
96			h_photonPt = createHistoFromTree( photonTree, "photon[0].pt", "", nBins, xMin, xMax )
97			h_jetPt = createHistoFromTree( photonJetTree, "photon[0].pt", "", nBins, xMin, xMax )
98
99			# h_jetPt is now histogram with w^{-1} = h_jetPt / h_photonPt
100			h_jetPt.Divide( h_photonPt )
101			h_jetPt.SetTitle(";p_{T,#gamma};w^{-1}")
102			#h_jetPt.Draw()
103			#raw_input()
104
105			weightTree = ROOT.TTree("QCDWeightTree", "my TFriend for weights")
106			import numpy
107
108			# a python float corresponds to a root double
109			weight = numpy.zeros(1, dtype=float)
110			weightError = numpy.zeros(1, dtype=float)
111			weightTree.Branch( "weight", weight, "weight/D" )
112			weightTree.Branch( "weightError", weightError, "weightError/D" )
113
114			for event in photonJetTree:
115			try:
116			pt = event.photon.at(0).pt
117			except:
118			pt = 0
119			bin = h_jetPt.FindBin( pt )
120			weight[0] = h_jetPt.GetBinContent( bin )
121			weightError[0] = h_jetPt.GetBinError( bin )
122			weightTree.Fill()
123
124			return weightTree
125
126			def splitCandidates( inputFileName, nExpected, isQCD, isEWK):
127	kiesel	1.1	outputFileName = "slim"+inputFileName
128
129			eventHisto = readHisto( inputFileName )
130			nGenerated = eventHisto.GetBinContent(1)
131
132			tree = readTree( inputFileName )
133
134			photonTree, photons = gammaSelectionClone( tree, "photonTree" )
135			photonJetTree, photonJets = gammaSelectionClone( tree, "photonJetTree" )
136			photonElectronTree, photonElectrons = gammaSelectionClone( tree, "photonElectronTree" )
137	kiesel	1.2	#genElectronTree = gammaSelectionClone( tree, "genElectronTree" )[0]
138	kiesel	1.1
139	kiesel	1.2	weight = numpy.zeros(1, dtype=float)
140			photonTree.SetBranchAddress("weight", weight)
141			photonElectronTree.SetBranchAddress("weight", weight)
142			photonJetTree.SetBranchAddress("weight", weight)
143			#genElectronTree.SetBranchAddress("weight", weight)
144	kiesel	1.1
145	kiesel	1.2	#genElectrons = ROOT.std.vector("tree::Particle")()
146			#photonElectronTree.SetBranchAddress("genElectron", genElectrons )
147	kiesel	1.1
148	kiesel	1.2
149			hist2 = ROOT.TH2F( inputFileName, "", 100, -.30, .30, 100, 0, 0.06 )
150	kiesel	1.1	hist2.SetTitle(";#frac{p_{T,rec}-p_{T,gen}}{p_{T,gen}};min #DeltaR")
151
152			for event in tree:
153	kiesel	1.2	if not event.GetReadEntry()%100000:
154			print 100*event.GetReadEntry()/event.GetEntries(), "%"
155			weight[0] = event.weight * nExpected / nGenerated
156	kiesel	1.1	photons.clear()
157			photonElectrons.clear()
158			photonJets.clear()
159
160			for gamma in event.photon:
161
162			# cuts for every object to rject spikes
163			if gamma.r9 < 1 \
164			and gamma.sigmaIetaIeta > 0.001 \
165			and abs(gamma.eta) < 1.479 \
166			and gamma.pt > 20:
167
168			# look for gamma and electrons
169			if gamma.hadTowOverEm < 0.05 \
170			and gamma.sigmaIetaIeta < 0.012 \
171			and gamma.chargedIso < 2.6 \
172			and gamma.neutralIso < 3.5 + 0.04*gamma.pt \
173			and gamma.photonIso < 1.3 + 0.005*gamma.pt:
174			if gamma.pixelseed:
175			photonElectrons.push_back( gamma )
176			else:
177			photons.push_back(gamma)
178
179			# QCD fake object definition
180	kiesel	1.2	if gamma.ptJet > 75 \
181	kiesel	1.1	and gamma.hadTowOverEm < 0.05 \
182			and gamma.sigmaIetaIeta < 0.014 \
183			and gamma.chargedIso < 15 \
184			and gamma.neutralIso < 3.5 + 0.04*gamma.pt \
185			and gamma.photonIso < 1.3 + 0.005*gamma.pt \
186			and not gamma.pixelseed \
187			and ( gamma.sigmaIetaIeta>=0.012 or gamma.chargedIso>=2.6):
188			photonJets.push_back( gamma )
189
190	kiesel	1.2	photons, photonElectrons, hist2 = genElectronMatching( photons, photonElectrons, event.genElectron, hist2 )
191	kiesel	1.1
192	kiesel	1.2	#genElectrons = recElectronMatching( event.genElectron, photonElectrons )
193	kiesel	1.1
194			if photons.size() > 0:
195			photonTree.Fill()
196			else:
197			if photonElectrons.size() > 0:
198			photonElectronTree.Fill()
199			if photonJets.size() > 0:
200			photonJetTree.Fill()
201	kiesel	1.2	#if (photons.size()>0 or photonElectrons.size()>0) and event.genElectron.size() > 0:
202			# genElectronTree.Fill()
203
204			can = ROOT.TCanvas()
205			can.cd()
206			hist2.Draw("colz")
207			can.SaveAs("pt_r_%s_new.pdf"%inName[0:5])
208	kiesel	1.1
209	kiesel	1.2	if isQCD:
210			qcdWeights = weightTree( photonTree, photonJetTree )
211	kiesel	1.1
212
213			f = ROOT.TFile( outputFileName, "recreate" )
214			f.cd()
215			photonTree.Write()
216			photonJetTree.Write()
217			photonElectronTree.Write()
218	kiesel	1.2	#genElectronTree.Write()
219			if isQCD:
220			qcdWeights.Write()
221	kiesel	1.1	f.Write()
222			f.Close()
223
224
225			if __name__ == "__main__":
226			# include knowledge about objects saved in the tree
227			ROOT.gSystem.Load("libTreeObjects.so")
228
229			arguments = argparse.ArgumentParser( description="Slim tree" )
230			arguments.add_argument("--input", default=["WJets_V01.03_tree.root"], nargs="+" )
231			opts = arguments.parse_args()
232
233	kiesel	1.2	ROOT.gROOT.SetBatch()
234	kiesel	1.1
235
236			datasetConf = ConfigParser.SafeConfigParser()
237			datasetConf.read( "dataset.cfg" )
238
239			integratedLumi = 19300 #pb
240
241			for inName in opts.input:
242	kiesel	1.2	shortName = None
243	kiesel	1.1	for configName in datasetConf.sections():
244			if inName.count( configName ):
245	kiesel	1.2	shortName = configName
246	kiesel	1.1	crosssection = datasetConf.getfloat( configName, "crosssection" )
247
248	kiesel	1.2	qcd = False
249			ewk = False
250			if shortName in ["QCD_250_500", "QCD_500_1000", "QCD_1000_inf","GJets"]:
251			qcd = True
252			if shortName in ["TTJets", "WJets"]:
253			ewk = True
254
255	kiesel	1.1	# N = L * sigma
256			nExpected = integratedLumi * crosssection
257
258	kiesel	1.2	splitCandidates( inName, nExpected, isQCD=qcd, isEWK=ewk )
259	kiesel	1.1