ViewVC Help

View File | Revision Log | Show Annotations | Root Listing

root/cvsroot/UserCode/STPol/sync/syncTopRefSel.py

Comparing UserCode/STPol/sync/syncTopRefSel.py (file contents):
Revision 1.1 by jpata, Fri Jul 20 08:06:25 2012 UTC vs.
Revision 1.2 by jpata, Mon Jul 23 08:55:35 2012 UTC

#	Line 1 | Line 1
1		import ROOT
2		from DataFormats.FWLite import Events, Handle
3		import sys
4	+	import time
5	+	import numpy
6
7	+	bTagCutCSVT = 0.898
8	+	bTagCutCSVM = 0.679
9	+	bTagCutCSVL = 0.244
10	+	jetTightPtCut = 40
11
12	<	#Load PyTables
13	<	from tables import *
12	>	#Define the input file
13	>	input_file = "~/singletop/sync_T_t_ntuple.root"
14	>	input_file = "/home/joosep/singletop/CMSSW/CMSSW_5_2_5_patch1/src/TopQuarkAnalysis/SingleTop/test/edmntuple_sync_T_t.root"
15	>	#input_file = "/hdfs/local/stpol/ntuples/WJets_part_1Merged.root"
16
17
18	<	#Define the event database row
19	<	class Event(IsDescription):
20	<	eventId = Int32Col()
21	<	runId = Int32Col()
22	<	lumiId = Int32Col()
23	<	nTightLeptons = Int32Col()
24	<	leptonFlavour = Int32Col()
25	<	leptonEta = Float32Col()
26	<	jetPt0 = Float32Col()
18	>	class Source:
19	>	def __init__(self, label, outLabel=None, datatype="vector<float>"):
20	>	self.datatype = datatype
21	>	self.handle = Handle(self.datatype)
22	>	self.label = label
23	>	self._data = None
24	>	self.outLabel = outLabel
25	>
26	>	def get(self, event):
27	>	event.getByLabel(self.label, self.handle)
28	>	#self._data = [x for x in self.handle.product()]
29	>	self._data = self.handle.product()
30	>
31	>	def getData(self):
32	>	return self._data
33	>
34	>	class TreeSink:
35	>	def __init__(self):
36	>	self.tree = ROOT.TTree("events", "events")
37	>	self.branches = []
38	>	self.branchVariables = {}
39	>
40	>	def addBranch(self, name, variables):
41	>	branchStr = ""
42	>	for var in variables:
43	>	branchStr += var + "/F:"
44	>	branchStr = branchStr[:-1]
45	>	self.branchVariables[name] = numpy.empty((1,len(variables)), dtype=numpy.float32)
46	>	self.tree.Branch(name, self.branchVariables[name], branchStr)
47	>
48	>	def fillBranchVec(self, name, values):
49	>	if not name in self.branchVariables.keys():
50	>	raise Exception("Branch %s not found" % name)
51	>
52	>	shape = self.branchVariables[name].shape
53	>
54	>	if len(values)!= shape[1]:
55	>	raise Exception("Branch fill failed: shapes don't match")
56	>
57	>	for i in range(len(values)):
58	>	self.branchVariables[name][0,i] = values[i]
59	>
60	>	def fillBranchRepr(self, name, rep):
61	>	vec = rep.getVec()
62	>	n = len(vec)
63	>	for i in range(n):
64	>	self.branchVariables[name][0,i] = vec[i]
65	>
66	>	def fillBranches(self, name, reprs, upTo=10):
67	>	i = 0
68	>	for r in reprs[0:upTo]:
69	>	self.fillBranchRepr("%s%d" % (name, i), r)
70	>	i += 1
71	>
72	>	def initBranches(self):
73	>	for (k,v) in self.branchVariables.items():
74	>	v.fill(float("nan"))
75	>
76	>	def fillTree(self):
77	>	self.tree.Fill()
78	>
79	>	class Repr(object):
80	>	varOrder = []
81	>	def __init__(self):
82	>	self.variables = {}
83	>
84	>	def getVec(self):
85	>	return [self.variables[name] for name in self.varOrder]
86	>
87	>	@staticmethod
88	>	def make(sources, outType):
89	>	n = len(sources.values()[0].getData())
90	>	products = [None]*n
91	>	for i in range(n):
92	>	prod = outType()
93	>	for (name, source) in sources.items():
94	>	prod.variables[source.outLabel] = source.getData()[i]
95	>	products[i] = prod
96	>	return products
97	>
98	>	#def __repr__(self):
99	>	#    return str(self.variables)
100	>
101	>	class ParticleRepr(Repr):
102	>	def __init__(self):
103	>	super(ParticleRepr, self).__init__()
104	>
105	>	class METRepr(ParticleRepr):
106	>	def __init__(self):
107	>	super(METRepr, self).__init__()
108	>
109	>	def calcValues(self):
110	>	self.variables["Px"] = self.variables["Pt"]*ROOT.TMath.Cos(self.variables["Phi"])
111	>	self.variables["Py"] = self.variables["Pt"]*ROOT.TMath.Sin(self.variables["Phi"])
112	>
113	>	def calcW_MT(self, lepton):
114	>	lepPx = lepton.variables["Pt"]*ROOT.TMath.Cos(lepton.variables["Phi"])
115	>	lepPy = lepton.variables["Pt"]*ROOT.TMath.Sin(lepton.variables["Phi"])
116	>	return ROOT.TMath.Sqrt(ROOT.TMath.Power(lepton.variables["Pt"] + self.variables["Pt"], 2) - \
117	>	ROOT.TMath.Power(lepPx + self.variables["Px"], 2)  - \
118	>	ROOT.TMath.Power(lepPy + self.variables["Py"], 2))
119	>
120	>	class JetRepr(ParticleRepr):
121	>
122	>	def __init__(self):
123	>	super(JetRepr, self).__init__()
124	>
125	>	def isBJet(self, cut):
126	>	return self.variables["bTag"] > cut
127	>
128	>	def isTightJet(self):
129	>	return self.variables["Pt"] > jetTightPtCut
130	>
131	>	class LeptonRepr(ParticleRepr):
132	>
133	>	def __init__(self):
134	>	super(LeptonRepr, self).__init__()
135	>
136	>	class MuonRepr(LeptonRepr):
137	>
138	>	def __init__(self):
139	>	super(MuonRepr, self).__init__()
140
141	<	#Create the database file
142	<	h5file = openFile("tutorial1.h5", mode = "w", title = "Test file")
141	>	def isTight(self):
142	>	pass
143
144	<	#Create table in the database
145	<	table = h5file.createTable("/", "EventTable", Event)
144	>	class ElectronRepr(LeptonRepr):
145	>	varOrder = LeptonRepr.varOrder
146
147	<	#Get a handle to the row
148	<	event_row = table.row
147	>	def __init__(self):
148	>	super(ElectronRepr, self).__init__()
149
150		ROOT.gROOT.SetBatch()        # don't pop up canvases
151
152	<	#Define the input file
153	<	input_file = "~/singletop/sync_T_t_ntuple.root"
152	>	def calcSumPt(mus, eles, jets):
153	>	sumPt = 0
154	>	for mu in mus:
155	>	sumPt += mu.variables["Pt"]
156	>	for ele in eles:
157	>	sumPt += ele.variables["Pt"]
158	>	for jet in jets:
159	>	sumPt += jet.variables["Pt"]
160	>	return sumPt
161
162
163		#Get the Event generator from the event file
164		events = Events(input_file)
165
166	<	# create handle outside of loop
39	<	handle  = Handle ("vector<float>")
40	<	muEtaLabel = ("nTupleMuons", "tightMuonsEta", "SingleTop")
41	<	eleEtaLabel = ("nTupleElectrons", "tightElectronsEta", "SingleTop")
42	<	jetPtLabel = ("nTupleTopJetsPF", "topJetsPFPt", "SingleTop")
166	>	processName = "SingleTop"
167
168	<	events.toBegin()
168	>	jetModuleName = "nTupleTopJetsPF"
169	>	jetProductPrefix = "topJetsPF"
170	>	jetSources = {}
171	>	jetInLabels = ["Pt", "Eta", "Phi", "E", "CombinedSecondaryVertexBJetTags", "Flavour"]
172	>	for s in jetInLabels:
173	>	jetSources[s] = Source(label=(jetModuleName, jetProductPrefix+s, processName), outLabel=s)
174	>	jetSources["CombinedSecondaryVertexBJetTags"].outLabel = "bTag"
175	>
176	>	muModuleName = "nTupleMuons"
177	>	muProductPrefix = "tightMuons"
178	>	muSources = {}
179	>	relIsoLabel = "PFDeltaCorrectedRelIso"
180	>	muInLabels = ["Pt", "Eta", "Phi", "E", relIsoLabel, "AbsoluteDB", "PVDz"]
181	>	for s in muInLabels:
182	>	muSources[s] = Source(label=(muModuleName, muProductPrefix+s, processName), outLabel=s)
183	>
184	>	eleModuleName = "nTupleElectrons"
185	>	eleProductPrefix = "tightElectrons"
186	>	eleSources = {}
187	>	eleInLabels = ["Pt", "Eta", "Phi", "E", relIsoLabel, "AbsoluteDB"]
188	>	for s in eleInLabels:
189	>	eleSources[s] = Source(label=(eleModuleName, eleProductPrefix+s, processName), outLabel=s)
190	>
191	>	for s in [muSources, eleSources]:
192	>	s["AbsoluteDB"].outLabel = "dB"
193	>	s[relIsoLabel].outLabel = "relIso"
194	>
195	>
196	>	metModuleName = "nTuplePatMETsPF"
197	>	metProductPrefix = "patMETsPF"
198	>	metSources = {}
199	>	metInLabels = ["Pt", "Phi"]
200	>	for s in metInLabels:
201	>	metSources[s] = Source(label=(metModuleName, metProductPrefix+s, processName), outLabel=s)
202	>
203	>	JetRepr.varOrder = [jetSources[k].outLabel for k in jetInLabels]
204	>	MuonRepr.varOrder = [muSources[k].outLabel for k in muInLabels]
205	>	ElectronRepr.varOrder = [eleSources[k].outLabel for k in eleInLabels]
206	>	METRepr.varOrder = [metSources[k].outLabel for k in metInLabels]
207	>
208	>	sources = jetSources.values()+muSources.values()+eleSources.values()+metSources.values()
209	>	print JetRepr.varOrder
210	>	print MuonRepr.varOrder
211	>	print ElectronRepr.varOrder
212	>
213	>	outFile = ROOT.TFile("out.root", "RECREATE")
214	>
215	>	#Create the output TTree
216	>	treeOut = TreeSink()
217	>
218	>	maxN_BJets = 2
219	>	for i in range(maxN_BJets):
220	>	treeOut.addBranch("bjet%d" % i, JetRepr.varOrder)
221	>
222	>	maxN_Jets = 3
223	>	for i in range(maxN_Jets):
224	>	treeOut.addBranch("jet%d" % i, JetRepr.varOrder)
225	>
226	>	treeOut.addBranch("id", ["run", "lumi", "event"])
227	>	treeOut.addBranch("fstate", ["nJets", "nBTagsCSVM", "nBTagsCSVL", "nMuons", "nEles", "mtwMass", "sumPt"])
228	>	treeOut.addBranch("trigger", ["passesHLT"])
229	>	treeOut.addBranch("met", ["Pt", "Phi"])
230	>
231	>	maxNMuons = 1
232	>	for i in range(maxNMuons):
233	>	treeOut.addBranch("mu%d"%i, MuonRepr.varOrder)
234	>	maxNEles = 3
235	>	for i in range(maxNEles):
236	>	treeOut.addBranch("ele%d"%i, ElectronRepr.varOrder)
237
238		# loop over event
239	+	totalEvents = events.size()
240	+	processedEvents = 0
241	+	print "Running over %d events" % events.size()
242	+	t1 = time.time()
243	+
244	+	class Triggers:
245	+	def __init__(self):
246	+	self.statusHandle = Handle("edm::TriggerResults")
247	+	#self.nameHandle = Handle("edm::TriggerNames")
248	+	self.trigLabel = ("TriggerResults", "", "HLT")
249	+
250	+	def get(self, event):
251	+	event.getByLabel(self.trigLabel, self.statusHandle)
252	+	self.triggerResults = self.statusHandle.product()
253	+	self.names = event.object().triggerNames(self.triggerResults)
254	+
255	+
256	+	def passesHLT(self, name, isPrecise=True):
257	+	if isPrecise:
258	+	return self.triggerResults.accept(self.names.triggerIndex(name))
259	+	else:
260	+	n = [[self.names.triggerIndex(n), n] for n in self.names.triggerNames()]
261	+	(idx, name) = filter(lambda x: name in x[1], n)[0]
262	+	return self.triggerResults.accept(idx)
263	+
264	+	triggers = Triggers()
265	+
266		for event in events:
267
268	<	#Get the event contents
269	<	event.getByLabel (muEtaLabel, handle)
270	<	muEtas = handle.product()
271	<	event.getByLabel (eleEtaLabel, handle)
272	<	eleEtas = handle.product()
54	<	event.getByLabel (jetPtLabel, handle)
55	<	jetPts = handle.product()
56	<
57	<	#Convert the std::vector to python list
58	<	jetPts = [x for x in jetPts]
59	<
60	<	#Get the first 1 jet PT-s
61	<	i = 0
62	<	for jetPt in jetPts[0:1]:
63	<	event_row["jetPt%d" % i] = jetPt
64	<	i += 1
65	<
66	<	nMu = len(muEtas)
67	<	nEle = len(eleEtas)
68	<	nTightLeptons = nMu + nEle
69	<	event_row["nTightLeptons"] = nTightLeptons
70	<
71	<	if nTightLeptons != 1:
72	<	leptonEta = float("nan")
73	<	leptonFlavour = 0
74	<	elif nEle==1:
75	<	leptonEta = eleEtas[0]
76	<	leptonFlavour = 1
77	<	elif nMu==1:
78	<	leptonEta = muEtas[0]
79	<	leptonFlavour = 2
80	<	#for val in vals:
81	<	#    print val
82	<	#for mu in muons:
83	<	#    print mu.userFloat("pt")
268	>	treeOut.initBranches()
269	>
270	>	for source in sources:
271	>	source.get(event)
272	>	triggers.get(event)
273
274		eventId = event.object().id().event()
86	–	lumiId = event.object().id().luminosityBlock()
275		runId = event.object().id().run()
276	<	event_row["eventId"] = eventId
89	<	event_row["runId"] = runId
90	<	event_row["lumiId"] = lumiId
91	<	event_row["leptonEta"] = leptonEta
92	<	event_row["leptonFlavour"] = leptonFlavour
93	<
94	<	#Put the event row into the table
95	<	event_row.append()
96	<	#print "%d:%d:%d" % (runId, lumiId, eventId)
97	<
98	<	#FLush the table from memory to disk
99	<	table.flush()
100	<
101	<	#Print out some debugging events
102	<	evs = [(x["eventId"], x["leptonEta"], x["leptonFlavour"], x["jetPt0"]) for x in table.where("(eventId >= 7802) & (eventId <= 7900) & (nTightLeptons==1)")]
276	>	lumiId = event.object().id().luminosityBlock()
277
278	<	for ev in evs:
279	<	print ev
278	>	jets = ParticleRepr.make(jetSources, JetRepr)
279	>	muons = ParticleRepr.make(muSources, MuonRepr)
280	>	eles = ParticleRepr.make(eleSources, ElectronRepr)
281	>
282	>	nMuons = len(muons)
283	>	nEles = len(eles)
284	>
285	>	if nMuons>=1:
286	>	isoLepton = muons[0]
287	>	elif nEles>=1:
288	>	isoLepton = eles[0]
289	>	else:
290	>	isoLepton = None
291	>
292	>	mets = ParticleRepr.make(metSources, METRepr)
293	>	mets[0].calcValues()
294	>
295	>	if isoLepton!=None:
296	>	mtwMass = mets[0].calcW_MT(isoLepton)
297	>	else:
298	>	mtwMass = float("nan")
299	>
300	>	tightJets = filter(lambda jet: jet.isTightJet(), jets)
301	>	nJets = len(jets)
302	>	looseBJets = filter(lambda jet: jet.isBJet(bTagCutCSVL), tightJets)
303	>	mediumBJets = filter(lambda jet: jet.isBJet(bTagCutCSVM), looseBJets)
304	>	nBTagsCSVM = len(mediumBJets)
305	>	nBTagsCSVL = len(looseBJets)
306	>
307	>	treeOut.fillBranchVec("id", [runId, lumiId, eventId])
308	>	treeOut.fillBranches("bjet", mediumBJets, upTo=maxN_BJets)
309	>	treeOut.fillBranches("jet", jets, upTo=maxN_Jets)
310	>	treeOut.fillBranches("ele", eles, upTo=maxNEles)
311	>	treeOut.fillBranches("mu", muons, upTo=maxNMuons)
312	>	treeOut.fillBranchVec("fstate", (nJets, nBTagsCSVM, nBTagsCSVL, nMuons, nEles, mtwMass, calcSumPt(muons, eles, jets)))
313	>	treeOut.fillBranchVec("trigger", [triggers.passesHLT("HLT_IsoMu24_eta2p1_v11")])
314	>	treeOut.fillBranchVec("met", (mets[0].variables["Pt"], mets[0].variables["Phi"]))
315	>	treeOut.fillTree()
316	>
317	>	if processedEvents%5000 == 0:
318	>	print "Done %d/%d in %.1f sec" % (processedEvents, totalEvents, time.time()-t1)
319	>
320	>	processedEvents += 1
321	>
322	>	t2 = time.time()
323	>
324	>	elapsed = t2-t1
325	>	print "Time elapsed: %.2f sec" % elapsed
326	>	print "events per second: %d" % (int(events.size()/elapsed))
327	>	outFile.Write()
328	>	print "Step HLT: %d" % (treeOut.tree.GetEntries("trigger.passesHLT==1.0"))
329	>	print "Step IsoMu: %d" % (treeOut.tree.GetEntries("trigger.passesHLT==1.0 && mu0.relIso<0.12 && mu0.Pt>26.0"))
330	>	print "Step Jet: %d" % (treeOut.tree.GetEntries("trigger.passesHLT==1.0 && mu0.relIso<0.12 && mu0.Pt>26.0 && fstate.nJets==2.0"))
331	>	print "Step W-mass: %d" % (treeOut.tree.GetEntries("trigger.passesHLT==1.0 && mu0.relIso<0.12 && mu0.Pt>26.0 && fstate.nJets==2.0 && fstate.mtwMass >= 40.0"))
332	>	print "Step B: %d" % (treeOut.tree.GetEntries("trigger.passesHLT==1.0 && mu0.relIso<0.12 && mu0.Pt>26.0 && fstate.nJets==2.0 && fstate.mtwMass >= 40.0 && fstate.nBTagsCSVM == 1.0"))
333	>	outFile.Close()

Diff Legend

–	Removed lines
+	Added lines
<	Changed lines
>	Changed lines