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
#	Content
1	#! /usr/bin/env python2
2	# -- coding: utf-8 --
3	from treeFunctions import *
4	import Styles
5	Styles.tdrStyle2D()
6	import numpy
7	import ROOT
8	import argparse
9	import ConfigParser
10	import sys
11
12	def deltaPhi( phi1, phi2):
13	"""Computes delta phi.
14	Due to the zylindrical form of CMS, it has be be corrected by a diffenence
15	of 2pi.
16	"""
17	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	"""Compute delta R of two objects.
27	Both objects must have eta and phi as member variables.
28	"""
29	import math
30	return math.sqrt( (object1.eta-object2.eta)2 + (deltaPhi(object1.phi,object2.phi))2 )
31
32	def gammaSelectionClone( tree, newTreeName, weight, type_="tree::Photon", name="photon" ):
33	"""Clones a tree and get access to photon vector"""
34	newTree = tree.CloneTree(0)
35	newTree.SetName( newTreeName )
36	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
53	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	minIndex = -1
77	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
114	tree = readTree( inputFileName )
115	eventHisto = readHisto( inputFileName )
116	if processNEvents == -1:
117	processNEvents = eventHisto.GetBinContent(1)
118
119	weight = numpy.zeros(1, dtype=float)
120	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
127	emObjects = ROOT.std.vector("tree::Photon")()
128
129	for event in tree:
130	if not event.GetReadEntry()%100000:
131	print '{0}%\r'.format(100*event.GetReadEntry()/event.GetEntries())
132
133	if event.GetReadEntry() > processNEvents:
134	break
135
136	weight[0] = event.weight * nExpected / processNEvents
137	emObjects.clear()
138	photons.clear()
139	photonElectrons.clear()
140	photonJets.clear()
141	if genMatching:
142	genElectrons.clear()
143
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	and abs(gamma.eta) < 1.4442 \
150	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	emObjects.push_back( gamma )
159
160	# QCD fake object definition
161	if gamma.ptJet > 75 \
162	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	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
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	if genMatching and genElectrons.size() > 0:
192	genElectronTree.Fill()
193
194
195	outputFileName = "slim"+inputFileName
196	f = ROOT.TFile( outputFileName, "recreate" )
197	f.cd()
198	photonTree.Write()
199	photonJetTree.Write()
200	photonElectronTree.Write()
201	if genMatching:
202	genElectronTree.Write()
203	draw_histogram_dict( histograms, shortName )
204	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	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	opts = arguments.parse_args()
216
217	processNEvents = 10000 if opts.test else -1
218
219
220	# disable open canvas
221	ROOT.gROOT.SetBatch()
222
223	datasetConfigName = "dataset.cfg"
224	datasetConf = ConfigParser.SafeConfigParser()
225	datasetConf.read( datasetConfigName )
226
227	integratedLumi = 19300 #pb
228
229	for inName in opts.input:
230	shortName = None
231	for configName in datasetConf.sections():
232	if inName.count( configName ):
233	shortName = configName
234	crosssection = datasetConf.getfloat( configName, "crosssection" )
235	if not shortName:
236	print "No configuration for input file {} defined in '{}'".format(
237	inName, datasetConfigName )
238	continue
239
240	# N = L * sigma
241	nExpected = integratedLumi * crosssection
242
243	splitCandidates( inName, shortName, nExpected, processNEvents, opts.genMatching )
244