Scripts/Python/material_recording.py

0001 #!/usr/bin/env python3
0002
0003 import argparse
0004 from pathlib import Path
0005
0006 import acts
0007 from acts.examples import (
0008     GaussianVertexGenerator,
0009     ParametricParticleGenerator,
0010     FixedMultiplicityGenerator,
0011     EventGenerator,
0012     RandomNumbers,
0013 )
0014
0015 import acts.examples.dd4hep
0016 import acts.examples.geant4
0017 import acts.examples.hepmc3
0018 from acts.examples.odd import getOpenDataDetector
0019 from acts.examples.root import RootMaterialTrackWriter
0020
0021 u = acts.UnitConstants
0022
0023
0024 def runMaterialRecording(
0025     detector,
0026     s,
0027     tracksPerEvent=10000,
0028     etaRange=(-4.0, 4.0),
0029     phiRange=(0.0, 360.0 * u.degree),
0030     materialTrackCollectionName="material_tracks",
0031     outputFileBase="geant4_material_tracks",
0032 ):
0033
0034     rnd = RandomNumbers(seed=228)
0035
0036     evGen = EventGenerator(
0037         level=acts.logging.INFO,
0038         generators=[
0039             EventGenerator.Generator(
0040                 multiplicity=FixedMultiplicityGenerator(n=1),
0041                 vertex=GaussianVertexGenerator(
0042                     stddev=acts.Vector4(0, 0, 0, 0),
0043                     mean=acts.Vector4(0, 0, 0, 0),
0044                 ),
0045                 particles=ParametricParticleGenerator(
0046                     pdg=acts.PdgParticle.eInvalid,
0047                     charge=0,
0048                     randomizeCharge=False,
0049                     mass=0,
0050                     p=(1 * u.GeV, 10 * u.GeV),
0051                     eta=etaRange,
0052                     phi=phiRange,
0053                     numParticles=tracksPerEvent,
0054                     etaUniform=True,
0055                 ),
0056             )
0057         ],
0058         randomNumbers=rnd,
0059     )
0060
0061     s.addReader(evGen)
0062
0063     hepmc3Converter = acts.examples.hepmc3.HepMC3InputConverter(
0064         level=acts.logging.INFO,
0065         inputEvent=evGen.config.outputEvent,
0066         outputParticles="particles_initial",
0067         outputVertices="vertices_initial",
0068         mergePrimaries=False,
0069     )
0070     s.addAlgorithm(hepmc3Converter)
0071
0072     g4Alg = acts.examples.geant4.Geant4MaterialRecording(
0073         level=acts.logging.INFO,
0074         detector=detector,
0075         randomNumbers=rnd,
0076         inputParticles=hepmc3Converter.config.outputParticles,
0077         outputMaterialTracks=materialTrackCollectionName,
0078     )
0079
0080     s.addAlgorithm(g4Alg)
0081
0082     s.addWriter(
0083         RootMaterialTrackWriter(
0084             prePostStep=True,
0085             recalculateTotals=True,
0086             inputMaterialTracks=materialTrackCollectionName,
0087             treeName=materialTrackCollectionName,
0088             filePath=str(outputFileBase) + ".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     p.add_argument(
0108         "--eta-range",
0109         nargs=2,
0110         type=float,
0111         metavar=("MIN", "MAX"),
0112         default=(-4.0, 4.0),
0113         help="Eta range for generated particles",
0114     )
0115     p.add_argument(
0116         "--phi-range",
0117         nargs=2,
0118         type=float,
0119         metavar=("MIN_DEG", "MAX_DEG"),
0120         default=(0.0, 360.0),
0121         help="Phi range in degree for generated particles",
0122     )
0123     p.add_argument(
0124         "--material-track-collection",
0125         type=str,
0126         default="material_tracks",
0127         help="Output material track collection name",
0128     )
0129     p.add_argument(
0130         "-o",
0131         "--output",
0132         type=str,
0133         default="material_geant4",
0134         help="Output file stem (without extension)",
0135     )
0136
0137     args = p.parse_args()
0138
0139     detector = None
0140     if args.input == "":
0141         detector = getOpenDataDetector()
0142     elif args.input.endswith(".gdml"):
0143         detector = acts.examples.geant4.GdmlDetector(path=args.input)
0144     elif args.input.endswith(".sqlite") or args.input.endswith(".db"):
0145         gmdConfig = acts.geomodel.GeoModelDetector.Config(path=args.input)
0146         detector = acts.geomodel.GeoModelDetector(gmdConfig)
0147
0148     runMaterialRecording(
0149         detector=detector,
0150         s=acts.examples.Sequencer(events=args.events, numThreads=1),
0151         tracksPerEvent=args.tracks,
0152         etaRange=tuple(args.eta_range),
0153         phiRange=(args.phi_range[0] * u.degree, args.phi_range[1] * u.degree),
0154         materialTrackCollectionName=args.material_track_collection,
0155         outputFileBase=args.output,
0156     ).run()
0157
0158
0159 if "__main__" == __name__:
0160     main()