kiesel/plotTree/splitCandidates.py

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

def deltaPhi( phi1, phi2):
        """Computes delta phi.
        Due to the zylindrical form of CMS, it has be be corrected by a diffenence
        of 2pi.
        """
        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 ):
        """Compute delta R of two objects.
        Both objects must have eta and phi as member variables.
        """
        import math
        return math.sqrt( (object1.eta-object2.eta)**2 + (deltaPhi(object1.phi,object2.phi))**2 )

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

def histoDefinition():
        """All histograms are defined here and put into an directory."""
        from random import randint
        histos = {}
        histos["dEtadPhiLarge"] = ROOT.TH2F( "%x"%randint(0,100000), ";#Delta#eta;#Delta#phi", 100, -4, 4, 100, -3.3, 3.3 )
        histos["dEtadPhiSmall"] = ROOT.TH2F( "%x"%randint(0,1000000), ";|#Delta#eta|;|#Delta#phi|", 100, 0, .012, 100, 0, .1 )
        histos["dEtadPhiSmallPhoton"] = ROOT.TH2F( "%x"%randint(0,1000000), ";|#Delta#eta|;|#Delta#phi|", 100, 0, .012, 100, 0, .1 )
        histos["dEtadPhiSmallElectron"] = ROOT.TH2F( "%x"%randint(0,1000000), ";|#Delta#eta|;|#Delta#phi|", 100, 0, .012, 100, 0, .1 )
        histos["nGenEnRec"] = ROOT.TH2F( "%x"%randint(0,10000000), ";generated e;reco EM-Objects", 5, -0.5, 4.5, 5, -0.5, 4.5 )
        histos["dPtdR"] = ROOT.TH2F( "%x"%randint(0,10000000), ";|(p_{T}-p_{T gen} )/ p_{T gen}|;#DeltaR", 100, 0, 1, 100, 0, .5 )
        return histos

def draw_histogram_dict( histograms, suffix="new" ):
        """Draw all histograms in this directory and save it as pdf."""
        can = ROOT.TCanvas()
        can.cd()
        dataset = ROOT.TPaveText(.4,.9,.6,.98, "ndc")
        dataset.SetFillColor(0)
        dataset.SetBorderSize(0)
        dataset.AddText( suffix )
        for name, histo in histograms.iteritems():
                name = name.replace(".","")
                histo.Draw("colz")
                dataset.Draw()
                can.SaveAs("plots/%s_%s.pdf"%(name,suffix))

def generalMatching( objects1, objects2, hist):
        """Matches for each object in objects1 an object in objects2.
        'hists' is a dict containing ROOT.TH* objects which will be filled and returned.
        The gen information will be storend in first list and returned.
        """
        hist["nGenEnRec"].Fill( objects1.size(), objects2.size() )
        for o1 in objects1:
                match = False
                minDeltaPt = sys.maxint
                minIndex = -1
                for i, o2 in enumerate(objects2):
                        DeltaEta = o1.eta - o2.eta
                        DeltaPhi = deltaPhi( o1.phi, o2.phi)
                        absDeltaEta = abs( DeltaEta )
                        absDeltaPhi = abs( DeltaPhi )
                        absDeltaPt = 2*abs( o1.pt - o2.pt ) / ( o1.pt + o2.pt )
                        hist["dEtadPhiLarge"].Fill( DeltaEta, DeltaPhi )
                        hist["dEtadPhiSmall"].Fill( absDeltaEta, absDeltaPhi )
                        if hasattr( o1, "pixelseed" ):
                                if o1.pixelseed:
                                        hist["dEtadPhiSmallElectron"].Fill( absDeltaEta, absDeltaPhi )
                                else:
                                        hist["dEtadPhiSmallPhoton"].Fill( absDeltaEta, absDeltaPhi )

                        if absDeltaEta < .008 and absDeltaPhi < .08 and absDeltaPt < .1:
                                match = True
                                if absDeltaPt < minDeltaPt:
                                        minDeltaPt = absDeltaPt
                                        minIndex = i

                if match:
                        if hasattr( o1, "genInformation" ):
                                o1.genInformation = 1
                        else: # use phi variable for gen information. TODO: fix that
                                # only fill gen matching for photons matching a gen electrons.
                                # Electrons are not needed in fake rate
                                if not objects2[minIndex].pixelseed:
                                        o1.phi = 5
        return objects1, hist

def splitCandidates( inputFileName, shortName, nExpected, processNEvents=-1, genMatching=False ):
        print "Processing file {}".format(inputFileName)

        #genMatching = not shortName in ["WJets", "DY", "TTJets"]
        if genMatching:
                print "Match generated objects"

        tree = readTree( inputFileName )
        eventHisto = readHisto( inputFileName )
        if processNEvents == -1:
                processNEvents = eventHisto.GetBinContent(1)

        weight = numpy.zeros(1, dtype=float)
        photonTree, photons, weight = gammaSelectionClone( tree, "photonTree", weight )
        photonJetTree, photonJets, weight = gammaSelectionClone( tree, "photonJetTree", weight )
        photonElectronTree, photonElectrons, weight = gammaSelectionClone( tree, "photonElectronTree", weight )
        if genMatching:
                genElectronTree, genElectrons, weight = gammaSelectionClone( tree, "genElectronTree", weight, "tree::Particle","genElectron" )
                histograms = histoDefinition()

        emObjects = ROOT.std.vector("tree::Photon")()

        for event in tree:
                if not event.GetReadEntry()%100000:
                        print '{0}%\r'.format(100*event.GetReadEntry()/event.GetEntries())

                if event.GetReadEntry() > processNEvents:
                        break

                weight[0] = event.weight * nExpected / processNEvents
                emObjects.clear()
                photons.clear()
                photonElectrons.clear()
                photonJets.clear()
                if genMatching:
                        genElectrons.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.4442 \
                        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:
                                        emObjects.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 )

                if genMatching:
                        emObjects, histograms = generalMatching( emObjects, event.genElectron, histograms )
                        for genE in event.genElectron:
                                if abs(genE.eta) < 1.4442:
                                        genElectrons.push_back( genE )
                        genElectrons, histograms = generalMatching( genElectrons, emObjects, histograms )

                for emObject in emObjects:
                        if emObject.pixelseed:
                                photonElectrons.push_back( emObject )
                        else:
                                photons.push_back( emObject )

                if photons.size() > 0:
                        photonTree.Fill()
                else:
                        if photonElectrons.size() > 0:
                                photonElectronTree.Fill()
                        if photonJets.size() > 0:
                                photonJetTree.Fill()
                if genMatching and genElectrons.size() > 0:
                        genElectronTree.Fill()


        outputFileName = "slim"+inputFileName
        f = ROOT.TFile( outputFileName, "recreate" )
        f.cd()
        photonTree.Write()
        photonJetTree.Write()
        photonElectronTree.Write()
        if genMatching:
                genElectronTree.Write()
                draw_histogram_dict( histograms, shortName )
        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.12_tree.root"], nargs="+" )
        arguments.add_argument("--test", action="store_true" )
        arguments.add_argument("--genMatching", action="store_true" )
        opts = arguments.parse_args()

        processNEvents = 10000 if opts.test else -1


        # disable open canvas
        ROOT.gROOT.SetBatch()

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

        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" )
                if not shortName:
                        print "No configuration for input file {} defined in '{}'".format(
                                        inName, datasetConfigName )
                        continue

                # N = L * sigma
                nExpected = integratedLumi * crosssection

                splitCandidates( inName, shortName, nExpected, processNEvents, opts.genMatching )

Revision:	1.3
Committed:	Wed May 15 14:20:17 2013 UTC (11 years, 11 months ago) by kiesel
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	HEAD
Changes since 1.2:	+139 -154 lines
Log Message:	nicer plot routines
#	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	kiesel	1.3	Styles.tdrStyle2D()
6	kiesel	1.1	import numpy
7			import ROOT
8			import argparse
9			import ConfigParser
10	kiesel	1.3	import sys
11	kiesel	1.1
12			def deltaPhi( phi1, phi2):
13	kiesel	1.3	"""Computes delta phi.
14			Due to the zylindrical form of CMS, it has be be corrected by a diffenence
15			of 2pi.
16			"""
17	kiesel	1.1	import math
18			result = phi1 - phi2
19			while result > math.pi:
20			result -= 2*math.pi
21			while result <= -math.pi:
22			result += 2*math.pi
23			return result
24
25			def deltaR( object1, object2 ):
26	kiesel	1.3	"""Compute delta R of two objects.
27			Both objects must have eta and phi as member variables.
28			"""
29	kiesel	1.1	import math
30			return math.sqrt( (object1.eta-object2.eta)2 + (deltaPhi(object1.phi,object2.phi))2 )
31
32	kiesel	1.3	def gammaSelectionClone( tree, newTreeName, weight, type_="tree::Photon", name="photon" ):
33			"""Clones a tree and get access to photon vector"""
34	kiesel	1.1	newTree = tree.CloneTree(0)
35			newTree.SetName( newTreeName )
36	kiesel	1.3	photons = ROOT.std.vector(type_)()
37			newTree.SetBranchAddress( name, photons )
38			newTree.SetBranchAddress("weight", weight )
39			return newTree, photons, weight
40
41			def histoDefinition():
42			"""All histograms are defined here and put into an directory."""
43			from random import randint
44			histos = {}
45			histos["dEtadPhiLarge"] = ROOT.TH2F( "%x"%randint(0,100000), ";#Delta#eta;#Delta#phi", 100, -4, 4, 100, -3.3, 3.3 )
46			histos["dEtadPhiSmall"] = ROOT.TH2F( "%x"%randint(0,1000000), ";\|#Delta#eta\|;\|#Delta#phi\|", 100, 0, .012, 100, 0, .1 )
47			histos["dEtadPhiSmallPhoton"] = ROOT.TH2F( "%x"%randint(0,1000000), ";\|#Delta#eta\|;\|#Delta#phi\|", 100, 0, .012, 100, 0, .1 )
48			histos["dEtadPhiSmallElectron"] = ROOT.TH2F( "%x"%randint(0,1000000), ";\|#Delta#eta\|;\|#Delta#phi\|", 100, 0, .012, 100, 0, .1 )
49			histos["nGenEnRec"] = ROOT.TH2F( "%x"%randint(0,10000000), ";generated e;reco EM-Objects", 5, -0.5, 4.5, 5, -0.5, 4.5 )
50			histos["dPtdR"] = ROOT.TH2F( "%x"%randint(0,10000000), ";\|(p_{T}-p_{T gen} )/ p_{T gen}\|;#DeltaR", 100, 0, 1, 100, 0, .5 )
51			return histos
52	kiesel	1.1
53	kiesel	1.3	def draw_histogram_dict( histograms, suffix="new" ):
54			"""Draw all histograms in this directory and save it as pdf."""
55			can = ROOT.TCanvas()
56			can.cd()
57			dataset = ROOT.TPaveText(.4,.9,.6,.98, "ndc")
58			dataset.SetFillColor(0)
59			dataset.SetBorderSize(0)
60			dataset.AddText( suffix )
61			for name, histo in histograms.iteritems():
62			name = name.replace(".","")
63			histo.Draw("colz")
64			dataset.Draw()
65			can.SaveAs("plots/%s_%s.pdf"%(name,suffix))
66
67			def generalMatching( objects1, objects2, hist):
68			"""Matches for each object in objects1 an object in objects2.
69			'hists' is a dict containing ROOT.TH* objects which will be filled and returned.
70			The gen information will be storend in first list and returned.
71			"""
72			hist["nGenEnRec"].Fill( objects1.size(), objects2.size() )
73			for o1 in objects1:
74			match = False
75			minDeltaPt = sys.maxint
76	kiesel	1.1	minIndex = -1
77	kiesel	1.3	for i, o2 in enumerate(objects2):
78			DeltaEta = o1.eta - o2.eta
79			DeltaPhi = deltaPhi( o1.phi, o2.phi)
80			absDeltaEta = abs( DeltaEta )
81			absDeltaPhi = abs( DeltaPhi )
82			absDeltaPt = 2*abs( o1.pt - o2.pt ) / ( o1.pt + o2.pt )
83			hist["dEtadPhiLarge"].Fill( DeltaEta, DeltaPhi )
84			hist["dEtadPhiSmall"].Fill( absDeltaEta, absDeltaPhi )
85			if hasattr( o1, "pixelseed" ):
86			if o1.pixelseed:
87			hist["dEtadPhiSmallElectron"].Fill( absDeltaEta, absDeltaPhi )
88			else:
89			hist["dEtadPhiSmallPhoton"].Fill( absDeltaEta, absDeltaPhi )
90
91			if absDeltaEta < .008 and absDeltaPhi < .08 and absDeltaPt < .1:
92			match = True
93			if absDeltaPt < minDeltaPt:
94			minDeltaPt = absDeltaPt
95			minIndex = i
96
97			if match:
98			if hasattr( o1, "genInformation" ):
99			o1.genInformation = 1
100			else: # use phi variable for gen information. TODO: fix that
101			# only fill gen matching for photons matching a gen electrons.
102			# Electrons are not needed in fake rate
103			if not objects2[minIndex].pixelseed:
104			o1.phi = 5
105			return objects1, hist
106
107			def splitCandidates( inputFileName, shortName, nExpected, processNEvents=-1, genMatching=False ):
108			print "Processing file {}".format(inputFileName)
109
110			#genMatching = not shortName in ["WJets", "DY", "TTJets"]
111			if genMatching:
112			print "Match generated objects"
113	kiesel	1.1
114	kiesel	1.3	tree = readTree( inputFileName )
115	kiesel	1.1	eventHisto = readHisto( inputFileName )
116	kiesel	1.3	if processNEvents == -1:
117			processNEvents = eventHisto.GetBinContent(1)
118	kiesel	1.1
119	kiesel	1.2	weight = numpy.zeros(1, dtype=float)
120	kiesel	1.3	photonTree, photons, weight = gammaSelectionClone( tree, "photonTree", weight )
121			photonJetTree, photonJets, weight = gammaSelectionClone( tree, "photonJetTree", weight )
122			photonElectronTree, photonElectrons, weight = gammaSelectionClone( tree, "photonElectronTree", weight )
123			if genMatching:
124			genElectronTree, genElectrons, weight = gammaSelectionClone( tree, "genElectronTree", weight, "tree::Particle","genElectron" )
125			histograms = histoDefinition()
126	kiesel	1.1
127	kiesel	1.3	emObjects = ROOT.std.vector("tree::Photon")()
128	kiesel	1.1
129	kiesel	1.3	for event in tree:
130			if not event.GetReadEntry()%100000:
131			print '{0}%\r'.format(100*event.GetReadEntry()/event.GetEntries())
132	kiesel	1.2
133	kiesel	1.3	if event.GetReadEntry() > processNEvents:
134			break
135	kiesel	1.1
136	kiesel	1.3	weight[0] = event.weight * nExpected / processNEvents
137			emObjects.clear()
138	kiesel	1.1	photons.clear()
139			photonElectrons.clear()
140			photonJets.clear()
141	kiesel	1.3	if genMatching:
142			genElectrons.clear()
143	kiesel	1.1
144			for gamma in event.photon:
145
146			# cuts for every object to rject spikes
147			if gamma.r9 < 1 \
148			and gamma.sigmaIetaIeta > 0.001 \
149	kiesel	1.3	and abs(gamma.eta) < 1.4442 \
150	kiesel	1.1	and gamma.pt > 20:
151
152			# look for gamma and electrons
153			if gamma.hadTowOverEm < 0.05 \
154			and gamma.sigmaIetaIeta < 0.012 \
155			and gamma.chargedIso < 2.6 \
156			and gamma.neutralIso < 3.5 + 0.04*gamma.pt \
157			and gamma.photonIso < 1.3 + 0.005*gamma.pt:
158	kiesel	1.3	emObjects.push_back( gamma )
159	kiesel	1.1
160			# QCD fake object definition
161	kiesel	1.2	if gamma.ptJet > 75 \
162	kiesel	1.1	and gamma.hadTowOverEm < 0.05 \
163			and gamma.sigmaIetaIeta < 0.014 \
164			and gamma.chargedIso < 15 \
165			and gamma.neutralIso < 3.5 + 0.04*gamma.pt \
166			and gamma.photonIso < 1.3 + 0.005*gamma.pt \
167			and not gamma.pixelseed \
168			and ( gamma.sigmaIetaIeta>=0.012 or gamma.chargedIso>=2.6):
169			photonJets.push_back( gamma )
170
171	kiesel	1.3	if genMatching:
172			emObjects, histograms = generalMatching( emObjects, event.genElectron, histograms )
173			for genE in event.genElectron:
174			if abs(genE.eta) < 1.4442:
175			genElectrons.push_back( genE )
176			genElectrons, histograms = generalMatching( genElectrons, emObjects, histograms )
177
178			for emObject in emObjects:
179			if emObject.pixelseed:
180			photonElectrons.push_back( emObject )
181			else:
182			photons.push_back( emObject )
183	kiesel	1.1
184			if photons.size() > 0:
185			photonTree.Fill()
186			else:
187			if photonElectrons.size() > 0:
188			photonElectronTree.Fill()
189			if photonJets.size() > 0:
190			photonJetTree.Fill()
191	kiesel	1.3	if genMatching and genElectrons.size() > 0:
192			genElectronTree.Fill()
193	kiesel	1.1
194
195	kiesel	1.3	outputFileName = "slim"+inputFileName
196	kiesel	1.1	f = ROOT.TFile( outputFileName, "recreate" )
197			f.cd()
198			photonTree.Write()
199			photonJetTree.Write()
200			photonElectronTree.Write()
201	kiesel	1.3	if genMatching:
202			genElectronTree.Write()
203			draw_histogram_dict( histograms, shortName )
204	kiesel	1.1	f.Close()
205
206
207			if __name__ == "__main__":
208			# include knowledge about objects saved in the tree
209			ROOT.gSystem.Load("libTreeObjects.so")
210
211			arguments = argparse.ArgumentParser( description="Slim tree" )
212	kiesel	1.3	arguments.add_argument("--input", default=["WJets_V01.12_tree.root"], nargs="+" )
213			arguments.add_argument("--test", action="store_true" )
214			arguments.add_argument("--genMatching", action="store_true" )
215	kiesel	1.1	opts = arguments.parse_args()
216
217	kiesel	1.3	processNEvents = 10000 if opts.test else -1
218
219
220			# disable open canvas
221	kiesel	1.2	ROOT.gROOT.SetBatch()
222	kiesel	1.1
223	kiesel	1.3	datasetConfigName = "dataset.cfg"
224	kiesel	1.1	datasetConf = ConfigParser.SafeConfigParser()
225	kiesel	1.3	datasetConf.read( datasetConfigName )
226	kiesel	1.1
227			integratedLumi = 19300 #pb
228
229			for inName in opts.input:
230	kiesel	1.2	shortName = None
231	kiesel	1.1	for configName in datasetConf.sections():
232			if inName.count( configName ):
233	kiesel	1.2	shortName = configName
234	kiesel	1.1	crosssection = datasetConf.getfloat( configName, "crosssection" )
235	kiesel	1.3	if not shortName:
236			print "No configuration for input file {} defined in '{}'".format(
237			inName, datasetConfigName )
238			continue
239	kiesel	1.2
240	kiesel	1.1	# N = L * sigma
241			nExpected = integratedLumi * crosssection
242
243	kiesel	1.3	splitCandidates( inName, shortName, nExpected, processNEvents, opts.genMatching )
244	kiesel	1.1