From 36cbda84cdb6a4d520bf764eb01efcbcb03893c0 Mon Sep 17 00:00:00 2001 From: Isaac Wang Date: Sat, 29 Oct 2022 22:45:30 -0400 Subject: [PATCH 1/2] deal with beam particles --- bin/hepmc2root.py | 341 +++++++++++++++++++++++++++------------------- 1 file changed, 203 insertions(+), 138 deletions(-) diff --git a/bin/hepmc2root.py b/bin/hepmc2root.py index 128ee85..67472b4 100755 --- a/bin/hepmc2root.py +++ b/bin/hepmc2root.py @@ -32,55 +32,61 @@ # 15-Apr-2019 HBP test that ROOT can be imported # 31-Jan-2020 HBP make compatible with Python 3 # ----------------------------------------------------------------------- -import os, sys +import os +import sys + try: import ROOT except: - sys.exit('\n\t*** This program uses ROOT! ***\n') + sys.exit("\n\t*** This program uses ROOT! ***\n") from math import sqrt from time import ctime + from pnames import particleName + + # ----------------------------------------------------------------------- def nameonly(s): import posixpath + return posixpath.splitext(posixpath.split(s)[1])[0] -TREENAME= "Events" + +TREENAME = "Events" MAXPART = 5000 debug = 0 + class hepmc2root: - def __init__(self, filename, outfilename=None, treename=TREENAME, complevel=2): # check that file exists - + if not os.path.exists(filename): - sys.exit("** hepmc2root.py: can't open file %s" % \ - filename) + sys.exit("** hepmc2root.py: can't open file %s" % filename) self.inp = open(filename) inp = self.inp # get version number of HepMC - - self.header = [] # cache HepMC header + + self.header = [] # cache HepMC header version = None for line in inp: self.header.append(line) version = str.strip(line) - if version == '': continue + if version == "": + continue token = str.split(version) - if token[0] == 'HepMC::Version': + if token[0] == "HepMC::Version": version = token[1] break else: - sys.exit("** hepmc2root.py: format problem in file %s" % \ - filename) - + sys.exit("** hepmc2root.py: format problem in file %s" % filename) + print("HepMC version: %s" % version) # skip start of listing - + for line in inp: self.header.append(line) break @@ -88,15 +94,14 @@ def __init__(self, filename, outfilename=None, treename=TREENAME, complevel=2): # open output root file if outfilename == None: - outfilename = '%s.root' % nameonly(filename) - + outfilename = "%s.root" % nameonly(filename) + self.file = ROOT.TFile(outfilename, "recreate") - self.tree = ROOT.TTree(treename, 'created: %s HepMC %s' % \ - (ctime(), version)) + self.tree = ROOT.TTree(treename, "created: %s HepMC %s" % (ctime(), version)) # define event struct - - self.struct = '''struct Bag { + + self.struct = """struct Bag { int Event_number; int Event_numberMP; double Event_scale; @@ -136,44 +141,48 @@ def __init__(self, filename, outfilename=None, treename=TREENAME, complevel=2): int Particle_status[%(size)d]; int Particle_d1[%(size)d]; int Particle_d2[%(size)d]; -};''' % {'size': MAXPART} +};""" % { + "size": MAXPART + } # indices to vertices - - self.pvertex = [0]*MAXPART - + + self.pvertex = [0] * MAXPART + # create struct - + ROOT.gROOT.ProcessLine(self.struct) from ROOT import Bag + self.bag = Bag() # create branches - + self.branch = [] - recs = str.split(self.struct, '\n')[1:-1] + recs = str.split(self.struct, "\n")[1:-1] for rec in recs: t = str.split(rec) - if len(t) == 0: continue - + if len(t) == 0: + continue + fmt, name = t T = str.upper(fmt[0]) - name = name[:-1] # skip ";" + name = name[:-1] # skip ";" # check for variable length array - if name[-1] == ']': - field = str.split(name, '[')[0] - fmt = '%s[Event_numberP]/%s' % (field, T) + if name[-1] == "]": + field = str.split(name, "[")[0] + fmt = "%s[Event_numberP]/%s" % (field, T) else: field = name - fmt = '%s/%s' % (field, T) - self.branch.append(self.tree.Branch(field, - ROOT.addressof(self.bag, field), - fmt)) + fmt = "%s/%s" % (field, T) + self.branch.append( + self.tree.Branch(field, ROOT.addressof(self.bag, field), fmt) + ) # list branches - + for ii, b in enumerate(self.branch): bname = b.GetName() - leaves= b.GetListOfLeaves() + leaves = b.GetListOfLeaves() if leaves == None: sys.exit("** hepmc2root: no list of leaves found for branch %s" % bname) leaf = leaves[0] @@ -181,46 +190,48 @@ def __init__(self, filename, outfilename=None, treename=TREENAME, complevel=2): sys.exit("** hepmc2root: no leaf found for branch %s" % bname) leafname = leaf.GetName() leaftype = leaf.GetTypeName() - print("%4d\t%-20s\t%s" % (ii+1, bname, leaftype)) - + print("%4d\t%-20s\t%s" % (ii + 1, bname, leaftype)) + def __del__(self): self.tree.Write("", ROOT.TObject.kOverwrite) - + def __str__(self, index): bag = self.bag - d = " <%4d, %4d>" % (bag.Particle_d1[index], bag.Particle_d2[index]) - px = bag.Particle_px[index] - py = bag.Particle_py[index] - pt = sqrt(px**2+py**2) - rec = '%-14s %7d %4d %3d %7.1f (%7.1f, %7.1f, %7.1f, %7.1f)%s' \ - % (particleName(bag.Particle_pid[index]), - bag.Particle_pid[index], - bag.Particle_barcode[index], - bag.Particle_status[index], - pt, - bag.Particle_energy[index], - bag.Particle_px[index], - bag.Particle_py[index], - bag.Particle_pz[index], - d) + d = " <%4d, %4d>" % (bag.Particle_d1[index], bag.Particle_d2[index]) + px = bag.Particle_px[index] + py = bag.Particle_py[index] + pt = sqrt(px**2 + py**2) + rec = "%-14s %7d %4d %3d %7.1f (%7.1f, %7.1f, %7.1f, %7.1f)%s" % ( + particleName(bag.Particle_pid[index]), + bag.Particle_pid[index], + bag.Particle_barcode[index], + bag.Particle_status[index], + pt, + bag.Particle_energy[index], + bag.Particle_px[index], + bag.Particle_py[index], + bag.Particle_pz[index], + d, + ) return rec def __call__(self): inp = self.inp bag = self.bag - self.event = [] # cache HepMC event in original format - + self.event = [] # cache HepMC event in original format + # find start of event - + token = None for line in inp: self.event.append(line) token = str.split(line) - key = token[0] - if key != 'E': continue + key = token[0] + if key != "E": + continue if debug > 0: - print('BEGIN event') + print("BEGIN event") break else: return False @@ -228,18 +239,18 @@ def __call__(self): if token == None: sys.exit("** hepmc2root.py: can't find start of event") - bag.Event_number = int(token[1]) - bag.Event_numberMP = int(token[2]) # number of multi-particle interactions - bag.Event_scale = float(token[3]) - bag.Event_alphaQCD = float(token[4]) - bag.Event_alphaQED = float(token[5]) - bag.Event_processID = int(token[6]) + bag.Event_number = int(token[1]) + bag.Event_numberMP = int(token[2]) # number of multi-particle interactions + bag.Event_scale = float(token[3]) + bag.Event_alphaQCD = float(token[4]) + bag.Event_alphaQED = float(token[5]) + bag.Event_processID = int(token[6]) bag.Event_barcodeSPV = int(token[7]) - bag.Event_numberV = int(token[8]) # number of vertices in event - bag.Event_barcodeBP1 = int(token[9]) # barcode beam particle 1 - bag.Event_barcodeBP2 = int(token[10]) # barcode beam particle 2 - bag.Event_numberP = 0 # number of particles - + bag.Event_numberV = int(token[8]) # number of vertices in event + bag.Event_barcodeBP1 = int(token[9]) # barcode beam particle 1 + bag.Event_barcodeBP2 = int(token[10]) # barcode beam particle 2 + bag.Event_numberP = 0 # number of particles + if debug > 0: print("\tbarcode 1: %d" % self.barcode1) print("\tbarcode 2: %d" % self.barcode2) @@ -251,84 +262,133 @@ def __call__(self): token = str.split(line) key = token[0] - if key == 'C': + if key == "C": # CROSS SECTION bag.Xsection_value = float(token[1]) bag.Xsection_error = float(token[2]) if debug > 0: - print("\tcross section: %10.3e +\- %10.3e pb" % \ - (bag.Xsection_value, bag.Xsection_error)) - - elif key == 'F': + print( + "\tcross section: %10.3e +\- %10.3e pb" + % (bag.Xsection_value, bag.Xsection_error) + ) + + elif key == "F": # PDF INFO - bag.PDF_parton1 = int(token[1]) - bag.PDF_parton2 = int(token[2]) - bag.PDF_x1 = float(token[3]) - bag.PDF_x2 = float(token[4]) - bag.PDF_Q2 = float(token[5]) - bag.PDF_x1f = float(token[6]) - bag.PDF_x2f = float(token[7]) - bag.PDF_id1 = int(token[8]) - bag.PDF_id2 = int(token[9]) + bag.PDF_parton1 = int(token[1]) + bag.PDF_parton2 = int(token[2]) + bag.PDF_x1 = float(token[3]) + bag.PDF_x2 = float(token[4]) + bag.PDF_Q2 = float(token[5]) + bag.PDF_x1f = float(token[6]) + bag.PDF_x2f = float(token[7]) + bag.PDF_id1 = int(token[8]) + bag.PDF_id2 = int(token[9]) if debug > 0: - print('\tfound PDF info') + print("\tfound PDF info") - elif key == 'V': + elif key == "V": # VERTEX vbarcode = int(token[1]) self.vertex[vbarcode] = [-1, -1] - x = float(token[3]) - y = float(token[4]) - z = float(token[5]) + x = float(token[3]) + y = float(token[4]) + z = float(token[5]) ctau = float(token[6]) nout = int(token[8]) if debug > 0: if debug > 1: print("\t%s" % token) - print('\tvertex(barcode): %10d' % vbarcode) - print('\tvertex(count): %10d' % nout) + print("\tvertex(barcode): %10d" % vbarcode) + print("\tvertex(count): %10d" % nout) + # Particles following the vertex are the particles associated + # with it. + # For the vertex with beam particles as incoming, the beam + # particle is also recorded after the vertex. + # Thus there are nout+1 number of particles + # after that vertex. + has_beam_particle = False + for line in inp: + token = str.split(line) + if ( + int(token[11]) == vbarcode + ): # This means that the incoming particles is also here. + # Record this particle, with x=y=z=0 + has_beam_particle = True + self.event.append(line) + index = bag.Event_numberP + bag.Event_numberP += 1 + bag.Particle_x[index] = 0 + bag.Particle_y[index] = 0 + bag.Particle_z[index] = 0 + bag.Particle_ctau[index] = ctau + + bag.Particle_barcode[index] = int(token[1]) + bag.Particle_pid[index] = int(token[2]) + bag.Particle_px[index] = float(token[3]) + bag.Particle_py[index] = float(token[4]) + bag.Particle_pz[index] = float(token[5]) + bag.Particle_energy[index] = float(token[6]) + bag.Particle_mass[index] = float(token[7]) + bag.Particle_status[index] = int(token[8]) + self.pvertex[index] = int(token[11]) + else: + # Then no beam particle here, a normal vertex + # Record this particle as the outgoing one. + self.event.append(line) + index = bag.Event_numberP + bag.Event_numberP += 1 + bag.Particle_x[index] = x + bag.Particle_y[index] = y + bag.Particle_z[index] = z + bag.Particle_ctau[index] = ctau - # particles pertaining to this vertex follow immediately - # after the vertex + bag.Particle_barcode[index] = int(token[1]) + bag.Particle_pid[index] = int(token[2]) + bag.Particle_px[index] = float(token[3]) + bag.Particle_py[index] = float(token[4]) + bag.Particle_pz[index] = float(token[5]) + bag.Particle_energy[index] = float(token[6]) + bag.Particle_mass[index] = float(token[7]) + bag.Particle_status[index] = int(token[8]) + self.pvertex[index] = int(token[11]) + self.vertex[vbarcode][0] == index + break + if ( + not has_beam_particle + ): # No beam particle. So one outgoing is already recorded. + nout -= 1 for ii in range(nout): for line in inp: + token = str.split(line) + key = token[0] self.event.append(line) - token = str.split(line) - if debug > 1: - print("\t%s" % token) - key = token[0] - if key != 'P': - sys.exit("** hepmc2root: faulty event record\n" + line) - - if bag.Event_numberP < MAXPART: - index = bag.Event_numberP - bag.Event_numberP += 1 - - bag.Particle_x[index] = x - bag.Particle_y[index] = y - bag.Particle_z[index] = z - bag.Particle_ctau[index] = ctau - - bag.Particle_barcode[index] = int(token[1]) - bag.Particle_pid[index] = int(token[2]) - bag.Particle_px[index] = float(token[3]) - bag.Particle_py[index] = float(token[4]) - bag.Particle_pz[index] = float(token[5]) - bag.Particle_energy[index] = float(token[6]) - bag.Particle_mass[index] = float(token[7]) - bag.Particle_status[index] = int(token[8]) - self.pvertex[index] = int(token[11]) + index = bag.Event_numberP + bag.Event_numberP += 1 + bag.Particle_x[index] = x + bag.Particle_y[index] = y + bag.Particle_z[index] = z + bag.Particle_ctau[index] = ctau + bag.Particle_barcode[index] = int(token[1]) + bag.Particle_pid[index] = int(token[2]) + bag.Particle_px[index] = float(token[3]) + bag.Particle_py[index] = float(token[4]) + bag.Particle_pz[index] = float(token[5]) + bag.Particle_energy[index] = float(token[6]) + bag.Particle_mass[index] = float(token[7]) + bag.Particle_status[index] = int(token[8]) + self.pvertex[index] = int(token[11]) + if has_beam_particle: if ii == 0: - self.vertex[vbarcode][0] = index + self.vertex[vbarcode][0] == index else: - self.vertex[vbarcode][1] = index - + self.vertex[vbarcode][1] == index + break + else: + self.vertex[vbarcode][1] == index break - else: - return False - + if len(self.vertex) >= bag.Event_numberV: for index in range(bag.Event_numberP): code = self.pvertex[index] @@ -341,32 +401,36 @@ def __call__(self): bag.Particle_d2[index] = -1 # fill ntuple - + self.file.cd() self.tree.Fill() - + return True else: return False - + def printTable(self): for ii in xrange(self.bag.Event_numberP): print("%4d\t%s" % (ii, self.__str__(ii))) -# ----------------------------------------------------------------------- + + +# ----------------------------------------------------------------------- def main(): argv = sys.argv[1:] argc = len(argv) if argc < 1: - sys.exit(''' + sys.exit( + """ Usage: ./hepmc2root.py [output root file = .root] - ''') + """ + ) filename = argv[0] if argc > 1: outfilename = argv[1] else: - outfilename = '%s.root' % nameonly(filename) + outfilename = "%s.root" % nameonly(filename) stream = hepmc2root(filename, outfilename) @@ -375,9 +439,10 @@ def main(): if ii % 1000 == 0: print(ii) ii += 1 + + # ----------------------------------------------------------------------- try: main() except KeyboardInterrupt: - print('\nciao!') - + print("\nciao!") From 87c214c400366c3e7a9959d5d9358c5aa4540e00 Mon Sep 17 00:00:00 2001 From: quarkquartet Date: Sat, 14 Jan 2023 19:16:00 -0500 Subject: [PATCH 2/2] pandas read daughters correctly --- bin/hepmc2pandas.py | 348 ++++++++++++++++++++++++++------------------ 1 file changed, 206 insertions(+), 142 deletions(-) diff --git a/bin/hepmc2pandas.py b/bin/hepmc2pandas.py index 88d21f5..1fa6135 100755 --- a/bin/hepmc2pandas.py +++ b/bin/hepmc2pandas.py @@ -1,7 +1,7 @@ #!/usr/bin/env python # ----------------------------------------------------------------------- # File: hepmc2pandas.py -# Description: write events in HepMC2 format to a pandas file +# Description: write events in HepMC2 format to a pandas file # # # status = +- (10 * i + j) @@ -30,6 +30,7 @@ # Updated: 04-Dec-2017 HBP add creation vertex (x,y,z) of particles. # ----------------------------------------------------------------------- import os, sys + try: import pandas as pd except: @@ -37,52 +38,56 @@ from math import sqrt from time import ctime from pnames import particleName + # ----------------------------------------------------------------------- def nameonly(s): import posixpath + return posixpath.splitext(posixpath.split(s)[1])[0] -TREENAME= "Events" + +TREENAME = "Events" MAXPART = 5000 debug = 0 + class hepmc2pandas: - def __init__(self, filename): # check that file exists - + if not os.path.exists(filename): sys.exit("** hepmc2pandas: can't open file %s" % filename) self.inp = open(filename) inp = self.inp # get version number of HepMC - - self.header = [] # cache HepMC header + + self.header = [] # cache HepMC header version = None for line in inp: self.header.append(line) version = str.strip(line) - if version == '': continue + if version == "": + continue token = str.split(version) - if token[0] == 'HepMC::Version': + if token[0] == "HepMC::Version": version = token[1] break else: sys.exit("** hepmc2pandas: format problem in file %s" % filename) - + print("HepMC version: %s" % version) # skip start of listing - + for line in inp: self.header.append(line) break # define event structure - - erec = '''{ + + erec = """{ 'Event_number': [], 'Event_numberMP': [], 'Event_scale': [], @@ -105,71 +110,71 @@ def __init__(self, filename): 'PDF_x2f': [], 'PDF_id1': [], 'PDF_id2': [] - }''' - prec = '''{ + }""" + prec = """{ 'Particle_name': [], - 'Particle_x': [], - 'Particle_y': [], - 'Particle_z': [], - 'Particle_ctau': [], - 'Particle_barcode': [], 'Particle_pid': [], + 'Particle_status': [], + 'Particle_energy': [], 'Particle_px': [], 'Particle_py': [], 'Particle_pz': [], - 'Particle_energy': [], 'Particle_mass': [], - 'Particle_status': [], - 'Particle_d1': [], - 'Particle_d2': [] - }''' + 'Particle_x': [], + 'Particle_y': [], + 'Particle_z': [], + 'Particle_daughters': [], + 'Particle_barcode': [], + 'Particle_ctau': [], + }""" self.event = eval(erec) self.particle = prec self.plist = [] - # indices to vertices - self.pvertex = [0]*MAXPART - + # indices to vertices + self.pvertex = [0] * MAXPART + def __del__(self): pass - + def __str__(self, index): p = self.pbag - d = " <%4d, %4d>" % (p['Particle_d1'][index], - p['Particle_d2'][index]) - px = p['Particle_px'][index] - py = p['Particle_py'][index] - pt = sqrt(px**2+py**2) - rec = '%-14s %7d %4d %3d %7.1f (%7.1f, %7.1f, %7.1f, %7.1f)%s' \ - % (particleName(p['Particle_pid'][index]), - p['Particle_pid'][index], - p['Particle_barcode'][index], - p['Particle_status'][index], - pt, - p['Particle_energy'][index], - p['Particle_px'][index], - p['Particle_py'][index], - p['Particle_pz'][index], - d) + d = " <%4d>" % (p["Particle_daughters"][index]) + px = p["Particle_px"][index] + py = p["Particle_py"][index] + pt = sqrt(px**2 + py**2) + rec = "%-14s %7d %4d %3d %7.1f (%7.1f, %7.1f, %7.1f, %7.1f)%s" % ( + particleName(p["Particle_pid"][index]), + p["Particle_pid"][index], + p["Particle_barcode"][index], + p["Particle_status"][index], + pt, + p["Particle_energy"][index], + p["Particle_px"][index], + p["Particle_py"][index], + p["Particle_pz"][index], + d, + ) return rec - + def __call__(self): inp = self.inp # new event self.pbag = eval(self.particle) - + e = self.event p = self.pbag - - # find start of event + + # find start of event token = None for line in inp: token = str.split(line) - key = token[0] - if key != 'E': continue + key = token[0] + if key != "E": + continue if debug > 0: - print('BEGIN event') + print("BEGIN event") break else: return False @@ -177,21 +182,21 @@ def __call__(self): if token == None: sys.exit("** hepmc2pandas: can't find start of event") - e['Event_number'].append(int(token[1])) - e['Event_numberMP'].append(int(token[2])) - e['Event_scale'].append(float(token[3])) - e['Event_alphaQCD'].append(float(token[4])) - e['Event_alphaQED'].append(float(token[5])) - e['Event_processID'].append(int(token[6])) - e['Event_barcodeSPV'].append(int(token[7])) - e['Event_numberV'].append(int(token[8])) # number of vertices in event - e['Event_barcodeBP1'].append(int(token[9])) # barcode beam particle 1 - e['Event_barcodeBP2'].append(int(token[10])) # barcode beam particle 2 - e['Event_numberP'].append(0) # number of particles - + e["Event_number"].append(int(token[1])) + e["Event_numberMP"].append(int(token[2])) + e["Event_scale"].append(float(token[3])) + e["Event_alphaQCD"].append(float(token[4])) + e["Event_alphaQED"].append(float(token[5])) + e["Event_processID"].append(int(token[6])) + e["Event_barcodeSPV"].append(int(token[7])) + e["Event_numberV"].append(int(token[8])) # number of vertices in event + e["Event_barcodeBP1"].append(int(token[9])) # barcode beam particle 1 + e["Event_barcodeBP2"].append(int(token[10])) # barcode beam particle 2 + e["Event_numberP"].append(0) # number of particles + if debug > 2: - print("\tbarcode 1: %d" % e['Event_barcodeBP1'][-1]) - print("\tbarcode 2: %d" % e['Event_barcodeBP2'][-1]) + print("\tbarcode 1: %d" % e["Event_barcodeBP1"][-1]) + print("\tbarcode 2: %d" % e["Event_barcodeBP2"][-1]) self.vertex = {} @@ -199,130 +204,188 @@ def __call__(self): token = str.split(line) key = token[0] - if key == 'C': + if key == "C": # CROSS SECTION - e['Xsection_value'].append(float(token[1])) - e['Xsection_error'].append(float(token[2])) + e["Xsection_value"].append(float(token[1])) + e["Xsection_error"].append(float(token[2])) if debug > 0: - print("\tcross section: %10.3e +\- %10.3e pb" % \ - (e['Xsection_value'][-1], e['Xsection_error'][-1])) - - elif key == 'F': + print( + "\tcross section: %10.3e +\- %10.3e pb" + % (e["Xsection_value"][-1], e["Xsection_error"][-1]) + ) + + elif key == "F": # PDF INFO - e['PDF_parton1'].append(int(token[1])) - e['PDF_parton2'].append(int(token[2])) - e['PDF_x1'].append(float(token[3])) - e['PDF_x2'].append(float(token[4])) - e['PDF_Q2'].append(float(token[5])) - e['PDF_x1f'].append(float(token[6])) - e['PDF_x2f'].append(float(token[7])) - e['PDF_id1'].append(int(token[8])) - e['PDF_id2'].append(int(token[9])) + e["PDF_parton1"].append(int(token[1])) + e["PDF_parton2"].append(int(token[2])) + e["PDF_x1"].append(float(token[3])) + e["PDF_x2"].append(float(token[4])) + e["PDF_Q2"].append(float(token[5])) + e["PDF_x1f"].append(float(token[6])) + e["PDF_x2f"].append(float(token[7])) + e["PDF_id1"].append(int(token[8])) + e["PDF_id2"].append(int(token[9])) if debug > 0: - print('\tfound PDF info') + print("\tfound PDF info") - elif key == 'V': + elif key == "V": # VERTEX vbarcode = int(token[1]) - self.vertex[vbarcode] = [-1, -1] - x = float(token[3]) - y = float(token[4]) - z = float(token[5]) + self.vertex[vbarcode] = [] + x = float(token[3]) + y = float(token[4]) + z = float(token[5]) ctau = float(token[6]) nout = int(token[8]) if debug > 0: if debug > 2: print("\t%s" % token) - print('\tvertex(barcode): %10d' % vbarcode) - print('\tvertex(count): %10d' % nout) + print("\tvertex(barcode): %10d" % vbarcode) + print("\tvertex(count): %10d" % nout) + has_beam_particle = False + for line in inp: + token = str.split(line) + if int(token[11]) == vbarcode and token[0] == "P": + has_beam_particle = True + index = e["Event_numberP"][-1] + e["Event_numberP"][-1] += 1 + + p["Particle_x"].append(0) + p["Particle_y"].append(0) + p["Particle_z"].append(0) + p["Particle_ctau"].append(ctau) + p["Particle_barcode"].append(int(token[1])) + p["Particle_pid"].append(int(token[2])) + + pid = p["Particle_pid"][-1] + p["Particle_name"].append(particleName(pid)) + p["Particle_daughters"].append([]) + p["Particle_px"].append(float(token[3])) + p["Particle_py"].append(float(token[4])) + p["Particle_pz"].append(float(token[5])) + p["Particle_energy"].append(float(token[6])) + p["Particle_mass"].append(float(token[7])) + p["Particle_status"].append(int(token[8])) + self.pvertex[index] = int(token[11]) + else: + index = e["Event_numberP"][-1] + e["Event_numberP"][-1] += 1 + + p["Particle_x"].append(x) + p["Particle_y"].append(y) + p["Particle_z"].append(z) + p["Particle_ctau"].append(ctau) + + p["Particle_barcode"].append(int(token[1])) + p["Particle_pid"].append(int(token[2])) + + pid = p["Particle_pid"][-1] + p["Particle_name"].append(particleName(pid)) + p["Particle_daughters"].append([]) + # p["Particle_d2"].append(-1) + p["Particle_px"].append(float(token[3])) + p["Particle_py"].append(float(token[4])) + p["Particle_pz"].append(float(token[5])) + p["Particle_energy"].append(float(token[6])) + p["Particle_mass"].append(float(token[7])) + p["Particle_status"].append(int(token[8])) + self.pvertex[index] = int(token[11]) + self.vertex[vbarcode].append(int(token[1])) + break + if not has_beam_particle: + nout -= 1 # particles pertaining to this vertex follow immediately # after the vertex for ii in range(nout): for line in inp: - token = str.split(line) - if debug > 0: - print("\t%s" % token) - key = token[0] - if key != 'P': - sys.exit("** hepmc2pandas: faulty event record\n" + line) - - if e['Event_numberP'][-1] < MAXPART: - index = e['Event_numberP'][-1] - e['Event_numberP'][-1] += 1 - - p['Particle_x'].append(x) - p['Particle_y'].append(y) - p['Particle_z'].append(z) - p['Particle_ctau'].append(ctau) - - p['Particle_barcode'].append(int(token[1])) - p['Particle_pid'].append(int(token[2])) - - pid = p['Particle_pid'][-1] - p['Particle_name'].append(particleName(pid)) - p['Particle_d1'].append(-1) - p['Particle_d2'].append(-1) - p['Particle_px'].append(float(token[3])) - p['Particle_py'].append(float(token[4])) - p['Particle_pz'].append(float(token[5])) - p['Particle_energy'].append(float(token[6])) - p['Particle_mass'].append(float(token[7])) - p['Particle_status'].append(int(token[8])) + token = str.split(line) + # if debug > 0: + # print("\t%s" % token) + key = token[0] + # if key != "P": + # sys.exit("** hepmc2pandas: faulty event record\n" + line) + + if e["Event_numberP"][-1] < MAXPART: + index = e["Event_numberP"][-1] + e["Event_numberP"][-1] += 1 + + p["Particle_x"].append(x) + p["Particle_y"].append(y) + p["Particle_z"].append(z) + p["Particle_ctau"].append(ctau) + + p["Particle_barcode"].append(int(token[1])) + p["Particle_pid"].append(int(token[2])) + + pid = p["Particle_pid"][-1] + p["Particle_name"].append(particleName(pid)) + p["Particle_daughters"].append([]) + # p["Particle_d2"].append(-1) + p["Particle_px"].append(float(token[3])) + p["Particle_py"].append(float(token[4])) + p["Particle_pz"].append(float(token[5])) + p["Particle_energy"].append(float(token[6])) + p["Particle_mass"].append(float(token[7])) + p["Particle_status"].append(int(token[8])) self.pvertex[index] = int(token[11]) + self.vertex[vbarcode].append(int(token[1])) - if ii == 0: - self.vertex[vbarcode][0] = index - else: - self.vertex[vbarcode][1] = index + # if ii == 0: + # self.vertex[vbarcode][0] = index + # else: + # self.vertex[vbarcode][1] = index break - else: - return False - - if len(self.vertex) >= e['Event_numberV'][-1]: - for index in range(e['Event_numberP'][-1]): + # else: + # return False + + if len(self.vertex) >= e["Event_numberV"][-1]: + for index in range(e["Event_numberP"][-1]): code = self.pvertex[index] if code in self.vertex: d = self.vertex[code] - - p['Particle_d1'][index] = d[0] - p['Particle_d2'][index] = d[1] + p["Particle_daughters"][index] = d + # p["Particle_d2"][index] = d[1] else: - p['Particle_d1'][index] = -1 - p['Particle_d2'][index] = -1 + p["Particle_daughters"][index] = [] + # p["Particle_d2"][index] = -1 self.plist.append(pd.DataFrame(p)) - + return True else: return False def save(self, filename): print("=> saving to file: %s" % filename) - self.event['Particle'] = self.plist + self.event["Particle"] = self.plist df = pd.DataFrame(self.event) df.to_pickle(filename) print("=> done!") - + def printTable(self): - for ii in range(self.event['Event_numberP'][-1]): + for ii in range(self.event["Event_numberP"][-1]): print("%4d\t%s" % (ii, self.__str__(ii))) -# ----------------------------------------------------------------------- + + +# ----------------------------------------------------------------------- def main(): argv = sys.argv[1:] argc = len(argv) if argc < 1: - sys.exit(''' + sys.exit( + """ Usage: ./hepmc2pandas.py [output pickle file = .pkl] - ''') + """ + ) filename = argv[0] if argc > 1: outfilename = argv[1] else: - outfilename = '%s.pkl' % nameonly(filename) + outfilename = "%s.pkl" % nameonly(filename) stream = hepmc2pandas(filename) @@ -332,9 +395,10 @@ def main(): print(ii) ii += 1 stream.save(outfilename) + + # ----------------------------------------------------------------------- try: main() except KeyboardInterrupt: - print('\nciao!') - + print("\nciao!")