Scripts/Python/material_recording.py

0001 #!/usr/bin/env python3
0002
0003 import os
0004 import warnings
0005 import argparse
0006
0007 import acts
0008 from acts.examples import (
0009     GaussianVertexGenerator,
0010     ParametricParticleGenerator,
0011     FixedMultiplicityGenerator,
0012     EventGenerator,
0013     RandomNumbers,
0014 )
0015
0016 import acts.examples.dd4hep
0017 import acts.examples.geant4
0018 from acts.examples.odd import getOpenDataDetector
0019
0020 try:
0021     import acts.examples.geant4.geomodel
0022 except ImportError:
0023     # geomodel is optional for this script
0024     pass
0025
0026 u = acts.UnitConstants
0027
0028 _material_recording_executed = False
0029
0030
0031 def runMaterialRecording(
0032     detector,
0033     outputDir,
0034     tracksPerEvent=10000,
0035     s=None,
0036     etaRange=(-4, 4),
0037 ):
0038     global _material_recording_executed
0039     if _material_recording_executed:
0040         warnings.warn("Material recording already ran in this process. Expect crashes")
0041     _material_recording_executed = True
0042
0043     rnd = RandomNumbers(seed=228)
0044
0045     evGen = EventGenerator(
0046         level=acts.logging.INFO,
0047         generators=[
0048             EventGenerator.Generator(
0049                 multiplicity=FixedMultiplicityGenerator(n=1),
0050                 vertex=GaussianVertexGenerator(
0051                     stddev=acts.Vector4(0, 0, 0, 0),
0052                     mean=acts.Vector4(0, 0, 0, 0),
0053                 ),
0054                 particles=ParametricParticleGenerator(
0055                     pdg=acts.PdgParticle.eInvalid,
0056                     charge=0,
0057                     randomizeCharge=False,
0058                     mass=0,
0059                     p=(1 * u.GeV, 10 * u.GeV),
0060                     eta=etaRange,
0061                     numParticles=tracksPerEvent,
0062                     etaUniform=True,
0063                 ),
0064             )
0065         ],
0066         outputParticles="particles_initial",
0067         outputVertices="vertices_initial",
0068         randomNumbers=rnd,
0069     )
0070
0071     s.addReader(evGen)
0072
0073     g4Alg = acts.examples.geant4.Geant4MaterialRecording(
0074         level=acts.logging.INFO,
0075         detector=detector,
0076         randomNumbers=rnd,
0077         inputParticles=evGen.config.outputParticles,
0078         outputMaterialTracks="material-tracks",
0079     )
0080
0081     s.addAlgorithm(g4Alg)
0082
0083     s.addWriter(
0084         acts.examples.RootMaterialTrackWriter(
0085             prePostStep=True,
0086             recalculateTotals=True,
0087             inputMaterialTracks="material-tracks",
0088             filePath=os.path.join(outputDir, "geant4_material_tracks.root"),
0089             level=acts.logging.INFO,
0090         )
0091     )
0092
0093     return s
0094
0095
0096 def main():
0097     p = argparse.ArgumentParser()
0098     p.add_argument(
0099         "-n", "--events", type=int, default=1000, help="Number of events to generate"
0100     )
0101     p.add_argument(
0102         "-t", "--tracks", type=int, default=100, help="Particle tracks per event"
0103     )
0104     p.add_argument(
0105         "-i", "--input", type=str, default="", help="input (GDML/SQL) file (optional)"
0106     )
0107
0108     args = p.parse_args()
0109
0110     detector = None
0111     if args.input == "":
0112         detector = getOpenDataDetector()
0113     elif args.input.endswith(".gdml"):
0114         detector = acts.examples.geant4.GdmlDetector(path=args.input)
0115     elif args.input.endswith(".sqlite") or args.input.endswith(".db"):
0116         detector = acts.examples.GeoModelDetector(path=args.input)
0117
0118     runMaterialRecording(
0119         detector=detector,
0120         tracksPerEvent=args.tracks,
0121         outputDir=os.getcwd(),
0122         s=acts.examples.Sequencer(events=args.events, numThreads=1),
0123     ).run()
0124
0125
0126 if "__main__" == __name__:
0127     main()