modify gen_particles.py

gen_particles.py

@@ -5,22 +5,22 @@ import argparse
 
 
 PARTICLES = {
-    111: 0.1349766,       # pi0
-    211: 0.13957018,      # pi+
-    -211: 0.13957018,     # pi-
-    311: 0.497648,        # K0
-    321: 0.493677,        # K+
-    -321: 0.493677,       # K-
-    2212: 0.938272,       # proton
-    2112: 0.939565,       # neutron
-    11: 0.51099895e-3,    # electron
-    -11: 0.51099895e-3,   # positron
-    22: 0,                # photon
+    "pion0": (111, 0.1349766),       # pi0
+    "pion+": (211, 0.13957018),      # pi+
+    "pion-": (-211, 0.13957018),     # pi-
+    "kaon0": (311, 0.497648),        # K0
+    "kaon+": (321, 0.493677),        # K+
+    "kaon-": (-321, 0.493677),       # K-
+    "proton": (2212, 0.938272),      # proton
+    "neutron": (2112, 0.939565),     # neutron
+    "electron": (11, 0.51099895e-3), # electron
+    "positron": (-11, 0.51099895e-3),# positron
+    "photon": (22, 0),               # photon
 }
 
 
 # p in GeV, angle in degree, vertex in mm
-def gen_event(prange=(8, 100), arange=(0, 20), phrange=(0., 360.), parts=[(p, m) for p, m in PARTICLES.items()]):
+def gen_event(prange=(8, 100), arange=(0, 20), phrange=(0., 360.), parts=[(p, m) for p, m in PARTICLES.values()]):
     evt = hm.GenEvent(momentum_unit=hm.Units.MomentumUnit.GEV, length_unit=hm.Units.LengthUnit.MM)
     pid, mass = parts[np.random.randint(len(parts))]
 
@@ -39,7 +39,7 @@ def gen_event(prange=(8, 100), arange=(0, 20), phrange=(0., 360.), parts=[(p, m)
 
     # 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)
+    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)
@@ -61,8 +61,8 @@ 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('--phmin', type=float, default=0.0, dest='angmin', help='minimum angle in degree')
-    parser.add_argument('--phmax', type=float, default=360.0, dest='angmax', help='maximum angle in degree')
+    parser.add_argument('--phmin', type=float, default=0.0, dest='phmin', help='minimum angle in degree')
+    parser.add_argument('--phmax', type=float, default=360.0, dest='phmax', help='maximum angle in degree')
     parser.add_argument('--particles', type=str, default='11', dest='particles', help='particles pdg code')
     parser.add_argument('-s', type=int, default=-1, dest='seed', help='seed for random generator')
 
@@ -80,11 +80,11 @@ if __name__ == "__main__":
 
     parts = []
     for pid in args.particles.split(','):
-        pid = int(pid.strip())
+        pid = pid.strip()
         if pid not in PARTICLES.keys():
             print('pid {:d} not found in dictionary, ignored.'.format(pid))
             continue
-        parts.append((pid, PARTICLES[pid]))
+        parts.append(PARTICLES[pid])
 
 
     if not args.parray: