Scripts/Python/material_validation.py

0001 #!/usr/bin/env python3
0002
0003 import os
0004 import argparse
0005
0006 import acts
0007 from acts.examples import Sequencer, RootMaterialTrackWriter
0008 from acts.examples.odd import getOpenDataDetector
0009 from acts.examples.simulation import addParticleGun, EtaConfig, ParticleConfig
0010
0011
0012 def runMaterialValidation(
0013     nevents,
0014     ntracks,
0015     trackingGeometry,
0016     decorators,
0017     field,
0018     outputDir,
0019     outputName="propagation-material",
0020     s=None,
0021 ):
0022     # Create a sequencer
0023     s = s or Sequencer(events=nevents, numThreads=-1)
0024
0025     for decorator in decorators:
0026         s.addContextDecorator(decorator)
0027
0028     nav = acts.Navigator(trackingGeometry=trackingGeometry)
0029     stepper = acts.StraightLineStepper()
0030     # stepper = acts.EigenStepper(field)
0031
0032     prop = acts.examples.ConcretePropagator(acts.Propagator(stepper, nav))
0033
0034     rnd = acts.examples.RandomNumbers(seed=42)
0035
0036     addParticleGun(
0037         s,
0038         ParticleConfig(num=ntracks, pdg=acts.PdgParticle.eMuon, randomizeCharge=True),
0039         EtaConfig(-4.0, 4.0),
0040         rnd=rnd,
0041     )
0042
0043     trkParamExtractor = acts.examples.ParticleTrackParamExtractor(
0044         level=acts.logging.INFO,
0045         inputParticles="particles_input",
0046         outputTrackParameters="params_particles_input",
0047     )
0048     s.addAlgorithm(trkParamExtractor)
0049
0050     alg = acts.examples.PropagationAlgorithm(
0051         propagatorImpl=prop,
0052         level=acts.logging.INFO,
0053         sterileLogger=True,
0054         recordMaterialInteractions=True,
0055         inputTrackParameters="params_particles_input",
0056         outputSummaryCollection="propagation_summary",
0057         outputMaterialCollection="material_tracks",
0058     )
0059     s.addAlgorithm(alg)
0060
0061     s.addWriter(
0062         RootMaterialTrackWriter(
0063             level=acts.logging.INFO,
0064             inputMaterialTracks=alg.config.outputMaterialCollection,
0065             filePath=os.path.join(outputDir, (outputName + ".root")),
0066             storeSurface=True,
0067             storeVolume=True,
0068         )
0069     )
0070
0071     return s
0072
0073
0074 if "__main__" == __name__:
0075     p = argparse.ArgumentParser()
0076
0077     p.add_argument(
0078         "-n", "--events", type=int, default=1000, help="Number of events to process"
0079     )
0080     p.add_argument(
0081         "-t", "--tracks", type=int, default=1000, help="Number of tracks per event"
0082     )
0083     p.add_argument(
0084         "-m", "--map", type=str, help="Input file (optional) for the material map"
0085     )
0086     p.add_argument(
0087         "-o",
0088         "--output",
0089         type=str,
0090         default="propagation-material",
0091         help="Output file name",
0092     )
0093
0094     args = p.parse_args()
0095
0096     materialDecorator = (
0097         acts.IMaterialDecorator.fromFile(args.map) if args.map != None else None
0098     )
0099
0100     detector = getOpenDataDetector(materialDecorator)
0101     trackingGeometry = detector.trackingGeometry()
0102     decorators = detector.contextDecorators()
0103
0104     field = acts.ConstantBField(acts.Vector3(0, 0, 0 * acts.UnitConstants.T))
0105
0106     runMaterialValidation(
0107         args.events,
0108         args.tracks,
0109         trackingGeometry,
0110         decorators,
0111         field,
0112         outputDir=os.getcwd(),
0113         outputName=args.output,
0114     ).run()