Skip to content
Snippets Groups Projects
Commit f97d6c63 authored by Chao Peng's avatar Chao Peng
Browse files

add particles into the vertex, add random seed control in generator

parent d27758dc
No related branches found
No related tags found
1 merge request!7add particles into the vertex, add random seed control in generator
This diff is collapsed.
import os
from pyHepMC3 import HepMC3 as hm
import numpy as np
import argparse
......@@ -19,7 +20,7 @@ PARTICLES = [
# p in GeV, angle in degree, vertex in mm
def gen_event(prange=(8, 100), arange=(0, 20)):
evt = hm.GenEvent()
evt = hm.GenEvent(momentum_unit=hm.Units.MomentumUnit.GEV, length_unit=hm.Units.LengthUnit.MM)
pid, mass = PARTICLES[np.random.randint(len(PARTICLES))]
# final state
......@@ -35,10 +36,16 @@ def gen_event(prange=(8, 100), arange=(0, 20)):
py = np.sin(phi)*np.sin(theta)
pz = np.cos(theta)
part = hm.GenParticle(hm.FourVector(px*p, py*p, pz*p, e0), pid, state)
# beam
pbeam = hm.GenParticle(hm.FourVector(0, 0, 0, 0.938272), 2212, 4)
ebeam = hm.GenParticle(hm.FourVector(0, 0, e0, np.sqrt(e0*e0 + 0.511e-3*0.511e-3)), -11, 4)
hout = hm.GenParticle(hm.FourVector(px*p, py*p, pz*p, e0), pid, state)
# evt.add_particle(part)
vert = hm.GenVertex()
vert.add_particle_out(part)
vert.add_particle_in(ebeam)
vert.add_particle_in(pbeam)
vert.add_particle_out(hout)
evt.add_vertex(vert)
return evt
......@@ -52,9 +59,15 @@ if __name__ == "__main__":
parser.add_argument('--pmax', type=float, default=100.0, dest='pmax', help='maximum momentum in GeV')
parser.add_argument('--angmin', type=float, default=0.0, dest='angmin', help='minimum angle in degree')
parser.add_argument('--angmax', type=float, default=20.0, dest='angmax', help='maximum angle in degree')
parser.add_argument('-s', type=int, default=-1, dest='seed', help='seed for random generator')
args = parser.parse_args()
# random seed
if args.seed < 0:
args.seed = os.environ.get('SEED', int.from_bytes(os.urandom(4), byteorder='big', signed=False))
np.random.seed(args.seed)
output = hm.WriterAscii(args.output);
if output.failed():
print("Cannot open file \"{}\"".format(args.output))
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment