| 1 |
jpata |
1.1 |
import ROOT
|
| 2 |
|
|
from DataFormats.FWLite import Events, Handle
|
| 3 |
|
|
import sys
|
| 4 |
jpata |
1.2 |
import time
|
| 5 |
|
|
import numpy
|
| 6 |
jpata |
1.3 |
import ctypes
|
| 7 |
jpata |
1.1 |
|
| 8 |
jpata |
1.3 |
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 |
jpata |
1.1 |
|
| 23 |
jpata |
1.2 |
#Define the input file
|
| 24 |
jpata |
1.3 |
#input_file = "~/singletop/sync_T_t_ntuple.root"
|
| 25 |
|
|
#input_file = "~/singletop/edmntuple_T_t_sync.root"
|
| 26 |
jpata |
1.2 |
#input_file = "/hdfs/local/stpol/ntuples/WJets_part_1Merged.root"
|
| 27 |
jpata |
1.3 |
#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 |
jpata |
1.1 |
|
| 80 |
|
|
|
| 81 |
jpata |
1.3 |
def initVars(self):
|
| 82 |
|
|
for (name, dType) in self._fields_:
|
| 83 |
|
|
self.__setattr__(name, BaseStruct.getNaN(dType))
|
| 84 |
|
|
|
| 85 |
jpata |
1.2 |
class Source:
|
| 86 |
jpata |
1.3 |
def __init__(self, label, datatype, outLabel = None):
|
| 87 |
jpata |
1.2 |
self.datatype = datatype
|
| 88 |
|
|
self.handle = Handle(self.datatype)
|
| 89 |
|
|
self.label = label
|
| 90 |
|
|
self._data = None
|
| 91 |
jpata |
1.3 |
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 |
jpata |
1.2 |
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 |
jpata |
1.3 |
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 |
jpata |
1.2 |
def __init__(self):
|
| 181 |
|
|
super(ParticleRepr, self).__init__()
|
| 182 |
|
|
|
| 183 |
jpata |
1.3 |
|
| 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 |
jpata |
1.2 |
class METRepr(ParticleRepr):
|
| 203 |
jpata |
1.3 |
class EvStruct(BaseStruct):
|
| 204 |
|
|
_fields_ = [('pt', ctypes.c_float), ('phi', ctypes.c_float), ('px', ctypes.c_float), ('py', ctypes.c_float)]
|
| 205 |
|
|
|
| 206 |
jpata |
1.2 |
def __init__(self):
|
| 207 |
|
|
super(METRepr, self).__init__()
|
| 208 |
|
|
|
| 209 |
|
|
def calcValues(self):
|
| 210 |
jpata |
1.3 |
self.variables["Px"] = self.pt()*ROOT.TMath.Cos(self.phi())
|
| 211 |
|
|
self.variables["Py"] = self.pt()*ROOT.TMath.Sin(self.phi())
|
| 212 |
jpata |
1.2 |
|
| 213 |
|
|
def calcW_MT(self, lepton):
|
| 214 |
jpata |
1.3 |
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 |
jpata |
1.2 |
|
| 284 |
|
|
class JetRepr(ParticleRepr):
|
| 285 |
jpata |
1.3 |
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 |
jpata |
1.2 |
|
| 295 |
|
|
def __init__(self):
|
| 296 |
|
|
super(JetRepr, self).__init__()
|
| 297 |
|
|
|
| 298 |
jpata |
1.3 |
def passesBTagCut(self, cut):
|
| 299 |
|
|
return self.bTag() > cut
|
| 300 |
jpata |
1.2 |
|
| 301 |
|
|
def isTightJet(self):
|
| 302 |
|
|
return self.variables["Pt"] > jetTightPtCut
|
| 303 |
|
|
|
| 304 |
jpata |
1.3 |
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 |
jpata |
1.2 |
class LeptonRepr(ParticleRepr):
|
| 371 |
jpata |
1.3 |
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 |
jpata |
1.2 |
|
| 378 |
|
|
def __init__(self):
|
| 379 |
|
|
super(LeptonRepr, self).__init__()
|
| 380 |
|
|
|
| 381 |
jpata |
1.3 |
def relIso(self):
|
| 382 |
|
|
return self.variables["relIso"]
|
| 383 |
|
|
|
| 384 |
jpata |
1.2 |
class MuonRepr(LeptonRepr):
|
| 385 |
|
|
|
| 386 |
|
|
def __init__(self):
|
| 387 |
|
|
super(MuonRepr, self).__init__()
|
| 388 |
jpata |
1.1 |
|
| 389 |
jpata |
1.2 |
def isTight(self):
|
| 390 |
|
|
pass
|
| 391 |
jpata |
1.1 |
|
| 392 |
jpata |
1.2 |
class ElectronRepr(LeptonRepr):
|
| 393 |
|
|
varOrder = LeptonRepr.varOrder
|
| 394 |
jpata |
1.1 |
|
| 395 |
jpata |
1.2 |
def __init__(self):
|
| 396 |
|
|
super(ElectronRepr, self).__init__()
|
| 397 |
jpata |
1.1 |
|
| 398 |
jpata |
1.2 |
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 |
jpata |
1.3 |
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 |
jpata |
1.2 |
|
| 658 |
jpata |
1.3 |
### 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 |
jpata |
1.1 |
|
| 684 |
jpata |
1.3 |
#Write the tree to the file
|
| 685 |
|
|
outFile.Write()
|
| 686 |
jpata |
1.2 |
|
| 687 |
jpata |
1.3 |
#We're done, close the output file
|
| 688 |
|
|
outFile.Close() |
