add scripts for the test

scripts/subdetector_tests/ff_det_names.json

@@ -0,0 +1,31 @@
+{
+  "B0APF_BeamlineMagnet": "Beamline (ion)",
+  "B0ECal": "B0 Detector",
+  "B0PF_BeamlineMagnet": "Beamline (ion)",
+  "B0Tracker": "B0 Detector",
+  "B0TrackerCompanion": "B0 Detector",
+  "B1APF_BeamlineMagnet": "Beamline (ion)",
+  "B1PF_BeamlineMagnet": "Beamline (ion)",
+  "B2PF_BeamlineMagnet": "Beamline (ion)",
+  "BeamPipeB0": "B0 Detector",
+  "ForwardOffMTracker_station_1": "Off-M Tracker",
+  "ForwardOffMTracker_station_2": "Off-M Tracker",
+  "ForwardOffMTracker_station_3": "Off-M Tracker",
+  "ForwardOffMTracker_station_4": "Off-M Tracker",
+  "ForwardRomanPot_Station_1": "Roman Pot",
+  "ForwardRomanPot_Station_2": "Roman Pot",
+  "Pipe_cen_to_pos": "Beamline ($e^-$)",
+  "Q0EF": "Beamline ($e^-$)",
+  "Q0EF_vac": "Beamline ($e^-$)",
+  "Q1APF_BeamlineMagnet": "Beamline (ion)",
+  "Q1BPF_BeamlineMagnet": "Beamline (ion)",
+  "Q1EF": "Beamline ($e^-$)",
+  "Q1EF_vac": "Beamline ($e^-$)",
+  "Q2PF_BeamlineMagnet": "Beamline (ion)",
+  "VacuumMagnetElement": "Vacuum Magnet Element",
+  "ZDC_1stSilicon": "ZDC",
+  "ZDC_Crystal": "ZDC",
+  "ZDC_PbSci": "ZDC",
+  "ZDC_PbSi": "ZDC",
+  "ZDC_WSi": "ZDC"
+}

scripts/subdetector_tests/ff_performance_test.py

@@ -0,0 +1,219 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+# Copyright (C) 2023 Chao Peng
+'''
+    A python script to benchmark the performance of far forward detectors with high energy proton/nuclei
+    It tests each far forward detector (see FF_MOTHER and FF_COMP) and reads the process time at the end of the simulation
+
+    Author: Chao Peng (ANL)
+    Date: 07/19/2023
+'''
+import os
+import re
+import sys
+import subprocess
+import argparse
+import shutil
+import numpy as np
+from matplotlib import pyplot as plt
+from matplotlib import colors as mcolors
+
+
+# script directory
+SDIR = os.path.dirname(os.path.realpath(__file__))
+# event generation script (particle gun is not working in this method, subprocess.run won't go through)
+GEN_SCRIPT = os.path.join(SDIR, 'gen_particles.py')
+
+
+def mrad2deg(th):
+    return th/1000./np.pi*180.
+
+
+def deg2mrad(ang):
+    return ang/180.*np.pi*1000.
+
+
+# the ff xml file that is used by the top-level xml file
+# this script will modify the file so try to use a copy of the original file
+FF_MOTHER = 'compact/far_forward_test.xml'
+# ff components
+# this list will be tested one-by-one (by modifying the mother xml file
+FF_COMP = np.array([
+        ('Beamline (Ion)', 'far_forward/ion_beamline.xml'),
+        ('Beamline ($e^-$)', 'far_forward/electron_beamline.xml'),
+        ('B0 Beampipe', 'far_forward/beampipe_hadron_B0.xml'),
+        ('B0 Tracker', 'far_forward/B0_tracker.xml'),
+        ('B0 ECal', 'far_forward/B0_ECal.xml'),
+        ('Off-M Tracker', 'offM_tracker.xml'),
+        ('ZDC', 'far_forward/ZDC.xml'),
+        ('Roman Pots', 'far_forward/roman_pots_eRD24_design.xml'),
+])
+
+
+# kwargs are used by gen_cmd and sim_cmd
+def sim_performance_test(**kwargs):
+    gen_file = 'ff_test_gen.hepmc'
+    # generate particles
+    gen_cmd = [
+        'python', GEN_SCRIPT, gen_file,
+        '-n={nev}',
+        '--angmin={angle_min}',
+        '--angmax={angle_max}',
+        '--pmin={p_min}',
+        '--pmax={p_max}',
+        '--phmin={phi_min}',
+        '--phmax={phi_max}',
+        '--particles={particles}',
+        ]
+    gen_cmd = [c.format(**kwargs) for c in gen_cmd]
+    print(' '.join(gen_cmd))
+    subprocess.run(gen_cmd, check=True)
+
+    # simulation
+    sim_cmd = [
+        'ddsim',
+        '--runType=batch',
+        '--part.minimalKineticEnergy=1*TeV',
+        '--filter.tracker=edep0',
+        # '-v=WARNING',
+        '--numberOfEvents={nev}',
+        # '--physics.list {physics_list}',
+        # '--enableGun',
+        # '--gun.particle=proton', '--gun.energy=275*GeV',
+        # '--gun.position=\"0.0 0.0 0.0*cm\"',
+        # '--gun.phiMin=0.', '--gun.phiMax=2.*pi',
+        # '--gun.thetaMin=0.003', '--gun.thetaMax=0.004',
+        # '--gun.distribution=uniform',
+        '--inputFiles={}'.format(gen_file),
+        '--outputFile={sim_file}',
+        '--compact={compact}',
+        ]
+    sim_cmd = [c.format(**kwargs) for c in sim_cmd]
+    print(' '.join(sim_cmd))
+    p = subprocess.run(sim_cmd, stdout=subprocess.PIPE)
+    lines = p.stdout.decode('utf-8').split('\n')[-4:]
+
+    pat = re.compile('Event Processing:\s+([-+]?(?:\d*\.*\d+))\s+s')
+    r = pat.search('\n'.join(lines), re.MULTILINE)
+    if not r:
+        print('Cannot find event processing time in:')
+        print('\n'.join(lines))
+        return None
+    return float(r.group(1))
+
+
+if __name__ == '__main__':
+    # argument parser
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('compact',
+            help='A Top-level XML file of the detector discription.'
+            )
+    parser.add_argument('--ff-compact',
+            default=FF_MOTHER,
+            help='The XML file for far-forward detectors, used by the top-level XML file. It will be modified by the script.'
+            )
+    parser.add_argument('--nev',
+            default=100, type=int,
+            help='Number of events.'
+            )
+    parser.add_argument('--particles',
+            default='proton',
+            help='Type of particles.'
+            )
+    parser.add_argument('--sim-file',
+            default='ff_test.edm4hep.root',
+            help='Temporary output root file.'
+            )
+    parser.add_argument('--theta-min',
+            default=3, type=float,
+            help='Min. polar angle (mrad).'
+            )
+    parser.add_argument('--theta-max',
+            default=3.5, type=float,
+            help='Max. polar angle (mrad).'
+            )
+    parser.add_argument('--p-min',
+            default=275, type=float,
+            help='Min. momentum of the particles (GeV).'
+            )
+    parser.add_argument('--p-max',
+            default=275, type=float,
+            help='Max. momentum of the particles (GeV).'
+            )
+    parser.add_argument('--phi-min',
+            default=0, type=float,
+            help='Min. phi angle of the particles (degree).'
+            )
+    parser.add_argument('--phi-max',
+            default=360, type=float,
+            help='Max. phi angle of the particles (degree).'
+            )
+
+    args = parser.parse_args()
+    kwargs = vars(args)
+    # convert mrad to angle
+    kwargs['angle_min'] = mrad2deg(args.theta_min)
+    kwargs['angle_max'] = mrad2deg(args.theta_max)
+
+    # all ff components
+    t = sim_performance_test(**kwargs)
+    result = [('All FF', t)]
+
+    # read ff-compact
+    file1 = open(args.ff_compact, 'r')
+    lines = file1.readlines()
+    file1.close()
+
+    for comp1, xml1 in FF_COMP:
+        print('Running test for ff component: {}'.format(comp1))
+        new_lines = []
+        for nl in lines:
+            is_needed = True
+            for comp2, xml2 in FF_COMP:
+                if xml2 in nl and comp2 != comp1:
+                    is_needed = False
+            if is_needed:
+                new_lines.append(nl)
+        file2 = open(args.ff_compact, 'w')
+        file2.writelines(new_lines)
+        file2.close()
+
+        # print(comp1)
+        # print(''.join(lines))
+        # print(''.join(new_lines))
+
+        t = sim_performance_test(**vars(args))
+        result.append((comp1, t))
+        print('{}: {:.2f} s for {:d} events'.format(comp1, t, 100))
+
+    # restore it back
+    file2 = open(args.ff_compact, 'w')
+    file2.writelines(lines)
+    file2.close()
+
+
+    # Build the plot
+    result = np.array(result)
+    p = (args.p_min + args.p_max)/2.    # GeV
+    theta = (args.theta_min + args.theta_max)/2.    # mrad
+    nev = float(args.nev)
+
+    prop_cycle = plt.rcParams['axes.prop_cycle']
+    colors = prop_cycle.by_key()['color']
+
+    fig, ax = plt.subplots(figsize=(8, 6), dpi=160,  gridspec_kw=dict(top=0.95, bottom=0.25, left=0.15, right=0.98))
+    x_pos = np.arange(len(result))
+    for i, (x, d, e) in enumerate(zip(x_pos, result.T[1].astype(float)/nev, result.T[1].astype(float)/np.sqrt(nev)/nev)):
+        ic = colors[i % len(colors)]
+        ax.bar(x, d, yerr=e, align='center', color=mcolors.to_rgba(ic, alpha=0.5), ec=ic, ecolor='black', capsize=10)
+    ax.set_ylabel('Process Time (s / event)', fontsize=16)
+    ax.set_xticks(x_pos)
+    ax.set_xticklabels(result.T[0], rotation=45, ha='right', fontsize=14)
+    ax.set_title('{:.0f} GeV/c {} @ ~{:.1f} mrad'.format(p, args.particles, theta), fontsize=16)
+    ax.yaxis.grid(ls=':')
+    ax.tick_params(labelsize=16)
+    ax.set_axisbelow(True)
+
+    # Save the figure
+    # fig.tight_layout()
+    fig.savefig('ff_test_{:.0f}_{}_{:.0f}_mrad.png'.format(p, args.particles, theta))

scripts/subdetector_tests/gen_particles.py

@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+# Copyright (C) 2023 Chao Peng
+'''
+    A simple script to generate single particles in HEPMC3 format
+
+    Author: Chao Peng (ANL)
+    Date: 07/19/2023
+'''
+import os
+import sys
+from pyHepMC3 import HepMC3 as hm
+import numpy as np
+import argparse
+
+
+PARTICLES = {
+    "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
+    "muon": (13, 105.6583755),       # muon
+}
+
+
+def gen_event(p, theta, phi, pid, mass):
+    evt = hm.GenEvent(hm.Units.MomentumUnit.GEV, hm.Units.LengthUnit.MM)
+    # final state
+    state = 1
+    e0 = np.sqrt(p*p + mass*mass)
+    px = np.cos(phi)*np.sin(theta)
+    py = np.sin(phi)*np.sin(theta)
+    pz = np.cos(theta)
+
+    # 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)
+
+    # out particle
+    hout = hm.GenParticle(hm.FourVector(px*p, py*p, pz*p, e0), pid, state)
+
+    # vertex
+    vert = hm.GenVertex()
+    vert.add_particle_in(ebeam)
+    vert.add_particle_in(pbeam)
+    vert.add_particle_out(hout)
+    evt.add_vertex(vert)
+    return evt
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('output', help='path to the output file')
+    parser.add_argument('-n', type=int, default=1000, dest='nev', help='number of events to generate')
+    parser.add_argument('-s', type=int, default=-1, dest='seed', help='seed for random generator')
+    parser.add_argument('--parray', type=str, default="", dest='parray',
+                        help='an array of momenta in GeV, separated by \",\"')
+    parser.add_argument('--pmin', type=float, default=8.0, dest='pmin', help='minimum momentum in GeV')
+    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='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='electron', dest='particles',
+                        help='particle names, support {}'.format(list(PARTICLES.keys())))
+
+    args = parser.parse_args()
+
+    # random seed (< 0 will get it from enviroment variable 'SEED', or a system random number)
+    if args.seed < 0:
+        args.seed = os.environ.get('SEED', int.from_bytes(os.urandom(4), byteorder='big', signed=False))
+    print("Random seed is {}".format(args.seed))
+    np.random.seed(args.seed)
+
+    output = hm.WriterAscii(args.output);
+    if output.failed():
+        print("Cannot open file \"{}\"".format(args.output))
+        sys.exit(2)
+
+    # build particle info
+    parts = []
+    for pid in args.particles.split(','):
+        pid = pid.strip()
+        if pid not in PARTICLES.keys():
+            print('pid {:d} not found in dictionary, ignored.'.format(pid))
+            continue
+        parts.append(PARTICLES[pid])
+
+    # p values
+    pvals = np.random.uniform(args.pmin, args.pmax, args.nev) if not args.parray else \
+            np.random.choice([float(p.strip()) for p in args.parray.split(',')], args.nev)
+    thvals = np.random.uniform(args.angmin, args.angmax, args.nev)/180.*np.pi
+    phivals = np.random.uniform(args.phmin, args.phmax, args.nev)/180.*np.pi
+    partvals = [parts[i] for i in np.random.choice(len(parts), args.nev)]
+
+    count = 0
+    for p, theta, phi, (pid, mass) in zip(pvals, thvals, phivals, partvals):
+        if (count % 1000 == 0):
+            print("Generated {} events".format(count), end='\r')
+        evt = gen_event(p, theta, phi, pid, mass)
+        output.write_event(evt)
+        evt.clear()
+        count += 1
+
+    print("Generated {} events".format(args.nev))
+    output.close()