def convertGENSIM(infiles, outfilename, Nmax=-1):
"""Loop over GENSIM events and save custom trees."""
start1 = time.time()
lqids = [9000008, 9000009]
xids = [9000006, 9000007]
dmids = [9000005]
print ">>> loading files..."
events = Events(infiles)
outfile = TFile(outfilename, 'RECREATE')
print ">>> creating trees and branches..."
tree_event = TTree('event', 'event')
tree_jet = TTree('jet', 'jet')
tree_jet_raw = TTree('jet_raw', 'jet_raw')
# EVENT
tree_event.addBranch('nbgen', 'i')
tree_event.addBranch('nbcut', 'i')
tree_event.addBranch('nbcut50', 'i')
tree_event.addBranch('njet', 'i')
tree_event.addBranch('ntau', 'i')
tree_event.addBranch('ntaucut', 'i')
tree_event.addBranch('ntaucut50', 'i')
tree_event.addBranch('ntaucut_vis', 'i')
tree_event.addBranch('ntaucut50_vis', 'i')
tree_event.addBranch('nlq', 'i')
tree_event.addBranch('ndm', 'i')
tree_event.addBranch('nx', 'i')
tree_event.addBranch('met_pt', 'f')
tree_event.addBranch('met_phi', 'f')
tree_event.addBranch('jet1_mass', 'f')
tree_event.addBranch('jet2_mass', 'f')
tree_event.addBranch('jet3_mass', 'f')
tree_event.addBranch('jet1_pt', 'f')
tree_event.addBranch('jet2_pt', 'f')
tree_event.addBranch('jet3_pt', 'f')
tree_event.addBranch('jet1_eta', 'f')
tree_event.addBranch('jet2_eta', 'f')
tree_event.addBranch('jet3_eta', 'f')
tree_event.addBranch('jet1_phi', 'f')
tree_event.addBranch('jet2_phi', 'f')
tree_event.addBranch('jet3_phi', 'f')
tree_event.addBranch('sumjet', 'f')
tree_event.addBranch('dphi_jj', 'f')
tree_event.addBranch('deta_jj', 'f')
tree_event.addBranch('dr_jj', 'f')
tree_event.addBranch('mjj', 'f')
tree_event.addBranch('lq1_mass', 'f')
tree_event.addBranch('lq2_mass', 'f')
tree_event.addBranch('lq1_pt', 'f')
tree_event.addBranch('lq2_pt', 'f')
tree_event.addBranch('lq1_eta', 'f')
tree_event.addBranch('lq2_eta', 'f')
tree_event.addBranch('lq1_phi', 'f')
tree_event.addBranch('lq2_phi', 'f')
tree_event.addBranch('dm1_mass', 'f')
tree_event.addBranch('dm2_mass', 'f')
tree_event.addBranch('dm1_pt', 'f')
tree_event.addBranch('dm2_pt', 'f')
tree_event.addBranch('dm1_eta', 'f')
tree_event.addBranch('dm2_eta', 'f')
tree_event.addBranch('dm1_phi', 'f')
tree_event.addBranch('dm2_phi', 'f')
tree_event.addBranch('b1_mass', 'f')
tree_event.addBranch('b2_mass', 'f')
tree_event.addBranch('b1_pt', 'f')
tree_event.addBranch('b2_pt', 'f')
tree_event.addBranch('b1_eta', 'f')
tree_event.addBranch('b2_eta', 'f')
tree_event.addBranch('b1_phi', 'f')
tree_event.addBranch('b2_phi', 'f')
tree_event.addBranch('x_mass', 'f')
tree_event.addBranch('x_pt', 'f')
tree_event.addBranch('x_eta', 'f')
tree_event.addBranch('x_phi', 'f')
tree_event.addBranch('m_lqlq', 'f')
tree_event.addBranch('tau1_mass', 'f')
tree_event.addBranch('tau2_mass', 'f')
tree_event.addBranch('tau1_pt', 'f')
tree_event.addBranch('tau2_pt', 'f')
tree_event.addBranch('tau1_eta', 'f')
tree_event.addBranch('tau2_eta', 'f')
tree_event.addBranch('tau1_phi', 'f')
tree_event.addBranch('tau2_phi', 'f')
tree_event.addBranch('tau1_y', 'f')
tree_event.addBranch('tau2_y', 'f')
tree_event.addBranch('tau1_massvis', 'f')
tree_event.addBranch('tau2_massvis', 'f')
tree_event.addBranch('tau1_ptvis', 'f')
tree_event.addBranch('tau2_ptvis', 'f')
tree_event.addBranch('tau1_etavis', 'f')
tree_event.addBranch('tau2_etavis', 'f')
tree_event.addBranch('tau1_phivis', 'f')
tree_event.addBranch('tau2_phivis', 'f')
tree_event.addBranch('ditau_dy', 'f')
tree_event.addBranch('ditau_chi', 'f')
tree_event.addBranch('m_tauvis1_jet1', 'f')
tree_event.addBranch('st', 'f') # scalar sum pT
tree_event.addBranch('st_met', 'f') # scalar sum pT with MET
# tree_event.addBranch('jets_pt' , 'f', standard=False)
# tree_event.addBranch('jets_eta' , 'f', standard=False)
# tree_event.addBranch('jets_phi' , 'f', standard=False)
# tree_event.addBranch('jets_mass', 'f', standard=False)
tree_event.addBranch('weight', 'f')
tree_event.addBranch('event_id', 'i')
# JETS raw
tree_jet_raw.addBranch('mass', 'f')
tree_jet_raw.addBranch('pt', 'f')
tree_jet_raw.addBranch('eta', 'f')
tree_jet_raw.addBranch('phi', 'f')
tree_jet_raw.addBranch('weight', 'f')
tree_jet_raw.addBranch('event_id', 'i')
# JETS cleaned
tree_jet.addBranch('mass', 'f')
tree_jet.addBranch('pt', 'f')
tree_jet.addBranch('eta', 'f')
tree_jet.addBranch('phi', 'f')
tree_jet.addBranch('weight', 'f')
tree_jet.addBranch('event_id', 'i')
handle_gps, label_gps = Handle(
'std::vector<reco::GenParticle>'), 'genParticles'
handle_jets, label_jets = Handle('std::vector<reco::GenJet>'), 'ak4GenJets'
handle_met, label_met = Handle('vector<reco::GenMET>'), 'genMetTrue'
handle_weight, label_weight = Handle('GenEventInfoProduct'), 'generator'
evtid = 0
sec_per_evt = 0.023 # seconds per event
Ntot = Nmax if Nmax > 0 else events.size()
print ">>> start processing %d events, ETA %s..." % (
Ntot, formatTimeShort(sec_per_evt * Ntot))
step = stepsize(Ntot)
start_proc = time.time()
# LOOP OVER EVENTS
for event in events:
# print ' --- NEW EVENT'
# print '='*30
# print evtid
if Nmax > 0 and evtid >= Nmax: break
if evtid > 0 and evtid % step == 0:
print ">>> processed %4s/%d events, ETA %s" % (
evtid, Ntot, ETA(start_proc, evtid + 1, Ntot))
evtid += 1
tree_event.event_id[0] = evtid
event.getByLabel(label_gps, handle_gps)
gps = handle_gps.product()
event.getByLabel(label_jets, handle_jets)
jets = handle_jets.product()
event.getByLabel(label_met, handle_met)
met = handle_met.product()
event.getByLabel(label_weight, handle_weight)
gweight = handle_weight.product()
weight = gweight.weight()
# GEN PARTICLES
gps_final = [
p for p in gps if isFinal(p) and abs(p.pdgId()) in [5, 6, 15, 16] +
lqids + dmids + xids
]
gps_mother = [
p for p in gps_final if abs(p.pdgId()) in lqids and p.status() > 60
]
gps_mother.sort(key=lambda p: p.pt(), reverse=True)
gps_dm = [p for p in gps_final if abs(p.pdgId()) in dmids]
gps_dm.sort(key=lambda p: p.pt(), reverse=True)
gps_x = [p for p in gps_final if abs(p.pdgId()) in xids]
gps_x.sort(key=lambda p: p.pt(), reverse=True)
gps_bgen = [
p for p in gps_final if abs(p.pdgId()) == 5 and p.status() == 71
]
gps_bgen.sort(key=lambda p: p.pt(), reverse=True)
gps_bfromNP = [
p for p in gps
if abs(p.pdgId()) == 5 and abs(p.mother(0).pdgId()) in lqids +
dmids + xids
]
gps_bfromNP.sort(key=lambda p: p.pt(), reverse=True)
gps_bcut = [p for p in gps_bgen if p.pt() > 20 and abs(p.eta()) < 2.5]
gps_bcut50 = [
p for p in gps_bgen if p.pt() > 50 and abs(p.eta()) < 2.5
]
gps_tau = [
p for p in gps_final if abs(p.pdgId()) == 15 and p.status() == 2
]
gps_tau.sort(key=lambda p: p.pt(), reverse=True)
gps_taufromLQ = [
p for p in gps
if abs(p.pdgId()) == 15 and abs(p.mother(0).pdgId()) in lqids
]
gps_taufromLQ.sort(key=lambda p: p.pt(), reverse=True)
gps_taucut = [p for p in gps_tau if p.pt() > 20 and abs(p.eta()) < 2.5]
gps_taucut50 = [
p for p in gps_tau if p.pt() > 50 and abs(p.eta()) < 2.5
]
gps_tau_vis = []
gps_taucut_vis = []
gps_taucut50_vis = []
# find taus that survive ptvis > 20 and >50 cuts
for p in gps_tau:
while p.status() != 2:
p = p.daughter(0)
findau = finalDaughters(p, [])
tauvisp4 = p4sumvis(findau)
thisptvis = tauvisp4.pt()
thisetavis = tauvisp4.eta()
gps_tau_vis.append(tauvisp4)
if thisptvis > 20 and abs(thisetavis) < 2.5:
gps_taucut_vis.append(tauvisp4)
if thisptvis > 50 and abs(thisetavis) < 2.5:
gps_taucut50_vis.append(tauvisp4)
gps_tau_vis.sort(key=lambda p: p.pt(), reverse=True)
gps_taucut_vis.sort(key=lambda p: p.pt(), reverse=True)
gps_taucut50_vis.sort(key=lambda p: p.pt(), reverse=True)
# if not len(gps_bgen) == 2:
# print 'len, gps_mother = ', len(gps_mother)
# print 'number of b\'s: %i' % len(gps_bgen)
# print 'number of b\'s from NP decays: %i' % len(gps_bfromNP)
# if len(gps_mother) > 0:
# print '=' * 15 + ' NEW EVENT' + '='*15
# printDecayChain(gps_mother)
# continue
# # REMOVE TOP QUARK if its final daughter is also in the list
# for top in gps_tgen[:]:
# dau = top
# while abs(dau.daughter(0).pdgId())==6:
# dau = dau.daughter(0)
# if dau!=top and dau in gps_tgen:
# gps_tgen.remove(top)
# # write raw jets to tree
# jets_pt = []
# jets_eta = []
# jets_phi = []
# jets_mass = []
# for jet in jets:
# jets_pt.append(jet.pt())
# jets_eta.append(jet.eta())
# jets_phi.append(jet.phi())
# jets_mass.append(jet.mass())
# tree_event.jets_pt[0] = jets_pt
# tree_event.jets_pt[0] = jets_pt
# tree_event.jets_eta[0] = jets_eta
# tree_event.jets_phi[0] = jets_phi
# tree_event.jets_mass[0] = jets_mass
# write raw jets to a tree
for jet in jets:
tree_jet_raw.mass[0] = jet.mass()
tree_jet_raw.pt[0] = jet.pt()
tree_jet_raw.eta[0] = jet.eta()
tree_jet_raw.phi[0] = jet.phi()
tree_jet_raw.weight[0] = weight
tree_jet_raw.event_id[0] = evtid
tree_jet_raw.Fill()
# REMOVE JET-LEPTON OVERLAP
jets, dummy = cleanObjectCollection(jets, gps_tau_vis, dRmin=0.4)
# print 'now cleaning from DM, now: %f' % (len(jets))
jets, dummy = cleanObjectCollection(jets, gps_dm, dRmin=0.4, mfrac=-1.)
# print 'after cleaning from DM, now: %f' % (len(jets))
jets.sort(key=lambda p: p.pt(), reverse=True)
njets = 0
sumjet = 0
jets30 = []
for jet in jets:
if jet.pt() > 30 and abs(jet.eta()) < 5:
sumjet += jet.pt()
njets += 1
tree_jet.mass[0] = jet.mass()
tree_jet.pt[0] = jet.pt()
tree_jet.eta[0] = jet.eta()
tree_jet.phi[0] = jet.phi()
tree_jet.weight[0] = weight
tree_jet.event_id[0] = evtid
tree_jet.Fill()
jets30.append(jet)
# MULTIPLICITIES
tree_event.nlq[0] = len(gps_mother)
tree_event.ndm[0] = len(gps_dm)
tree_event.nx[0] = len(gps_x)
tree_event.nbcut[0] = len(gps_bcut)
tree_event.nbcut50[0] = len(gps_bcut50)
tree_event.nbgen[0] = len(gps_bgen)
tree_event.njet[0] = njets
tree_event.ntau[0] = len(gps_tau)
tree_event.ntaucut[0] = len(gps_taucut)
tree_event.ntaucut50[0] = len(gps_taucut50)
tree_event.ntaucut_vis[0] = len(gps_taucut_vis)
tree_event.ntaucut50_vis[0] = len(gps_taucut50_vis)
# JETS
tree_event.met_pt[0] = met[0].pt()
tree_event.met_phi[0] = met[0].phi()
tree_event.sumjet[0] = sumjet
if len(jets30) >= 2:
tree_event.jet1_mass[0] = jets30[0].mass()
tree_event.jet2_mass[0] = jets30[1].mass()
tree_event.jet1_pt[0] = jets30[0].pt()
tree_event.jet2_pt[0] = jets30[1].pt()
tree_event.jet1_eta[0] = jets30[0].eta()
tree_event.jet2_eta[0] = jets30[1].eta()
tree_event.jet1_phi[0] = jets30[0].phi()
tree_event.jet2_phi[0] = jets30[1].phi()
if len(jets30) >= 3:
tree_event.jet3_mass[0] = jets30[2].mass()
tree_event.jet3_pt[0] = jets30[2].pt()
tree_event.jet3_eta[0] = jets30[2].eta()
tree_event.jet3_phi[0] = jets30[2].phi()
else:
tree_event.jet3_mass[0] = -1.
tree_event.jet3_pt[0] = -1.
tree_event.jet3_eta[0] = -9.
tree_event.jet3_phi[0] = -9.
tree_event.dphi_jj[0] = deltaPhi(jets30[0].phi(), jets30[1].phi())
tree_event.deta_jj[0] = jets30[0].eta() - jets30[1].eta()
tree_event.dr_jj[0] = deltaR(jets30[0].eta(), jets30[0].phi(),
jets30[1].eta(), jets30[1].phi())
dijetp4 = jets30[0].p4() + jets30[1].p4()
tree_event.mjj[0] = dijetp4.M()
elif len(jets30) == 1:
tree_event.jet1_mass[0] = jets30[0].mass()
tree_event.jet2_mass[0] = -1.
tree_event.jet3_mass[0] = -1.
tree_event.jet1_pt[0] = jets30[0].pt()
tree_event.jet2_pt[0] = -1.
tree_event.jet3_pt[0] = -1.
tree_event.jet1_eta[0] = jets30[0].eta()
tree_event.jet2_eta[0] = -9.
tree_event.jet3_eta[0] = -9.
tree_event.jet1_phi[0] = jets30[0].phi()
tree_event.jet2_phi[0] = -9.
tree_event.jet3_phi[0] = -9.
tree_event.dphi_jj[0] = -9.
tree_event.deta_jj[0] = -9.
tree_event.dr_jj[0] = -1.
tree_event.mjj[0] = -1.
else:
tree_event.jet1_mass[0] = -1.
tree_event.jet2_mass[0] = -1.
tree_event.jet3_mass[0] = -1.
tree_event.jet1_pt[0] = -1.
tree_event.jet2_pt[0] = -1.
tree_event.jet3_pt[0] = -1.
tree_event.jet1_eta[0] = -9.
tree_event.jet2_eta[0] = -9.
tree_event.jet3_eta[0] = -9.
tree_event.jet1_phi[0] = -9.
tree_event.jet2_phi[0] = -9.
tree_event.jet3_phi[0] = -9.
tree_event.dphi_jj[0] = -9.
tree_event.deta_jj[0] = -9.
tree_event.dr_jj[0] = -1.
tree_event.mjj[0] = -1.
if len(gps_bgen) >= 2:
tree_event.b1_mass[0] = gps_bgen[0].mass()
tree_event.b1_pt[0] = gps_bgen[0].pt()
tree_event.b1_eta[0] = gps_bgen[0].eta()
tree_event.b1_phi[0] = gps_bgen[0].phi()
tree_event.b2_mass[0] = gps_bgen[1].mass()
tree_event.b2_pt[0] = gps_bgen[1].pt()
tree_event.b2_eta[0] = gps_bgen[1].eta()
tree_event.b2_phi[0] = gps_bgen[1].phi()
elif len(gps_bgen) == 1:
tree_event.b1_mass[0] = gps_bgen[0].mass()
tree_event.b1_pt[0] = gps_bgen[0].pt()
tree_event.b1_eta[0] = gps_bgen[0].eta()
tree_event.b1_phi[0] = gps_bgen[0].phi()
tree_event.b2_mass[0] = -1.
tree_event.b2_pt[0] = -1.
tree_event.b2_eta[0] = -9.
tree_event.b2_phi[0] = -9.
else:
tree_event.b1_mass[0] = -1.
tree_event.b1_pt[0] = -1.
tree_event.b1_eta[0] = -9.
tree_event.b1_phi[0] = -9.
tree_event.b2_mass[0] = -1.
tree_event.b2_pt[0] = -1.
tree_event.b2_eta[0] = -9.
tree_event.b2_phi[0] = -9.
# SCALAR SUM PT
if len(gps_tau_vis) >= 1 and len(gps_bcut) >= 1:
st = 0
for tauvis in gps_tau_vis[:2]:
st += tauvis.pt()
for b in gps_bcut:
st += b.pt()
stmet = st + met[0].pt()
else:
st = -1
stmet = -1
if len(gps_tau) >= 2:
tree_event.tau1_mass[0] = gps_tau[0].mass()
tree_event.tau2_mass[0] = gps_tau[1].mass()
tree_event.tau1_pt[0] = gps_tau[0].pt()
tree_event.tau2_pt[0] = gps_tau[1].pt()
tree_event.tau1_eta[0] = gps_tau[0].eta()
tree_event.tau2_eta[0] = gps_tau[1].eta()
tree_event.tau1_phi[0] = gps_tau[0].phi()
tree_event.tau2_phi[0] = gps_tau[1].phi()
tree_event.tau1_y[0] = gps_tau[0].p4().Rapidity()
tree_event.tau2_y[0] = gps_tau[1].p4().Rapidity()
tree_event.tau1_massvis[0] = gps_tau_vis[0].mass()
tree_event.tau2_massvis[0] = gps_tau_vis[1].mass()
tree_event.tau1_ptvis[0] = gps_tau_vis[0].pt()
tree_event.tau2_ptvis[0] = gps_tau_vis[1].pt()
tree_event.tau1_etavis[0] = gps_tau_vis[0].eta()
tree_event.tau2_etavis[0] = gps_tau_vis[1].eta()
tree_event.tau1_phivis[0] = gps_tau_vis[0].phi()
tree_event.tau2_phivis[0] = gps_tau_vis[1].phi()
dy = abs(gps_tau[0].p4().Rapidity() - gps_tau[1].p4().Rapidity())
tree_event.ditau_dy[0] = dy
tree_event.ditau_chi[0] = exp(dy)
elif len(gps_tau) == 1:
tree_event.tau1_mass[0] = gps_tau[0].mass()
tree_event.tau2_mass[0] = -1.
tree_event.tau1_pt[0] = gps_tau[0].pt()
tree_event.tau2_pt[0] = -1.
tree_event.tau1_eta[0] = gps_tau[0].eta()
tree_event.tau2_eta[0] = -9.
tree_event.tau1_phi[0] = gps_tau[0].phi()
tree_event.tau2_phi[0] = -9.
tree_event.tau1_y[0] = gps_tau[0].p4().Rapidity()
tree_event.tau2_y[0] = -9.
tree_event.tau1_massvis[0] = gps_tau_vis[0].mass()
tree_event.tau2_massvis[0] = -1.
tree_event.tau1_ptvis[0] = gps_tau_vis[0].pt()
tree_event.tau2_ptvis[0] = -1.
tree_event.tau1_etavis[0] = gps_tau_vis[0].eta()
tree_event.tau2_etavis[0] = -9.
tree_event.tau1_phivis[0] = gps_tau_vis[0].phi()
tree_event.tau2_phivis[0] = -9.
tree_event.ditau_dy[0] = -9.
tree_event.ditau_chi[0] = -1.
else:
tree_event.tau1_mass[0] = -1.
tree_event.tau2_mass[0] = -1.
tree_event.tau1_pt[0] = -1.
tree_event.tau2_pt[0] = -1.
tree_event.tau1_eta[0] = -9.
tree_event.tau2_eta[0] = -9.
tree_event.tau1_phi[0] = -9.
tree_event.tau2_phi[0] = -9.
tree_event.tau1_y[0] = -9.
tree_event.tau2_y[0] = -9.
tree_event.tau1_massvis[0] = -1.
tree_event.tau2_massvis[0] = -1.
tree_event.tau1_ptvis[0] = -1.
tree_event.tau2_ptvis[0] = -1.
tree_event.tau1_etavis[0] = -9.
tree_event.tau2_etavis[0] = -9.
tree_event.tau1_phivis[0] = -9.
tree_event.tau2_phivis[0] = -9.
tree_event.ditau_dy[0] = -9.
tree_event.ditau_chi[0] = -1.
tree_event.st[0] = st
tree_event.st_met[0] = stmet
# M (tauvis1, jet1)
if len(gps_tau) >= 1 and len(jets30) >= 1:
tree_event.m_tauvis1_jet1[0] = (gps_tau[0].p4() +
jets30[0].p4()).M()
else:
tree_event.m_tauvis1_jet1[0] = -1.
tree_event.weight[0] = weight
if len(gps_mother) == 1:
tree_event.lq1_mass[0] = gps_mother[0].mass()
tree_event.lq2_mass[0] = -1.
tree_event.lq1_pt[0] = gps_mother[0].pt()
tree_event.lq2_pt[0] = -1.
tree_event.lq1_eta[0] = gps_mother[0].eta()
tree_event.lq2_eta[0] = -9.
tree_event.lq1_phi[0] = gps_mother[0].phi()
tree_event.lq2_phi[0] = -9.
tree_event.m_lqlq[0] = -1.
elif len(gps_mother) >= 2:
tree_event.lq1_mass[0] = gps_mother[0].mass()
tree_event.lq2_mass[0] = gps_mother[1].mass()
tree_event.lq1_pt[0] = gps_mother[0].pt()
tree_event.lq2_pt[0] = gps_mother[1].pt()
tree_event.lq1_eta[0] = gps_mother[0].eta()
tree_event.lq2_eta[0] = gps_mother[1].eta()
tree_event.lq1_phi[0] = gps_mother[0].phi()
tree_event.lq2_phi[0] = gps_mother[1].phi()
dilqp4 = gps_mother[0].p4() + gps_mother[1].p4()
tree_event.m_lqlq[0] = dilqp4.M()
else:
tree_event.lq1_mass[0] = -1.
tree_event.lq2_mass[0] = -1.
tree_event.lq1_pt[0] = -1.
tree_event.lq2_pt[0] = -1.
tree_event.lq1_eta[0] = -9.
tree_event.lq2_eta[0] = -9.
tree_event.lq1_phi[0] = -9.
tree_event.lq2_phi[0] = -9.
tree_event.m_lqlq[0] = -1.
if len(gps_dm) == 1:
tree_event.dm1_mass[0] = gps_dm[0].mass()
tree_event.dm2_mass[0] = -1.
tree_event.dm1_pt[0] = gps_dm[0].pt()
tree_event.dm2_pt[0] = -1.
tree_event.dm1_eta[0] = gps_dm[0].eta()
tree_event.dm2_eta[0] = -9.
tree_event.dm1_phi[0] = gps_dm[0].phi()
tree_event.dm2_phi[0] = -9.
elif len(gps_dm) >= 2:
tree_event.dm1_mass[0] = gps_dm[0].mass()
tree_event.dm2_mass[0] = gps_dm[1].mass()
tree_event.dm1_pt[0] = gps_dm[0].pt()
tree_event.dm2_pt[0] = gps_dm[1].pt()
tree_event.dm1_eta[0] = gps_dm[0].eta()
tree_event.dm2_eta[0] = gps_dm[1].eta()
tree_event.dm1_phi[0] = gps_dm[0].phi()
tree_event.dm2_phi[0] = gps_dm[1].phi()
else:
tree_event.dm1_mass[0] = -1.
tree_event.dm2_mass[0] = -1.
tree_event.dm1_pt[0] = -1.
tree_event.dm2_pt[0] = -1.
tree_event.dm1_eta[0] = -9.
tree_event.dm2_eta[0] = -9.
tree_event.dm1_phi[0] = -9.
tree_event.dm2_phi[0] = -9.
if len(gps_x) == 1:
tree_event.x_mass[0] = gps_x[0].mass()
tree_event.x_pt[0] = gps_x[0].pt()
tree_event.x_eta[0] = gps_x[0].eta()
tree_event.x_phi[0] = gps_x[0].phi()
else:
tree_event.x_mass[0] = -1.
tree_event.x_pt[0] = -1.
tree_event.x_eta[0] = -9.
tree_event.x_phi[0] = -9.
tree_event.Fill()
print ">>> processed %4s events in %s" % (
evtid, formatTime(time.time() - start_proc))
print ">>> writing to output file %s..." % (outfilename)
outfile.Write()
outfile.Close()
print ">>> done in in %s" % (formatTime(time.time() - start1))
def convertGENSIM(infiles, outfilename, Nmax=-1, isPythia=False):
"""Loop over GENSIM events and save custom trees."""
start1 = time.time()
lqids = [46] if isPythia else [9000002, 9000006]
print ">>> loading files..."
events = Events(infiles)
outfile = TFile(outfilename, 'RECREATE')
print ">>> creating trees and branches..."
tree_event = TTree('event', 'event')
tree_jet = TTree('jet', 'jet')
tree_mother = TTree('mother', 'mother')
tree_decay = TTree('decay', 'decay')
tree_assoc = TTree('assoc', 'assoc')
# EVENT
tree_event.addBranch('nbgen', 'i')
tree_event.addBranch('nbcut', 'i')
tree_event.addBranch('nbcut50', 'i')
tree_event.addBranch('ntgen', 'i')
tree_event.addBranch('njet', 'i')
tree_event.addBranch('nlepton', 'i')
tree_event.addBranch('ntau', 'i')
tree_event.addBranch('ntaucut', 'i')
tree_event.addBranch('ntaucut50', 'i')
tree_event.addBranch('ntaucut_vis', 'i')
tree_event.addBranch('ntaucut50_vis', 'i')
tree_event.addBranch('nnu', 'i')
tree_event.addBranch('nlq', 'i')
tree_event.addBranch('ntau_assoc', 'i')
tree_event.addBranch('ntau_decay', 'i')
tree_event.addBranch('nbgen_decay', 'i')
tree_event.addBranch('met', 'f')
tree_event.addBranch('jpt1', 'f')
tree_event.addBranch('jpt2', 'f')
tree_event.addBranch('sumjet', 'f')
tree_event.addBranch('dphi_jj', 'f')
tree_event.addBranch('deta_jj', 'f')
tree_event.addBranch('dr_jj', 'f')
tree_event.addBranch('ncentral', 'i')
tree_event.addBranch('mjj', 'f')
tree_event.addBranch('lq1_mass', 'f')
tree_event.addBranch('lq2_mass', 'f')
tree_event.addBranch('lq1_pt', 'f')
tree_event.addBranch('lq2_pt', 'f')
tree_event.addBranch('m_lqlq', 'f')
tree_event.addBranch('tau1_pt', 'f')
tree_event.addBranch('tau1_eta', 'f')
tree_event.addBranch('tau1_y', 'f')
tree_event.addBranch('tau1_ptvis', 'f')
tree_event.addBranch('tau1_etavis', 'f')
tree_event.addBranch('tau2_pt', 'f')
tree_event.addBranch('tau2_eta', 'f')
tree_event.addBranch('tau2_y', 'f')
tree_event.addBranch('tau2_ptvis', 'f')
tree_event.addBranch('tau2_etavis', 'f')
tree_event.addBranch('ditau_dy', 'f')
tree_event.addBranch('ditau_chi', 'f')
tree_event.addBranch('st', 'f') # scalar sum pT
tree_event.addBranch('st_met', 'f') # scalar sum pT with MET
tree_event.addBranch('weight', 'f')
# LQ DECAY
tree_mother.addBranch('pid', 'i')
tree_mother.addBranch('moth', 'i')
tree_mother.addBranch('status', 'i')
tree_mother.addBranch('pt', 'f')
tree_mother.addBranch('eta', 'f')
tree_mother.addBranch('phi', 'f')
tree_mother.addBranch('mass', 'f')
tree_mother.addBranch('inv', 'f')
tree_mother.addBranch('ndau', 'i')
tree_mother.addBranch('dau', 'i')
tree_mother.addBranch('dphi_ll', 'f')
tree_mother.addBranch('deta_ll', 'f')
tree_mother.addBranch('dr_ll', 'f')
tree_mother.addBranch('st', 'f') # scalar sum pT
tree_mother.addBranch('st_met', 'f') # scalar sum pT with MET
tree_mother.addBranch('weight', 'f')
# FROM LQ DECAY
tree_decay.addBranch('pid', 'i')
tree_decay.addBranch('pt', 'f')
tree_decay.addBranch('eta', 'f')
tree_decay.addBranch('phi', 'f')
tree_decay.addBranch('lq_mass', 'f')
tree_decay.addBranch('ptvis', 'f')
tree_decay.addBranch('type', 'i')
tree_decay.addBranch('isBrem', 'i')
tree_decay.addBranch('weight', 'f')
# NOT FROM LQ DECAY (ASSOCIATED)
tree_assoc.addBranch('pid', 'i')
tree_assoc.addBranch('moth', 'i')
tree_assoc.addBranch('pt', 'f')
tree_assoc.addBranch('ptvis', 'f')
tree_assoc.addBranch('eta', 'f')
tree_assoc.addBranch('phi', 'f')
tree_assoc.addBranch('weight', 'f')
# JETS
tree_jet.addBranch('pt', 'f')
tree_jet.addBranch('eta', 'f')
tree_jet.addBranch('phi', 'f')
tree_jet.addBranch('weight', 'f')
hist_LQ_decay = TH1F('LQ_decay', "LQ decay", 60, -30, 30)
handle_gps, label_gps = Handle(
'std::vector<reco::GenParticle>'), 'genParticles'
handle_jets, label_jets = Handle('std::vector<reco::GenJet>'), 'ak4GenJets'
handle_met, label_met = Handle('vector<reco::GenMET>'), 'genMetTrue'
handle_weight, label_weight = Handle('GenEventInfoProduct'), 'generator'
evtid = 0
sec_per_evt = 0.023 # seconds per event
Ntot = Nmax if Nmax > 0 else events.size()
print ">>> start processing %d events, ETA %s..." % (
Ntot, formatTimeShort(sec_per_evt * Ntot))
step = stepsize(Ntot)
start_proc = time.time()
# LOOP OVER EVENTS
for event in events:
# print ' --- NEW EVENT'
# print '='*30
# print evtid
if Nmax > 0 and evtid >= Nmax: break
if evtid > 0 and evtid % step == 0:
print ">>> processed %4s/%d events, ETA %s" % (
evtid, Ntot, ETA(start_proc, evtid + 1, Ntot))
evtid += 1
event.getByLabel(label_gps, handle_gps)
gps = handle_gps.product()
event.getByLabel(label_jets, handle_jets)
jets = handle_jets.product()
event.getByLabel(label_met, handle_met)
met = handle_met.product()
event.getByLabel(label_weight, handle_weight)
gweight = handle_weight.product()
weight = gweight.weight()
# GEN PARTICLES
gps_mother = [p for p in gps if isFinal(p) and abs(p.pdgId()) in [42]]
gps_final = [
p for p in gps
if isFinal(p) and abs(p.pdgId()) in [5, 6, 15, 16] + lqids
]
gps_mother = [
p for p in gps_final if abs(p.pdgId()) in lqids and p.status() > 60
] #not(moth.numberOfDaughters()==2 and abs(moth.daughter(0).pdgId()) in lqids)
gps_bgen = [
p for p in gps_final if abs(p.pdgId()) == 5 and p.status() == 71
]
gps_bcut = [p for p in gps_bgen if p.pt() > 20 and abs(p.eta()) < 2.5]
gps_bcut50 = [
p for p in gps_bgen if p.pt() > 50 and abs(p.eta()) < 2.5
]
gps_tgen = [p for p in gps_final if abs(p.pdgId()) == 6] #[-1:]
gps_nugen = [p for p in gps_final if abs(p.pdgId()) == 16]
gps_tau = [
p for p in gps_final if abs(p.pdgId()) == 15 and p.status() == 2
]
gps_tau.sort(key=lambda p: p.pt(), reverse=True)
gps_taucut = [p for p in gps_tau if p.pt() > 20 and abs(p.eta()) < 2.5]
gps_taucut50 = [
p for p in gps_tau if p.pt() > 50 and abs(p.eta()) < 2.5
]
gps_tau_vis = []
gps_taucut_vis = []
gps_taucut50_vis = []
# find taus that survive ptvis > 20 and >50 cuts
for p in gps_tau:
while p.status() != 2:
p = p.daughter(0)
findau = finalDaughters(p, [])
thisptvis = p4sumvis(findau).pt()
thisetavis = p4sumvis(findau).eta()
gps_tau_vis.append(p)
if thisptvis > 20 and abs(thisetavis) < 2.5:
gps_taucut_vis.append(p)
if thisptvis > 50 and abs(thisetavis) < 2.5:
gps_taucut50_vis.append(p)
#print '-'*10
#for p in gps_tgen:
# printParticle(p)
#if gps_tgen:
# print "has top"
#for p in gps_nugen:
# printParticle(p)
# REMOVE TOP QUARK if its final daughter is also in the list
for top in gps_tgen[:]:
dau = top
while abs(dau.daughter(0).pdgId()) == 6:
dau = dau.daughter(0)
if dau != top and dau in gps_tgen:
gps_tgen.remove(top)
# REMOVE JET-LEPTON OVERLAP
jets, dummy = cleanObjectCollection(jets, gps_tau, dRmin=0.5)
njets = 0
sumjet = 0
jets30 = []
for jet in jets:
if jet.pt() > 30 and abs(jet.eta()) < 5:
sumjet += jet.pt()
njets += 1
tree_jet.pt[0] = jet.pt()
tree_jet.eta[0] = jet.eta()
tree_jet.phi[0] = jet.phi()
tree_jet.weight[0] = weight
tree_jet.Fill()
jets30.append(jet)
# MULTIPLICITIES
tree_event.nlq[0] = len(gps_mother)
tree_event.nbcut[0] = len(gps_bcut)
tree_event.nbcut50[0] = len(gps_bcut50)
tree_event.nbgen[0] = len(gps_bgen)
tree_event.ntgen[0] = len(gps_tgen)
tree_event.njet[0] = njets
tree_event.nlepton[0] = len(gps_tau)
tree_event.ntau[0] = len(gps_tau)
tree_event.ntaucut[0] = len(gps_taucut)
tree_event.ntaucut50[0] = len(gps_taucut50)
tree_event.ntaucut_vis[0] = len(gps_taucut_vis)
tree_event.ntaucut50_vis[0] = len(gps_taucut50_vis)
tree_event.nnu[0] = len(gps_nugen)
# JETS
tree_event.met[0] = met[0].pt()
tree_event.sumjet[0] = sumjet
if len(jets30) >= 2:
centrajpt1s = findCentrajpt1s(jets30[:2], jets30[2:])
tree_event.ncentral[0] = len(centrajpt1s)
else:
tree_event.ncentral[0] = -9
if (len(jets30) >= 2):
tree_event.jpt1[0] = jets30[0].pt()
tree_event.jpt2[0] = jets30[1].pt()
tree_event.dphi_jj[0] = deltaPhi(jets30[0].phi(), jets30[1].phi())
tree_event.deta_jj[0] = jets30[0].eta() - jets30[1].eta()
tree_event.dr_jj[0] = deltaR(jets30[0].eta(), jets30[0].phi(),
jets30[1].eta(), jets30[1].phi())
dijetp4 = jets30[0].p4() + jets30[1].p4()
tree_event.mjj[0] = dijetp4.M()
elif (len(jets30) == 1):
tree_event.jpt1[0] = jets30[0].pt()
tree_event.jpt2[0] = -1
tree_event.dphi_jj[0] = -9
tree_event.deta_jj[0] = -9
tree_event.dr_jj[0] = -1
tree_event.mjj[0] = -1
else:
tree_event.jpt1[0] = -1
tree_event.jpt2[0] = -1
tree_event.dphi_jj[0] = -9
tree_event.deta_jj[0] = -9
tree_event.dr_jj[0] = -1
tree_event.mjj[0] = -1
# SCALAR SUM PT
# if len(gps_taucut)>=2 and len(gps_bcut)>=1:
if len(gps_tau_vis) >= 2 and len(gps_bcut) >= 1:
st = 0
#gps_taucut.sort(key=lambda p: p.pt(), reverse=True)
gps_bcut.sort(key=lambda p: p.pt(), reverse=True)
#taus_assoc.sort(key=lambda p: p.pt(), reverse=True)
#taus_decay.sort(key=lambda p: p.pt(), reverse=True)
#bgen_decay.sort(key=lambda p: p.pt(), reverse=True)
for tau in gps_tau_vis[:2]:
st += p4sumvis(finalDaughters(tau, [])).pt()
st += gps_bcut[0].pt()
# for part in gps_taucut[:2]+gps_bcut[:1]:
# st += part.pt()
stmet = st + met[0].pt()
else:
st = -1
stmet = -1
if len(gps_tau) > 0:
tree_event.tau1_pt[0] = gps_tau[0].pt()
tree_event.tau1_eta[0] = gps_tau[0].eta()
tree_event.tau1_y[0] = gps_tau[0].p4().Rapidity()
tree_event.tau1_ptvis[0] = p4sumvis(
finalDaughters(gps_tau_vis[0], [])).pt()
tree_event.tau1_etavis[0] = p4sumvis(
finalDaughters(gps_tau_vis[0], [])).eta()
if len(gps_tau) > 1:
tree_event.tau2_pt[0] = gps_tau[1].pt()
tree_event.tau2_eta[0] = gps_tau[1].eta()
tree_event.tau2_y[0] = gps_tau[1].p4().Rapidity()
tree_event.tau2_ptvis[0] = p4sumvis(
finalDaughters(gps_tau_vis[1], [])).pt()
tree_event.tau2_etavis[0] = p4sumvis(
finalDaughters(gps_tau_vis[1], [])).eta()
dy = abs(gps_tau[0].p4().Rapidity() - gps_tau[1].p4().Rapidity())
tree_event.ditau_dy[0] = dy
tree_event.ditau_chi[0] = exp(dy)
tree_event.st[0] = st
tree_event.st_met[0] = stmet
tree_mother.st[0] = st
tree_mother.st_met[0] = stmet
tree_event.weight[0] = weight
#print 'len, gps_mother = ', len(gps_mother)
#if len(gps_mother)==1:
# print gps_mother[0].pdgId(), gps_mother[0].status(), gps_mother[0].pt(), gps_mother[0].eta(), gps_mother[0].phi()
# print '1 (ndaughter, daughter pdgid) =', gps_mother[0].numberOfDaughters(), gps_mother[0].daughter(0).pdgId(), '(pdgId, status, pt, eta, phi) = ', gps_mother[0].pdgId(), gps_mother[0].status(), gps_mother[0].pt(), gps_mother[0].eta(), gps_mother[0].phi()
#if len(gps_mother)>=2:
# print '2 (ndaughter, daughter 1/2 pdgid) =', gps_mother[0].numberOfDaughters(), gps_mother[0].daughter(0).pdgId(), gps_mother[0].daughter(1).pdgId(), '(pdgId, status, pt, eta, phi) = ', gps_mother[0].pdgId(), gps_mother[0].status(), gps_mother[0].pt(), gps_mother[0].eta(), gps_mother[0].phi()
# print '2 (ndaughter, daughter 1/2 pdgid) =', gps_mother[1].numberOfDaughters(), gps_mother[1].daughter(0).pdgId(), gps_mother[1].daughter(1).pdgId(), '(pdgId, status, pt, eta, phi) = ', gps_mother[1].pdgId(), gps_mother[1].status(), gps_mother[1].pt(), gps_mother[1].eta(), gps_mother[1].phi()
# TAU
taus_assoc = []
for gentau in gps_tau:
while gentau.status() != 2:
gentau = gentau.daughter(0)
genfinDaughters = finalDaughters(gentau, [])
genptvis = p4sumvis(genfinDaughters).pt()
# CHECK MOTHER
taumoth = gentau.mother(0)
mothpid = abs(taumoth.pdgId())
from_LQ = False
#from_had = False # from hadron decay
#print '-'*30
while mothpid != 2212:
#print taumoth.pdgId()
if mothpid in lqids:
from_LQ = True
break
elif 100 < mothpid < 10000: #and mothpid!=2212:
#from_had = True
break
taumoth = taumoth.mother(0)
mothpid = abs(taumoth.pdgId())
# ASSOC
if not from_LQ:
tree_assoc.pt[0] = gentau.pt()
tree_assoc.ptvis[0] = genptvis
tree_assoc.eta[0] = gentau.eta()
tree_assoc.phi[0] = gentau.phi()
tree_assoc.pid[0] = gentau.pdgId()
tree_assoc.moth[0] = taumoth.pdgId()
tree_assoc.weight[0] = weight
tree_assoc.Fill()
#if not from_had:
taus_assoc.append(gentau)
# B QUARK
for genb in gps_bgen:
bmoth = genb.mother(0)
mothpid = abs(bmoth.pdgId())
from_LQ = False
while mothpid != 2212:
if mothpid in lqids:
from_LQ = True
break
bmoth = bmoth.mother(0)
mothpid = abs(bmoth.pdgId())
if not from_LQ:
tree_assoc.pt[0] = genb.pt()
tree_assoc.ptvis[0] = -1
tree_assoc.eta[0] = genb.eta()
tree_assoc.phi[0] = genb.phi()
tree_assoc.pid[0] = genb.pdgId()
tree_assoc.moth[0] = bmoth.pdgId()
tree_assoc.weight[0] = weight
tree_assoc.Fill()
# MOTHER LQ
#print '-'*80
taus_decay = []
bgen_decay = []
gps_mother.sort(key=lambda p: p.pt(), reverse=True)
for moth in gps_mother:
dau_pid = 0
pair = []
if moth.numberOfDaughters() == 2:
if moth.daughter(0).pdgId() in [
21, 22
] or moth.daughter(1).pdgId() in [21, 22]:
continue
if abs(moth.daughter(0).pdgId()
) in lqids: # single production with t-channel LQ
continue
lq_moth = moth.mother(0)
while abs(lq_moth.pdgId()) in lqids:
lq_moth = lq_moth.mother(0)
for i in range(moth.numberOfDaughters()):
#print '\t', dau.pdgId()
dau = moth.daughter(i)
# TAU
isBrem = False
if abs(dau.pdgId()) == 15:
while dau.status() != 2:
dau = dau.daughter(0)
if dau.numberOfDaughters() == 2 and abs(
dau.daughter(0).pdgId()) == 15 and dau.daughter(
1).pdgId() == 22:
#print "This is brems !?!"
isBrem = True
else:
taus_decay.append(dau)
# BOTTOM QUARK
elif abs(dau.pdgId()) == 5:
dau_pid = dau.pdgId()
bgen_decay.append(dau)
# TOP QUARK
elif abs(dau.pdgId()) == 6:
dau_pid = dau.pdgId()
newdau = dau
while abs(newdau.daughter(0).pdgId()) == 6:
newdau = newdau.daughter(0)
if isFinal(newdau):
dau = newdau
pair.append(dau.p4())
tree_decay.lq_mass[0] = moth.mass()
tree_decay.pid[0] = dau.pdgId()
tree_decay.pt[0] = dau.pt()
tree_decay.eta[0] = dau.eta()
tree_decay.phi[0] = dau.phi()
tree_decay.isBrem[0] = isBrem
if abs(dau.pdgId()) == 15:
finDaughters = finalDaughters(dau, [])
ptvis = p4sumvis(finDaughters).pt()
tree_decay.ptvis[0] = ptvis
decaymode = tauDecayMode(dau)
tree_decay.type[0] = decaydict[decaymode]
#print decaymode, 'vis pt = ', ptvis , 'tau pt = ', dau.pt()
if ptvis > dau.pt():
print "%s, vis pt = %s, tau pt = %s " % (
decaymode, ptvis, dau.pt()) + '!' * 30
else:
tree_decay.ptvis[0] = dau.pt()
tree_decay.type[0] = -1
tree_decay.weight[0] = weight
tree_decay.Fill()
if abs(moth.pdgId()) in lqids:
hist_LQ_decay.Fill(dau.pdgId())
if len(pair) == 2:
tree_mother.inv[0] = (pair[0] + pair[1]).mass()
tree_mother.dphi_ll[0] = deltaPhi(pair[0].phi(), pair[1].phi())
tree_mother.deta_ll[0] = pair[0].eta() - pair[1].eta()
tree_mother.dr_ll[0] = deltaR(pair[0].eta(), pair[0].phi(),
pair[1].eta(), pair[1].phi())
else:
tree_mother.inv[0] = -1
tree_mother.dphi_ll[0] = -99
tree_mother.deta_ll[0] = -99
tree_mother.dr_ll[0] = -99
tree_mother.pid[0] = moth.pdgId()
tree_mother.moth[0] = lq_moth.pdgId()
tree_mother.status[0] = moth.status()
tree_mother.mass[0] = moth.mass()
tree_mother.pt[0] = moth.pt()
tree_mother.eta[0] = moth.eta()
tree_mother.phi[0] = moth.phi()
tree_mother.ndau[0] = len(pair)
tree_mother.dau[0] = dau_pid # save PDG ID for quark daughter
tree_mother.weight[0] = weight
tree_mother.Fill()
if len(gps_mother) == 1:
tree_event.lq1_mass[0] = gps_mother[0].mass()
tree_event.lq1_pt[0] = gps_mother[0].pt()
tree_event.lq2_mass[0] = -1
tree_event.lq2_pt[0] = -1
tree_event.m_lqlq[0] = -1
elif len(gps_mother) >= 2:
tree_event.lq1_mass[0] = gps_mother[0].mass()
tree_event.lq1_pt[0] = gps_mother[0].pt()
tree_event.lq2_mass[0] = gps_mother[1].mass()
tree_event.lq2_pt[0] = gps_mother[1].pt()
dilqp4 = gps_mother[0].p4() + gps_mother[1].p4()
tree_event.m_lqlq[0] = dilqp4.M()
else:
tree_event.lq1_mass[0] = -1
tree_event.lq1_pt[0] = -1
tree_event.lq2_mass[0] = -1
tree_event.lq2_pt[0] = -1
tree_event.m_lqlq[0] = -1
tree_event.ntau_assoc[0] = len(taus_assoc)
tree_event.ntau_decay[0] = len(taus_decay)
tree_event.nbgen_decay[0] = len(bgen_decay)
tree_event.Fill()
print ">>> processed %4s events in %s" % (
evtid, formatTime(time.time() - start_proc))
print ">>> writing to output file %s..." % (outfilename)
outfile.Write()
outfile.Close()
print ">>> done in in %s" % (formatTime(time.time() - start1))
