Merge branch 'master' into 'master'

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

See merge request EIC/datasets!7

rich_pid.py

+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))