def lhe_to_csv(importfile, exportfile, particlei=0, particlef=None):
"""
A barebones way to extract all the information from each event from a
.LHE file and write it to a .csv file
I choose to not do the xml processing myself, but to use pylhe, which is
available at https://github.com/lukasheinrich/pylhe
Note that pylhe uses Python2, not Python3!
if particlei or particlef are non-zero not all of the particles in the
event will be written to file
"""
LHEgen = pylhe.readLHE(importfile)
with open(exportfile, 'w') as file1:
writes = csv.writer(file1, delimiter=',')
for event in LHEgen:
eventlist = []
for particle in event['particles'][particlei:particlef]:
eventlist += particle.values()
writes.writerow(eventlist)
output_tree.Branch("spin", spin, "spin/F")
output_tree.Branch("pdgid", pdgid, "pdgid/F")
h = ROOT.TH1F('invmass', 'Invariant Mass of Final State', 150, 0, 175)
h.SetFillColor(38)
pos = ROOT.TH1F('pos', 'Invariant Mass of Final pos State', 150, 0, 175)
neg = ROOT.TH1F('neg', 'Invariant Mass of Final neg State', 150, 0, 175)
def invariant_mass(p1, p2):
return math.sqrt(
sum((1 if mu == 'e' else -1) *
(getattr(p1, mu) + getattr(p2, mu))**2
for mu in ['e', 'px', 'py', 'pz']))
num = 0
for e in pylhe.readLHE(args.input_file):
num = num + 1
if num % 1000 == 0:
print "Processing event number: ", num
w = 1
if "0" in args.pdf_type:
w = 1
if "1" in args.pdf_type:
w = float(e.pdfwe[319]) / float(e.pdfwe[1])
if "2" in args.pdf_type:
w = float(e.pdfwe[379]) / float(e.pdfwe[1])
if "3" in args.pdf_type:
w = float(e.pdfwe[1009]) / float(e.pdfwe[1])
if "4" in args.pdf_type:
w = float(e.pdfwe[430]) / float(e.pdfwe[1])
if "5" in args.pdf_type:
import pylhe
import ROOT
for event in pylhe.readLHE(
"/home/hannah/background_decay/Events/run_01/unweighted_events.lhe"):
print event
particles = event.particles
print particles
for particle in particles:
mass = particle["m"]
print mass
# check on visible transverse momentum
if (momentaLeptonsOut[0].pT+momentaLeptonsOut[1].pT+\
ev.outgoingBottom[0].pT+ev.outgoingBottom[1].pT)<130:
return False
# checks on bottom quarks
for i in ev.outgoingBottom:
if i.pT < 25: return False
if abs(i.rap) > 2.4: return False
for j in momentaLeptonsOut:
if i.R2(j) < 0.4**2: return False
return True # only gets reached, if no check failed.
evtFile = pylhe.readLHE(eventFile)
fullResultLL = ngl.Hist(nbins, tmax, errorHistCalc=True)
fullNGL1Loop = 0.
fullNGL1LoopSq = 0.
fullNGL2Loop = 0.
fullNGL2LoopSq = 0.
eventWeight = 0.
numberEvents = 0
numberValidEvents = 0
timeStart = time.time()
for event in evtFile:
tr.Branch('theta2', theta2, 'theta2/F')
tr.Branch('Phi', Phi, 'Phi/F')
tr.Branch('Phi1', Phi1, 'Phi1/F')
tr.Branch('thetaStar', thetaStar, 'thetaStar/F')
tr.Branch('thetaStar2', thetaStar2, 'thetaStar2/F')
tr.Branch('mZZ', mZZ, 'mZZ/F')
tr.Branch('PtZZ', PtZZ, 'PtZZ/F')
tr.Branch('Ptll', Ptll, 'Ptll/F')
tr.Branch('Ptqq', Ptqq, 'Ptqq/F')
print('Processing "', inputDire + args[1] + '.lhe', '....')
if args[1] == 'Zjets': isbkg = True
else: isbkg = False
for e in pylhe.readLHE(inputDire + args[1] + '.lhe'):
X = TLorentzVector(0., 0., 0., 0.)
Z1 = TLorentzVector(0., 0., 0., 0.)
Z2 = TLorentzVector(0., 0., 0., 0.)
l1 = TLorentzVector(0., 0., 0., 0.)
l2 = TLorentzVector(0., 0., 0., 0.)
q1 = TLorentzVector(0., 0., 0., 0.)
q2 = TLorentzVector(0., 0., 0., 0.)
isMuon = False
qcount = 0
for i in range(len(e['particles'])):
# Check if final state
if e['particles'][i]['status'] == 1:
pdgid = e['particles'][i]['id']
Px = e['particles'][i]['px']
Py = e['particles'][i]['py']
import time
import numpy as np
#number of jets per event
# fig = plt.figure()
c = ROOT.TCanvas()
h = ROOT.TH1F('jets', 'Distribution of Jets at Parton Level', 7, 0, 7)
# events = [np.zeroes(10000])
events = []
x = 0
y = 0
for e in pylhe.readLHE(
'/Users/margaretlazarovits/MG5_aMC_v2_6_6/PROC_sm_0/Events/run_01/unweighted_events.lhe'
):
# x+=1
# print(e)
njets = 0
for p_idx in range(len(e['particles'])):
if e['particles'][p_idx]['id'] == 21.0 and e['particles'][p_idx][
'mother1'] != 0 and e['particles'][p_idx]['mother2'] != 0:
# h.Fill(e['particles'][p_idx]['id'])
# print('particle id: ', get_id(e['particles'][p_idx]))
njets += 1
events.append(njets)
# print('\n')
# (e['particles'][p_idx]['id'] >= 1.0 and e['particles'][p_idx]['id'] <= 4.0) or (e['particles'][p_idx]['id'] <= -1.0 and e['particles'][p_idx]['id'] >= -4.0) or
def main(infname, outfname, rwfunc):
import ROOT
import pylhe
import itertools as it
outf = ROOT.TFile(outfname, "CREATE")
hpz_g1 = ROOT.TH1F("hpz_g1", "hpz_g1", 100, 0, 4 * TeV)
hpz_g2 = ROOT.TH1F("hpz_g2", "hpz_g2", 100, 0, 4 * TeV)
hmgg = ROOT.TH1F("hmgg", "hmgg", 100, 0, 4 * TeV)
hpzgg = ROOT.TH1F("hpzgg", "hpzgg", 100, 0, 4 * TeV)
hptv1 = ROOT.TH1F("hptv1", "hptv1", 100, 0, 2 * TeV)
hpzv1 = ROOT.TH1F("hpzv1", "hpzv1", 100, 0, 2 * TeV)
hptt1 = ROOT.TH1F("hptt1", "hptt1", 100, 0, 1 * TeV)
hptt2 = ROOT.TH1F("hptt2", "hptt2", 100, 0, 1 * TeV)
hptt3 = ROOT.TH1F("hptt3", "hptt3", 100, 0, 1 * TeV)
hptt4 = ROOT.TH1F("hptt4", "hptt4", 100, 0, 1 * TeV)
habsetat1 = ROOT.TH1F("habsetat1", "habsetat1", 100, 0, 5)
habsetat2 = ROOT.TH1F("habsetat2", "habsetat2", 100, 0, 5)
habsetat3 = ROOT.TH1F("habsetat3", "habsetat3", 100, 0, 5)
habsetat4 = ROOT.TH1F("habsetat4", "habsetat4", 100, 0, 5)
hmtt_centraltops = ROOT.TH1F("hmtt_centraltops", "hmtt_centraltops", 100,
0, 2 * TeV)
hmtt_leadingtops = ROOT.TH1F("hmtt_leadingtops", "hmtt_leadingtops", 100,
0, 2 * TeV)
hmtt_fromV = ROOT.TH1F("hmtt_fromV", "hmtt_fromV", 100, 0, 2 * TeV)
hpt_t1_fromV = ROOT.TH1F("hpt_t1_fromV", "hpt_t1_fromV", 100, 0, 1 * TeV)
hpt_t2_fromV = ROOT.TH1F("hpt_t2_fromV", "hpt_t2_fromV", 100, 0, 1 * TeV)
hpt_t1_notFromV = ROOT.TH1F("hpt_t1_notFromV", "hpt_t1_notFromV", 100, 0,
1 * TeV)
hpt_t2_notFromV = ROOT.TH1F("hpt_t2_notFromV", "hpt_t2_notFromV", 100, 0,
1 * TeV)
hpz_t1_fromV = ROOT.TH1F("hpz_t1_fromV", "hpz_t1_fromV", 100, 0, 1 * TeV)
hpz_t2_fromV = ROOT.TH1F("hpz_t2_fromV", "hpz_t2_fromV", 100, 0, 1 * TeV)
hpz_t1_notFromV = ROOT.TH1F("hpz_t1_notFromV", "hpz_t1_notFromV", 100, 0,
1 * TeV)
hpz_t2_notFromV = ROOT.TH1F("hpz_t2_notFromV", "hpz_t2_notFromV", 100, 0,
1 * TeV)
habseta_t1_fromV = ROOT.TH1F("habseta_t1_fromV", "habseta_t1_fromV", 100,
0, 5)
habseta_t2_fromV = ROOT.TH1F("habseta_t2_fromV", "habseta_t2_fromV", 100,
0, 5)
habseta_t1_notFromV = ROOT.TH1F("habseta_t1_notFromV",
"habseta_t1_notFromV", 100, 0, 5)
habseta_t2_notFromV = ROOT.TH1F("habseta_t2_notFromV",
"habseta_t2_notFromV", 100, 0, 5)
hpt_t3_notFromV = ROOT.TH1F("hpt_t3_notFromV", "hpt_t3_notFromV", 100, 0,
1 * TeV)
hpt_t4_notFromV = ROOT.TH1F("hpt_t4_notFromV", "hpt_t4_notFromV", 100, 0,
1 * TeV)
habseta_t3_notFromV = ROOT.TH1F("habseta_t3_notFromV",
"habseta_t3_notFromV", 100, 0, 5)
habseta_t4_notFromV = ROOT.TH1F("habseta_t4_notFromV",
"habseta_t4_notFromV", 100, 0, 5)
allHists = \
[ hpz_g1, hpz_g2, hmgg
, hptv1, hpzv1
, hptt1, hptt2, hptt3, hptt4
, habsetat1, habsetat2 , habsetat3, habsetat4
, hmtt_centraltops, hmtt_leadingtops, hmtt_fromV
, hpt_t1_fromV, hpt_t2_fromV
, hpt_t1_notFromV, hpt_t2_notFromV
, hpz_t1_fromV, hpz_t2_fromV
, hpz_t1_notFromV, hpz_t2_notFromV
, habseta_t1_fromV, habseta_t2_fromV
, habseta_t1_notFromV, habseta_t2_notFromV
, hpt_t3_notFromV, hpt_t4_notFromV
, habseta_t3_notFromV, habseta_t4_notFromV
]
rwf = ROOT.TFormula("rwf", rwfunc)
for evt in pylhe.readLHE(infname):
glus = [p for p in evt.particles if abs(p.id) == 21]
v1s = [p for p in evt.particles if abs(p.id) == 6000055]
tops = [p for p in evt.particles if abs(p.id) == 6]
# attach the TLV to the particles
for p in glus:
p.tlv = ROOT.TLorentzVector(p.px, p.py, p.pz, p.e)
for p in v1s:
p.tlv = ROOT.TLorentzVector(p.px, p.py, p.pz, p.e)
for p in tops:
p.tlv = ROOT.TLorentzVector(p.px, p.py, p.pz, p.e)
tops.sort(key=lambda t: t.tlv.Pt(), reverse=True)
wgt = evt.eventinfo.weight
if len(v1s) == 1:
wgt *= rwf.Eval(v1s[0].tlv.M())
hptv1.Fill(v1s[0].tlv.Pt(), wgt)
hpzv1.Fill(abs(v1s[0].tlv.Pz()), wgt)
if len(glus) == 2:
hpz_g1.Fill(abs(glus[0].tlv.Pz()), wgt)
hpz_g2.Fill(abs(glus[1].tlv.Pz()), wgt)
gg = glus[0].tlv + glus[1].tlv
hmgg.Fill(gg.M(), wgt)
hpzgg.Fill(gg.Pz(), wgt)
if len(tops) != 4:
print "woops!"
continue
hptt1.Fill(tops[0].tlv.Pt(), wgt)
hptt2.Fill(tops[1].tlv.Pt(), wgt)
hptt3.Fill(tops[2].tlv.Pt(), wgt)
hptt4.Fill(tops[3].tlv.Pt(), wgt)
habsetat1.Fill(abs(tops[0].tlv.Eta()), wgt)
habsetat2.Fill(abs(tops[1].tlv.Eta()), wgt)
habsetat3.Fill(abs(tops[2].tlv.Eta()), wgt)
habsetat4.Fill(abs(tops[3].tlv.Eta()), wgt)
hmtt_leadingtops.Fill((tops[0].tlv + tops[1].tlv).M(), wgt)
eta_ordered_tops = sorted(map(lambda t: (abs(t.tlv.Eta()), t), tops))
hmtt_centraltops.Fill(
(eta_ordered_tops[0][1].tlv + eta_ordered_tops[1][1].tlv).M(), wgt)
(topsFromV, topsNotFromV) = list(
partition(lambda p: hasAncestor(p, 6000055), tops))
if len(topsFromV) == 2:
hpt_t1_fromV.Fill(topsFromV[0].tlv.Pt(), wgt)
hpt_t2_fromV.Fill(topsFromV[1].tlv.Pt(), wgt)
hpz_t1_fromV.Fill(topsFromV[0].tlv.Pz(), wgt)
hpz_t2_fromV.Fill(topsFromV[1].tlv.Pz(), wgt)
habseta_t1_fromV.Fill(abs(topsFromV[0].tlv.Eta()), wgt)
habseta_t2_fromV.Fill(abs(topsFromV[1].tlv.Eta()), wgt)
hmtt_fromV.Fill((topsFromV[0].tlv + topsFromV[1].tlv).M(), wgt)
if len(topsNotFromV) >= 2:
hpt_t1_notFromV.Fill(topsNotFromV[0].tlv.Pt(), wgt)
hpt_t2_notFromV.Fill(topsNotFromV[1].tlv.Pt(), wgt)
hpz_t1_notFromV.Fill(topsNotFromV[0].tlv.Pz(), wgt)
hpz_t2_notFromV.Fill(topsNotFromV[1].tlv.Pz(), wgt)
habseta_t1_notFromV.Fill(abs(topsNotFromV[0].tlv.Eta()), wgt)
habseta_t2_notFromV.Fill(abs(topsNotFromV[1].tlv.Eta()), wgt)
if len(topsNotFromV) >= 4:
hpt_t3_notFromV.Fill(topsNotFromV[2].tlv.Pt(), wgt)
hpt_t4_notFromV.Fill(topsNotFromV[3].tlv.Pt(), wgt)
habseta_t3_notFromV.Fill(abs(topsNotFromV[2].tlv.Eta()), wgt)
habseta_t4_notFromV.Fill(abs(topsNotFromV[3].tlv.Eta()), wgt)
# scale by 100/fb
for h in allHists:
h.Scale(100e3)
outf.Write()
outf.Close()
return
def main():
t0 = time.time()
# Check user has inputted variables or not
# Parse in arguments
parser = argparse.ArgumentParser(
description=
'This parses the .lhe Les Houches Events file and produces .root TTree ntuples. Run as ./lhe_to_ntuples -i unweighted_events.lhe -o out_tree.root'
)
parser.add_argument('-i',
'--ifile',
type=str,
nargs='?',
help='Input LHE file',
default='unweighted_events.lhe')
parser.add_argument('-o',
'--ofile',
type=str,
nargs='?',
help='Name of output .root file',
default='out_tree.root')
args = parser.parse_args()
if args.ifile:
in_file = args.ifile
print('Input file: {0}'.format(in_file))
if 'lhe.gz' in in_file:
print('Unzipping LHE file to read')
os.system('gunzip {0}'.format(in_file))
in_file = in_file.replace('.lhe.gz', '.lhe')
if args.ofile:
out_file = args.ofile
print('Output file: {0}'.format(out_file))
#-------------------------------------------------
# Input LHE file
#-------------------------------------------------
print('Input .lhe file: {0}'.format(in_file))
# Use pylhe to parse the LHE file
lhef = pylhe.readLHE(in_file)
#-------------------------------------------------
# Output tree
#-------------------------------------------------
outFile = TFile(out_file, "recreate")
outTree = TTree('tree', 'tree')
#-------------------------------------------------
# Prepare variables to write to TTree branches
#-------------------------------------------------
h1_M = array('f', [0])
h1_Pt = array('f', [0])
h1_Eta = array('f', [0])
h1_Phi = array('f', [0])
h2_M = array('f', [0])
h2_Pt = array('f', [0])
h2_Eta = array('f', [0])
h2_Phi = array('f', [0])
m_hh = array('f', [0])
pT_hh = array('f', [0])
eta_hh = array('f', [0])
phi_hh = array('f', [0])
dEta_hh = array('f', [0])
# Declare TTree branches
outTree.Branch('h1_M', h1_M, 'h1_M/F')
outTree.Branch('h1_Pt', h1_Pt, 'h1_Pt/F')
outTree.Branch('h1_Eta', h1_Eta, 'h1_Eta/F')
outTree.Branch('h1_Phi', h1_Phi, 'h1_Phi/F')
outTree.Branch('h2_M', h2_M, 'h2_M/F')
outTree.Branch('h2_Pt', h2_Pt, 'h2_Pt/F')
outTree.Branch('h2_Eta', h2_Eta, 'h2_Eta/F')
outTree.Branch('h2_Phi', h2_Phi, 'h2_Phi/F')
outTree.Branch('m_hh', m_hh, 'm_hh/F')
outTree.Branch('pT_hh', pT_hh, 'pT_hh/F')
outTree.Branch('eta_hh', eta_hh, 'eta_hh/F')
outTree.Branch('phi_hh', phi_hh, 'phi_hh/F')
outTree.Branch('dEta_hh', dEta_hh, 'dEta_hh/F')
# Declare TLorentzVectors
tlv_unordered_higgs1 = TLorentzVector()
tlv_unordered_higgs2 = TLorentzVector()
tlv_ordered_higgs1 = TLorentzVector()
tlv_ordered_higgs2 = TLorentzVector()
tlv_DiHiggs = TLorentzVector()
#-------------------------------------------------
# Loop through all events in the parsed LHE file
#-------------------------------------------------
for count, event in enumerate(lhef):
# debug break
#if count > 10: break
if count % 1000 == 0:
sys.stdout.write('\r Event {0}'.format(count))
sys.stdout.flush()
count_higgs = 0
#-------------------------------------------------
# Loop through all particles
#-------------------------------------------------
for part in event.particles:
pid = int(part.id)
# PDG ID of Higgs boson is 25
if abs(pid) == 25:
count_higgs += 1
if count_higgs == 1:
tlv_unordered_higgs1.SetPxPyPzE(part.px, part.py, part.pz,
part.e)
if count_higgs == 2:
tlv_unordered_higgs2.SetPxPyPzE(part.px, part.py, part.pz,
part.e)
#-------------------------------------------------
# Order Higgs by pT
#-------------------------------------------------
# if Higgs2Pt is has larger pT than Higgs1Pt, invert the tlv assignment
if tlv_unordered_higgs2.Pt() > tlv_unordered_higgs1.Pt():
tlv_ordered_higgs1 = tlv_unordered_higgs2
tlv_ordered_higgs2 = tlv_unordered_higgs1
else:
tlv_ordered_higgs1 = tlv_unordered_higgs1
tlv_ordered_higgs2 = tlv_unordered_higgs2
tlv_DiHiggs = tlv_ordered_higgs1 + tlv_ordered_higgs2
#-------------------------------------------------
# Set variables for TTree
#-------------------------------------------------
h1_M[0] = tlv_ordered_higgs1.M()
h1_Pt[0] = tlv_ordered_higgs1.Pt()
h1_Eta[0] = tlv_ordered_higgs1.Eta()
h1_Phi[0] = tlv_ordered_higgs1.Phi()
h2_M[0] = tlv_ordered_higgs2.M()
h2_Pt[0] = tlv_ordered_higgs2.Pt()
h2_Eta[0] = tlv_ordered_higgs2.Eta()
h2_Phi[0] = tlv_ordered_higgs2.Phi()
m_hh[0] = tlv_DiHiggs.M()
pT_hh[0] = tlv_DiHiggs.Pt()
eta_hh[0] = tlv_DiHiggs.Eta()
phi_hh[0] = tlv_DiHiggs.Phi()
dEta_hh[0] = tlv_ordered_higgs1.Eta() - tlv_ordered_higgs2.Eta()
# Fill the tree
outTree.Fill()
print('\nFinished filling tree, closing files...')
# Close up shop
outFile.Write()
outFile.Close()
os.system('gzip {0}'.format(in_file))
dt = time.time() - t0
print('Finished in {0:.2f} seconds.'.format(dt))
output_tree.Branch("pz",pz , "pz/F")
output_tree.Branch("energy",energy , "energy/F")
output_tree.Branch("mass",mass , "mass/F")
output_tree.Branch("zmass",zmass , "zmass/F")
output_tree.Branch("spin",spin , "spin/F")
output_tree.Branch("pdgid", pdgid, "pdgid/F")
h = ROOT.TH1F('invmass','Invariant Mass of Final State',150,0,175)
h.SetFillColor(38)
def invariant_mass(p1,p2):
return math.sqrt(sum((1 if mu=='e' else -1)*(getattr(p1,mu)+getattr(p2,mu))**2 for mu in ['e','px','py','pz']))
for e in pylhe.readLHE('unweighted_events_80GeVCut.lhe'):
h.Fill(invariant_mass(e.particles[-1],e.particles[-2]),e.eventinfo.weight)
zmass[0]=invariant_mass(e.particles[-1],e.particles[-2]);
for p in e.particles:
pdgid[0] = getattr(p,'id')
px[0] = getattr(p,'px')
py[0] = getattr(p,'py')
pz[0] = getattr(p,'pz')
energy[0] = getattr(p,'e')
mass[0] = getattr(p,'m')
spin[0] = getattr(p,'spin')
output_tree.Fill()
output_tree.Write()