def write(self, data: vectors.ParticlePool):
# I like f-strings, but this is the fastest way to make a string
self.__file_handle.write("%i\n" % data.particle_count)
for p in data.iter_particles():
self.__file_handle.write(
"%d %d %.20f %.20f %.20f %.20f\n" %
(p.id, p.charge, p.x[0], p.y[0], p.z[0], p.e[0]))
def get_event_mass(collection: vectors.ParticlePool) -> npy.ndarray:
found_photon, found_proton = 0, 0
vector_sum = vectors.FourVector(collection.event_count)
for event_particle in collection.iter_particles():
if not found_photon and event_particle.id == 1:
found_photon = True
elif not found_proton and event_particle.id == 14:
found_proton = True
else:
vector_sum += event_particle
return vector_sum.get_mass()