rich/scripts/rich_data_gen.py

0001 import os
0002 from pyHepMC3 import HepMC3 as hm
0003 import numpy as np
0004 import argparse
0005
0006
0007 PARTICLES = [
0008 #    (111, 0.134.9766),      # pi0
0009     (211, 0.13957018),      # pi+
0010 #    (-211, 0.13957018),     # pi-
0011 #    (311, 0.497648),        # K0
0012     (321, 0.493677),        # K+
0013 #    (-321, 0.493677),       # K-
0014     (2212, 0.938272),       # proton
0015 #    (2112, 0.939565),       # neutron
0016 #    (11, 0.51099895e-3),    # electron
0017 #    (-11, 0.51099895e-3),   # positron
0018 ]
0019
0020
0021 # p in GeV, angle in degree, vertex in mm
0022 def gen_event(prange=(8, 100), arange=(0, 20)):
0023     evt = hm.GenEvent(hm.Units.MomentumUnit.GEV, hm.Units.LengthUnit.MM)
0024     pid, mass = PARTICLES[np.random.randint(len(PARTICLES))]
0025
0026     # final state
0027     state = 1
0028
0029     # momentum, angles, energy
0030     p = np.random.uniform(*prange)
0031     theta = np.random.uniform(*arange)*np.pi/180.
0032     phi = np.random.uniform(0., 2.*np.pi)
0033     e0 = np.sqrt(p*p + mass*mass)
0034
0035     px = np.cos(phi)*np.sin(theta)
0036     py = np.sin(phi)*np.sin(theta)
0037     pz = np.cos(theta)
0038
0039     # beam
0040     pbeam = hm.GenParticle(hm.FourVector(0, 0, 0, 0.938272), 2212, 4)
0041     ebeam = hm.GenParticle(hm.FourVector(0, 0, e0, np.sqrt(e0*e0 + 0.511e-3*0.511e-3)), -11, 4)
0042
0043     hout = hm.GenParticle(hm.FourVector(px*p, py*p, pz*p, e0), pid, state)
0044     # evt.add_particle(part)
0045     vert = hm.GenVertex()
0046     vert.add_particle_in(ebeam)
0047     vert.add_particle_in(pbeam)
0048     vert.add_particle_out(hout)
0049     evt.add_vertex(vert)
0050     return evt
0051
0052
0053 if __name__ == "__main__":
0054     parser = argparse.ArgumentParser('RICH dataset generator')
0055
0056     parser.add_argument('output', help='path to the output file')
0057     parser.add_argument('-n', type=int, default=1000, dest='nev', help='number of events to generate')
0058     parser.add_argument('--pmin', type=float, default=8.0, dest='pmin', help='minimum momentum in GeV')
0059     parser.add_argument('--pmax', type=float, default=100.0, dest='pmax', help='maximum momentum in GeV')
0060     parser.add_argument('--angmin', type=float, default=0.0, dest='angmin', help='minimum angle in degree')
0061     parser.add_argument('--angmax', type=float, default=20.0, dest='angmax', help='maximum angle in degree')
0062     parser.add_argument('-s', type=int, default=-1, dest='seed', help='seed for random generator')
0063
0064     args = parser.parse_args()
0065
0066     # random seed
0067     if args.seed < 0:
0068         args.seed = os.environ.get('SEED', int.from_bytes(os.urandom(4), byteorder='big', signed=False))
0069     np.random.seed(args.seed)
0070
0071     output = hm.WriterAscii(args.output);
0072     if output.failed():
0073         print("Cannot open file \"{}\"".format(args.output))
0074         sys.exit(2)
0075
0076     count = 0
0077     while count < args.nev:
0078         if (count % 1000 == 0):
0079             print("Generated {} events".format(count), end='\r')
0080         evt = gen_event((args.pmin, args.pmax), (args.angmin, args.angmax))
0081         output.write_event(evt)
0082         evt.clear()
0083         count += 1
0084
0085     print("Generated {} events".format(args.nev))
0086     output.close()
0087