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() |