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