Scripts/Python/event_recording.py

0001 #!/usr/bin/env python3
0002
0003 import os
0004
0005 import acts
0006 import acts.examples
0007 import acts.examples.hepmc3
0008 import acts.examples.dd4hep
0009 import acts.examples.geant4
0010 import acts.examples.geant4.hepmc3
0011 from acts.examples.odd import getOpenDataDetector
0012
0013
0014 u = acts.UnitConstants
0015
0016
0017 def runEventRecording(detector, outputDir, s=None):
0018     hepmc_dir = os.path.join(outputDir, "hepmc3")
0019     if not os.path.exists(hepmc_dir):
0020         os.mkdir(hepmc_dir)
0021
0022     s = s or acts.examples.Sequencer(
0023         events=int(os.environ.get("NEVENTS", 100)), numThreads=1
0024     )
0025
0026     rnd = acts.examples.RandomNumbers(seed=42)
0027     evGen = acts.examples.EventGenerator(
0028         level=acts.logging.INFO,
0029         generators=[
0030             acts.examples.EventGenerator.Generator(
0031                 multiplicity=acts.examples.FixedMultiplicityGenerator(n=2),
0032                 vertex=acts.examples.GaussianVertexGenerator(
0033                     stddev=acts.Vector4(0, 0, 0, 0), mean=acts.Vector4(0, 0, 0, 0)
0034                 ),
0035                 particles=acts.examples.ParametricParticleGenerator(
0036                     p=(1 * u.GeV, 10 * u.GeV),
0037                     eta=(-2, 2),
0038                     phi=(0, 90 * u.degree),
0039                     randomizeCharge=True,
0040                     numParticles=4,
0041                 ),
0042             )
0043         ],
0044         outputParticles="particles_input",
0045         outputVertices="vertices_input",
0046         randomNumbers=rnd,
0047     )
0048
0049     s.addReader(evGen)
0050
0051     erAlgCfg = acts.examples.geant4.hepmc3.EventRecording.Config(
0052         inputParticles=evGen.config.outputParticles,
0053         outputHepMcTracks="geant-event",
0054         seed1=43,
0055         seed2=44,
0056         detector=detector,
0057     )
0058
0059     erAlg = acts.examples.geant4.hepmc3.EventRecording(
0060         config=erAlgCfg, level=acts.logging.INFO
0061     )
0062
0063     s.addAlgorithm(erAlg)
0064
0065     s.addWriter(
0066         acts.examples.hepmc3.HepMC3AsciiWriter(
0067             level=acts.logging.INFO,
0068             outputDir=hepmc_dir,
0069             outputStem="events",
0070             inputEvents=erAlg.config.outputHepMcTracks,
0071         )
0072     )
0073
0074     return s
0075
0076
0077 if "__main__" == __name__:
0078     detector = getOpenDataDetector()
0079
0080     runEventRecording(
0081         detector=detector,
0082         outputDir=os.getcwd(),
0083     ).run()