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.3 by jpata, Thu Jul 26 19:32:45 2012 UTC

#	Line 1 | Line 1
1		import ROOT
2		from DataFormats.FWLite import Events, Handle
3		import sys
4	<
5	<
6	<	#Load PyTables
7	<	from tables import *
8	<
9	<
10	<	#Define the event database row
11	<	class Event(IsDescription):
12	<	eventId = Int32Col()
13	<	runId = Int32Col()
14	<	lumiId = Int32Col()
15	<	nTightLeptons = Int32Col()
16	<	leptonFlavour = Int32Col()
17	<	leptonEta = Float32Col()
18	<	jetPt0 = Float32Col()
19	<
20	<	#Create the database file
21	<	h5file = openFile("tutorial1.h5", mode = "w", title = "Test file")
22	<
23	<	#Create table in the database
24	<	table = h5file.createTable("/", "EventTable", Event)
25	<
26	<	#Get a handle to the row
27	<	event_row = table.row
28	<
29	<	ROOT.gROOT.SetBatch()        # don't pop up canvases
4	>	import time
5	>	import numpy
6	>	import ctypes
7	>
8	>	class Cuts:
9	>	class bTag:
10	>	""" Based on https://twiki.cern.ch/twiki/bin/view/CMS/BTagPerformanceOP """
11	>	class CSV:
12	>	tight = 0.898
13	>	medium = 0.679
14	>	loose = 0.244
15	>	class TCHP:
16	>	tight = 3.41
17	>	medium = 1.93
18	>
19	>	class NaNs:
20	>	floatNAN = float("nan")
21	>	intNAN = 99999
22
23		#Define the input file
24	<	input_file = "~/singletop/sync_T_t_ntuple.root"
25	<
26	<
27	<	#Get the Event generator from the event file
28	<	events = Events(input_file)
29	<
30	<	# create handle outside of loop
31	<	handle  = Handle ("vector<float>")
32	<	muEtaLabel = ("nTupleMuons", "tightMuonsEta", "SingleTop")
33	<	eleEtaLabel = ("nTupleElectrons", "tightElectronsEta", "SingleTop")
34	<	jetPtLabel = ("nTupleTopJetsPF", "topJetsPFPt", "SingleTop")
35	<
36	<	events.toBegin()
37	<
38	<	# loop over event
39	<	for event in events:
40	<
41	<	#Get the event contents
42	<	event.getByLabel (muEtaLabel, handle)
43	<	muEtas = handle.product()
44	<	event.getByLabel (eleEtaLabel, handle)
45	<	eleEtas = handle.product()
46	<	event.getByLabel (jetPtLabel, handle)
47	<	jetPts = handle.product()
48	<
49	<	#Convert the std::vector to python list
50	<	jetPts = [x for x in jetPts]
51	<
52	<	#Get the first 1 jet PT-s
53	<	i = 0
54	<	for jetPt in jetPts[0:1]:
55	<	event_row["jetPt%d" % i] = jetPt
56	<	i += 1
57	<
58	<	nMu = len(muEtas)
59	<	nEle = len(eleEtas)
60	<	nTightLeptons = nMu + nEle
61	<	event_row["nTightLeptons"] = nTightLeptons
62	<
63	<	if nTightLeptons != 1:
64	<	leptonEta = float("nan")
65	<	leptonFlavour = 0
66	<	elif nEle==1:
67	<	leptonEta = eleEtas[0]
68	<	leptonFlavour = 1
69	<	elif nMu==1:
70	<	leptonEta = muEtas[0]
71	<	leptonFlavour = 2
72	<	#for val in vals:
73	<	#    print val
74	<	#for mu in muons:
75	<	#    print mu.userFloat("pt")
76	<
77	<	eventId = event.object().id().event()
78	<	lumiId = event.object().id().luminosityBlock()
79	<	runId = event.object().id().run()
80	<	event_row["eventId"] = eventId
81	<	event_row["runId"] = runId
82	<	event_row["lumiId"] = lumiId
83	<	event_row["leptonEta"] = leptonEta
84	<	event_row["leptonFlavour"] = leptonFlavour
85	<
86	<	#Put the event row into the table
87	<	event_row.append()
88	<	#print "%d:%d:%d" % (runId, lumiId, eventId)
89	<
90	<	#FLush the table from memory to disk
91	<	table.flush()
24	>	#input_file = "~/singletop/sync_T_t_ntuple.root"
25	>	#input_file = "~/singletop/edmntuple_T_t_sync.root"
26	>	#input_file = "/hdfs/local/stpol/ntuples/WJets_part_1Merged.root"
27	>	#input_file = "/hdfs/local/stpol/joosep/ntuples_orso_new/WJets_part_1Merged.root"
28	>	#channel = "TTBar"
29	>
30	>	class InputOutputConf:
31	>	@staticmethod
32	>	def inFile(channel, part):
33	>	return "/hdfs/local/stpol/ntuples/%s_part_%dMerged.root" % (channel, part)
34	>
35	>	@staticmethod
36	>	def outFile(channel, part):
37	>	return "trees_%s_%d.root" % (channel, part)
38	>
39	>	channel = sys.argv[1]
40	>
41	>	if channel != "sync":
42	>	part = int(sys.argv[2])
43	>	input_file = InputOutputConf.inFile(channel, part)
44	>	output_file = InputOutputConf.outFile(channel, part)
45	>	else:
46	>	input_file = "~/singletop/edmntuple_T_t_sync.root"
47	>	output_file = "trees_T_t_SYNC.root"
48	>
49	>	class BaseStruct(ctypes.Structure):
50	>
51	>	def rowStr(self):
52	>	out = ""
53	>	for (name, dType) in self._fields_:
54	>	out += name + "/" + BaseStruct.ctypeToROOT(dType) + ":"
55	>	out = out[:-1]
56	>	return out
57	>
58	>	@staticmethod
59	>	def ctypeToROOT(dType):
60	>	if dType == ctypes.c_int:
61	>	return "I"
62	>	elif dType == ctypes.c_float:
63	>	return "F"
64	>	elif dType == ctypes.c_uint:
65	>	return "i"
66	>	else:
67	>	raise Exception("Can't find suitable ROOT data type for %s" % dType.__name__)
68	>
69	>	@staticmethod
70	>	def getNaN(dType):
71	>	if dType == ctypes.c_int:
72	>	return NaNs.intNAN
73	>	elif dType == ctypes.c_float:
74	>	return NaNs.floatNAN
75	>	elif dType == ctypes.c_uint:
76	>	return NaNs.intNAN
77	>	else:
78	>	raise Exception("NaN not defined for datatype %s" % dType.__name__)
79	>
80	>
81	>	def initVars(self):
82	>	for (name, dType) in self._fields_:
83	>	self.__setattr__(name, BaseStruct.getNaN(dType))
84	>
85	>	class Source:
86	>	def __init__(self, label, datatype, outLabel = None):
87	>	self.datatype = datatype
88	>	self.handle = Handle(self.datatype)
89	>	self.label = label
90	>	self._data = None
91	>	if outLabel==None:
92	>	if not isinstance(label, basestring) and isinstance(label, (list, tuple)): #label is a list
93	>	outLabel = label[1]
94	>	else:
95	>	outLabel = label
96	>	self.outLabel = outLabel
97	>
98	>	def get(self, event):
99	>	event.getByLabel(self.label, self.handle)
100	>	self._data = self.handle.product()
101	>
102	>	def getData(self):
103	>	return self._data
104	>
105	>	class TreeSink:
106	>	def __init__(self):
107	>	self.tree = ROOT.TTree("events", "events")
108	>	self.branches = []
109	>	self.branchVariables = {}
110	>
111	>	def addBranch(self, name, variables):
112	>	branchStr = ""
113	>	for var in variables:
114	>	branchStr += var + "/F:"
115	>	branchStr = branchStr[:-1]
116	>	self.branchVariables[name] = numpy.empty((1,len(variables)), dtype=numpy.float32)
117	>	self.tree.Branch(name, self.branchVariables[name], branchStr)
118	>
119	>	def fillBranchVec(self, name, values):
120	>	if not name in self.branchVariables.keys():
121	>	raise Exception("Branch %s not found" % name)
122	>
123	>	shape = self.branchVariables[name].shape
124	>
125	>	if len(values)!= shape[1]:
126	>	raise Exception("Branch fill failed: shapes don't match")
127	>
128	>	for i in range(len(values)):
129	>	self.branchVariables[name][0,i] = values[i]
130	>
131	>	def fillBranchRepr(self, name, rep):
132	>	vec = rep.getVec()
133	>	n = len(vec)
134	>	for i in range(n):
135	>	self.branchVariables[name][0,i] = vec[i]
136	>
137	>	def fillBranches(self, name, reprs, upTo=10):
138	>	i = 0
139	>	for r in reprs[0:upTo]:
140	>	self.fillBranchRepr("%s%d" % (name, i), r)
141	>	i += 1
142	>
143	>	def initBranches(self):
144	>	for (k,v) in self.branchVariables.items():
145	>	v.fill(float("nan"))
146	>
147	>	def fillTree(self):
148	>	self.tree.Fill()
149	>
150	>	class Repr(object):
151	>	varOrder = []
152	>	def __init__(self):
153	>	self.variables = {}
154	>
155	>	def getVec(self):
156	>	return [self.variables[name] for name in self.varOrder]
157	>
158	>	@staticmethod
159	>	def make(sources, outType):
160	>	n = len(sources.values()[0].getData())
161	>	products = [None]*n
162	>	for i in range(n):
163	>	prod = outType()
164	>	for (name, source) in sources.items():
165	>	prod.variables[source.outLabel] = source.getData()[i]
166	>	products[i] = prod
167	>	return products
168	>
169	>	#def __repr__(self):
170	>	#    return str(self.variables)
171	>
172	>	class ParticleRepr(Repr):
173	>	class EvStruct(BaseStruct):
174	>	_fields_ = [('pt', ctypes.c_float),
175	>	('eta', ctypes.c_float),
176	>	('phi', ctypes.c_float),
177	>	('E', ctypes.c_float),
178	>	]
179	>
180	>	def __init__(self):
181	>	super(ParticleRepr, self).__init__()
182	>
183	>
184	>	def pt(self):
185	>	return self.variables["Pt"]
186	>
187	>	def eta(self):
188	>	return self.variables["Eta"]
189	>
190	>	def phi(self):
191	>	return self.variables["Phi"]
192	>
193	>	def E(self):
194	>	return self.variables["E"]
195	>
196	>	def fillStruct(self, s):
197	>	s.pt = self.variables["Pt"]
198	>	s.eta = self.variables["Eta"]
199	>	s.phi = self.variables["Phi"]
200	>	s.E = self.variables["E"]
201	>
202	>	class METRepr(ParticleRepr):
203	>	class EvStruct(BaseStruct):
204	>	_fields_ = [('pt', ctypes.c_float), ('phi', ctypes.c_float), ('px', ctypes.c_float), ('py', ctypes.c_float)]
205	>
206	>	def __init__(self):
207	>	super(METRepr, self).__init__()
208	>
209	>	def calcValues(self):
210	>	self.variables["Px"] = self.pt()*ROOT.TMath.Cos(self.phi())
211	>	self.variables["Py"] = self.pt()*ROOT.TMath.Sin(self.phi())
212	>
213	>	def calcW_MT(self, lepton):
214	>	lepPx = lepton.pt()*ROOT.TMath.Cos(lepton.phi())
215	>	lepPy = lepton.pt()*ROOT.TMath.Sin(lepton.phi())
216	>	return ROOT.TMath.Sqrt(ROOT.TMath.Power(lepton.pt() + self.pt(), 2) - ROOT.TMath.Power(lepPx + self.px(), 2)  - ROOT.TMath.Power(lepPy + self.py(), 2))
217	>
218	>	def calc_topMass(self, lepton, jet):
219	>	MW = 80.399
220	>	lepVec = ROOT.TLorentzVector()
221	>	lepVec.SetPtEtaPhiE(lepton.pt(), lepton.eta(), lepton.phi(), lepton.E())
222	>	jetVec = ROOT.TLorentzVector()
223	>	jetVec.SetPtEtaPhiE(jet.pt_corr(), jet.eta(), jet.phi(), jet.E_corr())
224	>
225	>	metVec = ROOT.TVector2()
226	>	metVec.SetMagPhi(met.pt(), met.phi())
227	>
228	>	Lambda = pow(MW, 2) / 2 + lepVec.Px()*metVec.Px() + lepVec.Py()*metVec.Py()
229	>	#a = Lambda*lepVec.Pz() / lepVec.Pt()**2
230	>	#b = (lepVec.E()**2*metVec.Mod()**2-Lambda**2)/lepVec.Pt()**2
231	>	Delta = pow(lepVec.E(), 2) * (pow(Lambda,2) - pow(lepVec.Pt()*metVec.Mod(),2) )
232	>	sk1 = lepVec.Pt()*metVec.Mod()
233	>	sk2 = metVec.Px()*lepVec.Px() + metVec.Py()*lepVec.Py()
234	>	nuVec = ROOT.TLorentzVector()
235	>	if Delta > 0:
236	>	r = ROOT.TMath.Sqrt(Delta)
237	>	A = (Lambda*lepVec.Pz()+r)/pow(lepVec.Pt(), 2)
238	>	B = (Lambda*lepVec.Pz()-r)/pow(lepVec.Pt(), 2)
239	>
240	>	#get the root with the minimum absolute value
241	>	p_nu_z = min((A, abs(A)), (B, abs(B)), key=lambda x:x[1])[0]
242	>	E_nu = ROOT.TMath.Sqrt(metVec.Mod()**2 + p_nu_z**2)
243	>	nuVec.SetPxPyPzE(metVec.Px(), metVec.Py(), p_nu_z, E_nu)
244	>	top = ROOT.TLorentzVector()
245	>	top = lepVec+jetVec+nuVec
246	>	return top.M()
247	>
248	>	else:
249	>	MW_new = ROOT.TMath.Sqrt(2*(sk1-sk2))
250	>	Lambda_new = pow(MW_new, 2) / 2.0 + sk2
251	>	Delta_new = (pow(Lambda_new, 2) - pow(sk1, 2))*pow(lepVec.E(), 2)
252	>	p_nu_z = (Lambda_new*lepVec.Pz())/pow(lepVec.Pt(), 2)
253	>
254	>	E_nu = ROOT.TMath.Sqrt(metVec.Mod()**2 + p_nu_z**2)
255	>	nuVec.SetPxPyPzE(metVec.Px(), metVec.Py(), p_nu_z, E_nu)
256	>	top = ROOT.TLorentzVector()
257	>	top = lepVec+jetVec+nuVec
258	>	return -top.M()
259	>
260	>
261	>	E_nu = ROOT.TMath.Sqrt(metVec.Mod()**2 + p_nu_z**2)
262	>	nuVec.SetPxPyPzE(metVec.Px(), metVec.Py(), p_nu_z, E_nu)
263	>	top = ROOT.TLorentzVector()
264	>	top = lepVec+jetVec+nuVec
265	>	return top.M()
266	>
267	>	def fillStruct(self, s):
268	>	""" Fills the given MET EvStruct with the variable values. DON'T call the ParticleRepr.fillStruct(), since MET doesn't have eta/E. """
269	>	#super(METRepr, self).fillStruct(s)
270	>	s.pt = self.variables["Pt"]
271	>	s.phi = self.variables["Phi"]
272	>	s.px = self.variables["Px"]
273	>	s.py = self.variables["Py"]
274	>
275	>	def px(self):
276	>	return self.variables["Px"]
277	>
278	>	def py(self):
279	>	return self.variables["Py"]
280	>
281	>	def __str__(self):
282	>	return "MET pt: %.5f phi: %.5f" % (self.pt(), self.phi())
283	>
284	>	class JetRepr(ParticleRepr):
285	>	class EvStruct(BaseStruct):
286	>	_fields_ = ParticleRepr.EvStruct._fields_ + [('bTagTCHP', ctypes.c_float), ('pt_corr', ctypes.c_float), ('E_corr', ctypes.c_float), ('bTagCSV', ctypes.c_float)]
287	>
288	>	def fillStruct(self, s):
289	>	super(JetRepr, self).fillStruct(s)
290	>	s.bTagTCHP = self.bTagTCHP()
291	>	s.bTagCSV = self.bTagCSV()
292	>	s.pt_corr = self.pt_corr()
293	>	s.E_corr = self.E_corr()
294	>
295	>	def __init__(self):
296	>	super(JetRepr, self).__init__()
297	>
298	>	def passesBTagCut(self, cut):
299	>	return self.bTag() > cut
300	>
301	>	def isTightJet(self):
302	>	return self.variables["Pt"] > jetTightPtCut
303	>
304	>	def bTagTCHP(self):
305	>	return self.variables["bTagTCHP"]
306	>
307	>	def bTagCSV(self):
308	>	return self.variables["bTagCSV"]
309	>
310	>	def __str__(self):
311	>	bTagType = "N"
312	>	#if self.passesBTagCut(Cuts.bTagCutCSVL): bTagType = "L"
313	>	#if self.passesBTagCut(Cuts.bTagCutCSVM): bTagType = "M"
314	>	#if self.passesBTagCut(Cuts.bTagCutCSVT): bTagType = "T"
315	>	if self.passesBTagCut(Cuts.bTag.TCHP.tight): bTagType = "T"
316	>	return "jet pt: %.5f (%.5f) eta: %.5f phi: %.5f E: %.5f (%.5f) bTag: %.5f %s" % (self.pt(), self.pt_corr(), self.eta(), self.phi(), self.E(), self.E_corr(), self.bTag(), bTagType)
317	>
318	>	@staticmethod
319	>	def smearFactor(jet):
320	>	"""
321	>	Based on https://twiki.cern.ch/twiki/bin/view/CMS/JetResolution
322	>	"""
323	>	eta = ROOT.TMath.Abs(jet.eta())
324	>
325	>	#2010 recommendation
326	>	# if eta>0.0 and eta<1.1:
327	>	#   return 1.066
328	>	# elif eta>1.1 and eta<1.7:
329	>	#   return 1.191
330	>	# elif eta>1.7 and eta<2.3:
331	>	#   return 1.096
332	>	# elif eta>2.3 and eta<5.0:
333	>	#   return 1.166
334	>	# else:
335	>	#   return float("nan")
336	>
337	>	#2011 recommendation
338	>	if eta>0.0 and eta<0.5:
339	>	return 1.052
340	>	elif eta>0.5 and eta<1.1:
341	>	return 1.057
342	>	elif eta>1.1 and eta<1.7:
343	>	return 1.096
344	>	elif eta>1.7 and eta<2.3:
345	>	return 1.134
346	>	elif eta>2.3 and eta<5.0:
347	>	return 1.288
348	>	else:
349	>	return NaNs.floatNAN
350	>
351	>	def pt_corr(self):
352	>	return self.variables["pt_corr"]
353	>
354	>	def smear(self):
355	>	ptGen = self.variables["genJetPt"]
356	>	pt = self.pt()
357	>	resolutionSF = JetRepr.smearFactor(self)
358	>	smear = (max(0.0, ptGen + resolutionSF*(pt-ptGen)))/pt
359	>	return smear
360	>
361	>	def doMCSmearing(self):
362	>	s = self.smear()
363	>	self.variables["pt_corr"] = s*self.pt()
364	>	self.variables["E_corr"] = s*self.E()
365	>	return
366	>
367	>	def E_corr(self):
368	>	return self.variables["E_corr"]
369	>
370	>	class LeptonRepr(ParticleRepr):
371	>	class EvStruct(BaseStruct):
372	>	_fields_ = ParticleRepr.EvStruct._fields_ + [('relIso', ctypes.c_float)]
373	>
374	>	def fillStruct(self, s):
375	>	super(LeptonRepr, self).fillStruct(s)
376	>	s.relIso = self.variables["relIso"]
377	>
378	>	def __init__(self):
379	>	super(LeptonRepr, self).__init__()
380	>
381	>	def relIso(self):
382	>	return self.variables["relIso"]
383	>
384	>	class MuonRepr(LeptonRepr):
385	>
386	>	def __init__(self):
387	>	super(MuonRepr, self).__init__()
388	>
389	>	def isTight(self):
390	>	pass
391	>
392	>	class ElectronRepr(LeptonRepr):
393	>	varOrder = LeptonRepr.varOrder
394	>
395	>	def __init__(self):
396	>	super(ElectronRepr, self).__init__()
397	>
398	>	class Triggers:
399	>	def __init__(self):
400	>	self.statusHandle = Handle("edm::TriggerResults")
401	>	#self.nameHandle = Handle("edm::TriggerNames")
402	>	self.trigLabel = ("TriggerResults", "", "HLT")
403	>
404	>	def get(self, event):
405	>	event.getByLabel(self.trigLabel, self.statusHandle)
406	>	self.triggerResults = self.statusHandle.product()
407	>	self.names = event.object().triggerNames(self.triggerResults)
408	>
409	>
410	>	def passesHLT(self, name, isPrecise=True):
411	>	if isPrecise:
412	>	return self.triggerResults.accept(self.names.triggerIndex(name))
413	>	else:
414	>	n = [[self.names.triggerIndex(n), n] for n in self.names.triggerNames()]
415	>	(idx, name) = filter(lambda x: name in x[1], n)[0]
416	>	return self.triggerResults.accept(idx)
417	>
418	>	class EventStruct(BaseStruct):
419	>	_fields_ = [('eventID', ctypes.c_uint),
420	>	('runID', ctypes.c_uint),
421	>	('lumiID', ctypes.c_uint),
422	>	]
423	>
424	>	class FinalStateStruct(BaseStruct):
425	>	_fields_ = [('nJets', ctypes.c_int),
426	>	('nBTagsCSVL', ctypes.c_int),
427	>	('nBTagsCSVM', ctypes.c_int),
428	>	('nBTagsCSVT', ctypes.c_int),
429	>	('nBTagsTCHPT', ctypes.c_int),
430	>	('nMuons', ctypes.c_int),
431	>	('nEles', ctypes.c_int),
432	>	('sumJetPt', ctypes.c_float),
433	>	('M_W', ctypes.c_float),
434	>	('topMass', ctypes.c_float),
435	>	]
436	>
437	>	class TriggerStateStruct(BaseStruct):
438	>	_fields_ = [('HLT_Mu', ctypes.c_int),
439	>	('HLT_Ele', ctypes.c_int),
440	>	]
441	>
442	>	def calcSumPt(mus, eles, jets):
443	>	sumPt = 0
444	>	for mu in mus:
445	>	sumPt += mu.variables["Pt"]
446	>	for ele in eles:
447	>	sumPt += ele.variables["Pt"]
448	>	for jet in jets:
449	>	sumPt += jet.variables["Pt"]
450	>	return sumPt
451	>
452	>
453	>	debugging = False
454	>	doSelEv = False
455	>
456	>	if __name__=="__main__":
457	>
458	>	ROOT.gROOT.SetBatch()        # don't pop up canvases
459	>
460	>	#Get the Event generator from the event file
461	>	events = Events(input_file)
462	>
463	>	processName = "SingleTop"
464	>
465	>	jetModuleName = "nTupleTopJetsPF"
466	>	jetProductPrefix = "topJetsPF"
467	>	jetSources = {}
468	>	bTagSource = "TrackCountingHighPur"
469	>	jetInLabels = ["Pt", "Eta", "Phi", "E", "TrackCountingHighPur", "Flavour", "CombinedSecondaryVertexBJetTags"]
470	>	for s in jetInLabels:
471	>	jetSources[s] = Source(label=(jetModuleName, jetProductPrefix+s, processName), datatype="vector<float>", outLabel=s)
472	>
473	>	jetSources["genJetPt"] = Source(label=("genJetsPF", "genJetsPt", processName), datatype="vector<double>", outLabel="genJetPt")
474	>	jetSources["genJetEta"] = Source(label=("genJetsPF", "genJetsEta", processName), datatype="vector<double>", outLabel="genJetEta")
475	>	jetSources["TrackCountingHighPur"].outLabel = "bTagTCHP"
476	>	jetSources["CombinedSecondaryVertexBJetTags"].outLabel = "bTagCSV"
477	>
478	>	muModuleName = "nTupleMuons"
479	>	muProductPrefix = "tightMuons"
480	>	muSources = {}
481	>	muInLabels = ["Pt", "Eta", "Phi", "E", "PFDeltaCorrectedRelIso", "AbsoluteDB", "PVDz"]
482	>	for s in muInLabels:
483	>	muSources[s] = Source(label=(muModuleName, muProductPrefix+s, processName), datatype="vector<float>",  outLabel=s)
484	>
485	>	eleModuleName = "nTupleElectrons"
486	>	eleProductPrefix = "tightElectrons"
487	>	eleSources = {}
488	>	eleInLabels = ["Pt", "Eta", "Phi", "E", "PFRhoCorrectedRelIso", "AbsoluteDB"]
489	>	for s in eleInLabels:
490	>	eleSources[s] = Source(label=(eleModuleName, eleProductPrefix+s, processName), datatype="vector<float>",  outLabel=s)
491	>
492	>	muSources["AbsoluteDB"].outLabel = "dB"
493	>	muSources["PFDeltaCorrectedRelIso"].outLabel = "relIso"
494	>	eleSources["AbsoluteDB"].outLabel = "dB"
495	>	eleSources["PFRhoCorrectedRelIso"].outLabel = "relIso"
496	>
497	>	metModuleName = "nTuplePatMETsPF"
498	>	metProductPrefix = "patMETsPF"
499	>	metSources = {}
500	>	metInLabels = ["Pt", "Phi"]
501	>	for s in metInLabels:
502	>	metSources[s] = Source(label=(metModuleName, metProductPrefix+s, processName), datatype="vector<float>  ", outLabel=s)
503	>
504	>	sources = jetSources.values()+muSources.values()+eleSources.values()+metSources.values()
505	>
506	>	totalEvents = events.size()
507	>	print "Running over %d events" % events.size()
508	>	t1 = time.time()
509	>
510	>	triggers = Triggers()
511	>
512	>	outFile = ROOT.TFile(output_file, "RECREATE")
513	>	outTree = ROOT.TTree("events", "My Event Tree")
514	>
515	>	jetStructs = {}
516	>	maxN_Jets = 5
517	>	for i in range(maxN_Jets):
518	>	brName = "jet%d" % i
519	>	jetStructs[brName] = JetRepr.EvStruct()
520	>	outTree.Branch(brName, jetStructs[brName], jetStructs[brName].rowStr())
521	>
522	>	muStruct = MuonRepr.EvStruct()
523	>	outTree.Branch("mu", muStruct, muStruct.rowStr())
524	>
525	>	eleStruct = ElectronRepr.EvStruct()
526	>	outTree.Branch("ele", eleStruct, eleStruct.rowStr())
527	>
528	>	metBranch = METRepr.EvStruct()
529	>	outTree.Branch("met", metBranch, metBranch.rowStr())
530	>
531	>	eventBranch = EventStruct()
532	>	outTree.Branch("id", eventBranch, eventBranch.rowStr())
533	>
534	>	fstateBranch = FinalStateStruct()
535	>	outTree.Branch("fstate", fstateBranch, fstateBranch.rowStr())
536	>
537	>	signalLeptonBranch = LeptonRepr.EvStruct()
538	>	outTree.Branch("sigLepton", signalLeptonBranch, signalLeptonBranch.rowStr())
539	>
540	>	triggerBranch = TriggerStateStruct()
541	>	outTree.Branch("trigger", triggerBranch, triggerBranch.rowStr())
542	>
543	>	###------------###
544	>	### EVENT LOOP ###
545	>	###------------###
546	>	nProcessedEvents = 0
547	>	for event in events:
548	>
549	>	if not debugging:
550	>	if nProcessedEvents%1000==0: #Write a . every 1k lines
551	>	sys.stdout.write(".")
552	>	if nProcessedEvents%10000==0: #Write a newline every 10k lines
553	>	sys.stdout.write("\n%d " % nProcessedEvents)
554	>	sys.stdout.flush() #Make sure the dots and newlines get written to stdout
555	>
556	>
557	>	### VARIABLE INITIALIZATION
558	>	for (name, struct) in jetStructs.items():
559	>	struct.initVars()
560	>	muStruct.initVars()
561	>	eleStruct.initVars()
562	>	metBranch.initVars()
563	>	eventBranch.initVars()
564	>	fstateBranch.initVars()
565	>	signalLeptonBranch.initVars()
566	>	triggerBranch.initVars()
567	>
568	>	### EVENT->SOURCES
569	>	for source in sources:
570	>	source.get(event)
571	>	triggers.get(event)
572	>
573	>	### EVENTID
574	>	eventID = int(event.object().id().event())
575	>	runID = int(event.object().id().run())
576	>	lumiID = int(event.object().id().luminosityBlock())
577	>	eventBranch.eventID = eventID
578	>	eventBranch.runID = runID
579	>	eventBranch.lumiID = lumiID
580	>
581	>	if debugging and doSelEv:
582	>	if eventID not in selEvents:
583	>	continue
584	>
585	>	### TRIGGERS
586	>	triggerBranch.HLT_Mu = int(triggers.passesHLT("HLT_IsoMu24_eta2p1_v11"))
587	>	triggerBranch.HLT_Ele = int(True)
588	>
589	>	### SOURCE -> REPR
590	>	jets = ParticleRepr.make(jetSources, JetRepr)
591	>	muons = ParticleRepr.make(muSources, MuonRepr)
592	>	eles = ParticleRepr.make(eleSources, ElectronRepr)
593	>	met = ParticleRepr.make(metSources, METRepr)[0]
594	>
595	>	### JET SMEARING
596	>	for j in jets:
597	>	j.doMCSmearing()
598	>
599	>	### JET FILTERING
600	>	jets = filter(lambda x: x.pt_corr()>40, jets)
601	>	looseCSVBTags = filter(lambda x: x.bTagCSV() > Cuts.bTag.CSV.loose, jets)
602	>	mediumCSVBTags = filter(lambda x: x.bTagCSV() > Cuts.bTag.CSV.medium, jets)
603	>	tightCSVBTags = filter(lambda x: x.bTagCSV() > Cuts.bTag.CSV.tight, jets)
604	>	tightTCHPBTags = filter(lambda x: x.bTagTCHP() > Cuts.bTag.TCHP.tight, jets)
605	>
606	>	### Lepton filtering
607	>	muons = filter(lambda x: x.pt()>26 and x.relIso()<0.12, muons)
608	>	eles = filter(lambda x: x.pt()>30 and x.relIso()<0.20, eles)
609	>
610	>	### MET
611	>	met.calcValues()
612	>	met.fillStruct(metBranch)
613	>
614	>	###DEBUGGING PRINTOUT
615	>	if debugging:
616	>	print 80*"-"
617	>	print "eventID %d" % eventID
618	>	for j in jets:
619	>	print str(j)
620	>	for b in tightBTags:
621	>	print "b" + str(b)
622	>	print met
623	>
624	>	### FINAL STATE
625	>	fstateBranch.nJets = len(jets)
626	>	fstateBranch.nBTagsCSVL = len(looseCSVBTags)
627	>	fstateBranch.nBTagsCSVM = len(mediumCSVBTags)
628	>	fstateBranch.nBTagsCSVT = len(tightCSVBTags)
629	>	fstateBranch.nBTagsTCHPT = len(tightTCHPBTags)
630	>	fstateBranch.nMuons = len(muons)
631	>	fstateBranch.nEles = len(eles)
632	>	fstateBranch.sumJetPt = sum(map(lambda x: x.pt(), jets))
633	>
634	>	### JETS
635	>	i = 0
636	>	for j in jets[0:maxN_Jets]:
637	>	brName = "jet%d" % i
638	>	j.fillStruct(jetStructs[brName])
639	>	i += 1
640	>
641	>	### MUONS
642	>	if len(muons)>0:
643	>	muons[0].fillStruct(muStruct)
644	>
645	>	### ELECTRONS
646	>	if len(eles)>0:
647	>	eles[0].fillStruct(eleStruct)
648	>
649	>	### ISOLATED LEPTON
650	>	if len(muons)==1 and len(eles)==0: #We have one isolated muon
651	>	signalLepton = muons[0]
652	>	muons[0].fillStruct(signalLeptonBranch)
653	>	elif len(eles)==1 and len(muons)==0: #We have one isolated electron
654	>	signalLepton = eles[0]
655	>	else:
656	>	signalLepton = None
657	>
658	>	### RECONSTRUCT THE W
659	>	if signalLepton != None:
660	>	signalLepton.fillStruct(signalLeptonBranch)
661	>	mtwMass = met.calcW_MT(signalLepton)
662	>	fstateBranch.M_W = mtwMass
663	>
664	>	### RECONSTRUCT THE TOP
665	>	if len(tightTCHPBTags)>0:
666	>	sortedBTags = sorted(tightTCHPBTags, key=lambda x: x.bTagTCHP())
667	>	tightestBTagJet = sortedBTags[-1]
668	>	if debugging: print "jet 4 top: %s" % str(tightestBTagJet)
669	>	topMass = met.calc_topMass(signalLepton, tightestBTagJet)
670	>	fstateBranch.topMass = topMass
671	>
672	>	outTree.Fill()
673	>
674	>	nProcessedEvents += 1
675	>
676	>	sys.stdout.write("\n")
677	>	t2 = time.time()
678	>
679	>	elapsed = t2-t1
680	>	print "time elapsed: %.1f sec" % elapsed
681	>	evRate = int(float(totalEvents)/elapsed)
682	>	print "event processing rate: %d/s" % evRate
683
684	<	#Print out some debugging events
685	<	evs = [(x["eventId"], x["leptonEta"], x["leptonFlavour"], x["jetPt0"]) for x in table.where("(eventId >= 7802) & (eventId <= 7900) & (nTightLeptons==1)")]
684	>	#Write the tree to the file
685	>	outFile.Write()
686
687	<	for ev in evs:
688	<	print ev
687	>	#We're done, close the output file
688	>	outFile.Close()

Diff Legend

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