ViewVC Help
View File | Revision Log | Show Annotations | Root Listing
root/cvsroot/UserCode/STPol/sync/syncTopRefSel.py
Revision: 1.3
Committed: Thu Jul 26 19:32:45 2012 UTC (12 years, 9 months ago) by jpata
Content type: text/x-python
Branch: MAIN
CVS Tags: HEAD
Changes since 1.2: +533 -178 lines
Error occurred while calculating annotation data.
Log Message:
changed skimming

File Contents

# Content
1 import ROOT
2 from DataFormats.FWLite import Events, Handle
3 import sys
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 #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 #Write the tree to the file
685 outFile.Write()
686
687 #We're done, close the output file
688 outFile.Close()