def test_open(evt):
with hep.open("test_read_write_file.dat", "w") as f:
f.write(evt)
with hep.open("test_read_write_file.dat") as f:
evt2 = f.read()
assert evt == evt2
os.unlink("test_read_write_file.dat")
def test_pythonic_read_write():
evt1 = prepare_event()
oss = hep.stringstream()
with hep.open(oss, "w") as f:
f.write(evt1)
evt2 = None
with hep.open(oss) as f:
evt2 = f.read()
assert evt1.particles == evt2.particles
assert evt1.vertices == evt2.vertices
assert evt1 == evt2
def processEvents(self):
print( "Processing {}".format(self.infile) )
# Reads the file
with hep.open(self.infile) as f:
# Just keeps looping
while True :
# Try to get an event
evt = f.read()
if not evt : break # we're at the end of the file
# Stop if we're past max events
if evt.event_number >= self.maxEvents :
break
# Study the mediator and final state charged particles
self.getScalarInfo(evt.particles)
# End looping because we found the scalar
self.saveTrackInfo(evt.particles)
if self.verbose :
print("Event",evt.event_number)
# Mediator info
print("scalar: PtEtaPhiM[{:.2f}, {:.2f}, {:.2f}, {:.2f}]".format(self.scalarPt[-1],self.scalarEta[-1],self.scalarPhi[-1],self.scalarM[-1]) )
# Status=1 particles
print("Some final state charged particles: ")
particles = getStatusOne(evt.particles)
i=0
for part in particles :
if not hasScalarParent(part): continue
if part.momentum.pt() < 0.7: continue
if not isCharged(part): continue
i+=1
trk=part.momentum
print("pdgID {} : PtEtaPhiM[{:.2f}, {:.2f}, {:.2f}, {:.2f}]".format(part.pid,trk.pt(),trk.eta(),trk.phi(),trk.m()))
if i>15: break
# Event level info
print("Number of status=1 particles:",len(particles))
print("Number of charged particles:",self.nCharged[-1])
print("Number of truth tracks:",self.nTracks[-1])
print("Isotropy: {:.2f}".format(self.isotropy[-1] ))
print("HT: {:.1f}".format(self.ht[-1] ))
return
def main(file_name, min_event, max_event, verbose):
print("Reading", file_name, "between", min_event, "and", max_event,
"events")
def print_evt(evt):
def msg(*m):
if verbose:
print(*m)
msg("event_number:", evt.event_number)
msg("Units:", "momentum_unit:", evt.momentum_unit, "length_unit:",
evt.length_unit)
msg(len(evt.particles), "particles:")
for i in enumerate(evt.particles):
pdg = i[1].pid
if pdg in poi:
poi[pdg] = poi[pdg] + 1
pdg = f"{pdg} is of interest!!!"
msg(i, "PDG code", pdg)
msg(len(evt.vertices), "vertices:")
for i in enumerate(evt.vertices):
msg("Vertex:", i)
vertex_pdgs = []
msg("Input particles")
for j in i[1].particles_in:
msg("\t", j, "pdg", j.pid)
msg("Output particles")
for j in i[1].particles_out:
msg("\t", j, "pdg", j.pid)
vertex_pdgs.append(j.pid)
if 2212 in vertex_pdgs and 2112 in vertex_pdgs:
print(evt.event_number, "Has both")
print(i)
for j in i[1].particles_out:
print(j)
with hep.open(file_name) as f:
while True:
e = f.read()
if not e:
break
if e.event_number < min_event:
continue
print_evt(e)
if e.event_number >= max_event:
break
for i in poi:
print("Number of", poi_names[i] + "s", poi[i])
def test_open_with_writer(evt, writer):
filename = "test_open_%s.dat" % writer.__name__
with writer(filename) as f:
f.write(evt)
with hep.open(filename) as f:
evt2 = f.read()
# ReaderHEPEVT adds arbitrary weight to evt2, so we must add that to evt as well
if isinstance(writer, hep.WriterHEPEVT):
evt.run_info = hep.GenRunInfo(weight_names=["0"])
evt.weights = [1]
assert evt.run_info == evt2.run_info
assert evt == evt2
os.unlink(filename)
def eventDisplay(self,event=0,nbins=25,save=False):
#
# Makes SUEP event display
# quantites for plotting
part_pt = np.array([],dtype=float)
part_eta= np.array([],dtype=float)
part_phi= np.array([],dtype=float)
ht=0
# masks for plotting
fromScalar = np.array([],dtype=bool)
fromIsr = np.array([],dtype=bool)
is_mu = np.array([],dtype=bool)
is_e = np.array([],dtype=bool)
is_gam = np.array([],dtype=bool)
is_had = np.array([],dtype=bool)
# Read the file
with hep.open(self.infile) as f:
while True :
evt = f.read()
if not evt : break # we're at the end of the file
if evt.event_number!=event: continue
particles = getStatusOne(evt.particles)
for particle in particles:
# get track info
mom = particle.momentum
#if mom.pt() < 0.1 : continue
if mom.pt() < self.trackPtCut : continue
#if abs(mom.eta()) > self.maxEtaCut : continue
part_pt = np.append(part_pt , mom.pt())
part_phi= np.append(part_phi, mom.phi())
part_eta= np.append(part_eta, mom.eta())
ht+=mom.pt()
# handle masks
parent = hasScalarParent(particle)
fromScalar=np.append(fromScalar,parent)
fromIsr =np.append(fromIsr,not parent)
is_e =np.append(is_e , abs(particle.pid) == 11)
is_mu =np.append(is_mu , abs(particle.pid) == 13)
is_gam=np.append(is_gam, abs(particle.pid) == 22)
is_had=np.append(is_had, abs(particle.pid) > 100)
break # done
# 2D scatter plot
fig, ax = plt.subplots(figsize =(8, 6))
def scale(vec_pt):
s = [50.*pt for pt in vec_pt ]
#s = [2500.*pt/ht for pt in vec_pt ]
return s
def draw(mask,mark='o',col='xkcd:blue'):
ax.scatter(part_eta[mask], part_phi[mask], s=scale(part_pt[mask]),marker=mark,c=col)
return
draw(fromScalar & is_e ,mark='o',col='xkcd:blue')#
draw(fromScalar & is_mu ,mark='v',col='xkcd:blue')#
draw(fromScalar & is_gam,mark='s',col='xkcd:blue')#
draw(fromScalar & is_had,mark='*',col='xkcd:blue')#
draw(fromIsr & is_e ,mark='o',col='xkcd:magenta')#
draw(fromIsr & is_mu ,mark='v',col='xkcd:magenta')#
draw(fromIsr & is_gam,mark='s',col='xkcd:magenta')#
draw(fromIsr & is_had,mark='*',col='xkcd:magenta')#
ax.set_xlabel("$\eta_{particle}$")
ax.set_ylabel("$\phi_{particle}$")
ax.set_xlim(-4.0,4.0)
ax.set_ylim(-3.5,3.5)
# Legend 1 is particle type
line1 = ax.scatter([-100], [-100], label='$e$',marker='o', c='xkcd:black')
line2 = ax.scatter([-100], [-100], label='$\mu$', marker='v', c='xkcd:black')
line3 = ax.scatter([-100], [-100], label='$\gamma$', marker='s', c='xkcd:black')
line4 = ax.scatter([-100], [-100], label='hadron', marker='*', c='xkcd:black')
first_legend = plt.legend(handles=[line1, line2, line3, line4],
loc='upper right', fontsize=12)
ax.add_artist(first_legend)
# Legend 2 is particle origin
blue_patch = mpatches.Patch(color='xkcd:blue', label='from scalar')
red_patch = mpatches.Patch(color='xkcd:magenta', label='not from scalar')
plt.legend(handles=[blue_patch, red_patch],loc='upper left')
if (save) : plt.savefig("eventdisplay_{}_{}.png".format(event,self.infile.strip(".hepmc")))
else : plt.show()
return