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