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         recordElementFractions=False,
0079     )
0080
0081     s.addAlgorithm(g4Alg)
0082
0083     s.addWriter(
0084         RootMaterialTrackWriter(
0085             prePostStep=True,
0086             recalculateTotals=True,
0087             inputMaterialTracks=materialTrackCollectionName,
0088             treeName=materialTrackCollectionName,
0089             filePath=str(outputFileBase) + ".root",
0090             level=acts.logging.INFO,
0091         )
0092     )
0093
0094     return s
0095
0096
0097 def main():
0098     p = argparse.ArgumentParser()
0099     p.add_argument(
0100         "-n", "--events", type=int, default=1000, help="Number of events to generate"
0101     )
0102     p.add_argument(
0103         "-t", "--tracks", type=int, default=100, help="Particle tracks per event"
0104     )
0105     p.add_argument(
0106         "-i", "--input", type=str, default="", help="input (GDML/SQL) file (optional)"
0107     )
0108     p.add_argument(
0109         "--eta-range",
0110         nargs=2,
0111         type=float,
0112         metavar=("MIN", "MAX"),
0113         default=(-4.0, 4.0),
0114         help="Eta range for generated particles",
0115     )
0116     p.add_argument(
0117         "--phi-range",
0118         nargs=2,
0119         type=float,
0120         metavar=("MIN_DEG", "MAX_DEG"),
0121         default=(0.0, 360.0),
0122         help="Phi range in degree for generated particles",
0123     )
0124     p.add_argument(
0125         "--material-track-collection",
0126         type=str,
0127         default="material_tracks",
0128         help="Output material track collection name",
0129     )
0130     p.add_argument(
0131         "-o",
0132         "--output",
0133         type=str,
0134         default="material_geant4",
0135         help="Output file stem (without extension)",
0136     )
0137
0138     args = p.parse_args()
0139
0140     detector = None
0141     if args.input == "":
0142         detector = getOpenDataDetector()
0143     elif args.input.endswith(".gdml"):
0144         detector = acts.examples.geant4.GdmlDetector(path=args.input)
0145     elif args.input.endswith(".sqlite") or args.input.endswith(".db"):
0146         gmdConfig = acts.geomodel.GeoModelDetector.Config(path=args.input)
0147         detector = acts.geomodel.GeoModelDetector(gmdConfig)
0148
0149     runMaterialRecording(
0150         detector=detector,
0151         s=acts.examples.Sequencer(events=args.events, numThreads=1),
0152         tracksPerEvent=args.tracks,
0153         etaRange=tuple(args.eta_range),
0154         phiRange=(args.phi_range[0] * u.degree, args.phi_range[1] * u.degree),
0155         materialTrackCollectionName=args.material_track_collection,
0156         outputFileBase=args.output,
0157     ).run()
0158
0159
0160 if "__main__" == __name__:
0161     main()