Scripts/Python/material_validation_itk.py

0001 #!/usr/bin/env python3
0002 import os
0003 import argparse
0004 from pathlib import Path
0005
0006 from acts.examples import Sequencer
0007 from acts.examples.simulation import addParticleGun, EtaConfig, ParticleConfig
0008 from acts.examples.root import RootMaterialTrackWriter, RootPropagationStepsWriter
0009
0010 import acts
0011
0012
0013 def runMaterialValidation(
0014     nevents,
0015     ntracks,
0016     trackingGeometry,
0017     decorators,
0018     field,
0019     outputDir,
0020     outputName="propagation_material",
0021     dumpPropagationSteps=False,
0022     s=None,
0023 ):
0024     # Create a sequencer
0025     s = s or Sequencer(events=nevents, numThreads=-1)
0026
0027     rnd = acts.examples.RandomNumbers(seed=42)
0028
0029     for decorator in decorators:
0030         s.addContextDecorator(decorator)
0031
0032     nav = acts.Navigator(
0033         trackingGeometry=trackingGeometry,
0034         resolveSensitive=True,
0035         resolveMaterial=True,
0036         resolvePassive=True,
0037     )
0038
0039     # stepper = acts.StraightLineStepper()
0040     stepper = acts.EigenStepper(field)
0041
0042     prop = acts.examples.ConcretePropagator(acts.Propagator(stepper, nav))
0043
0044     addParticleGun(
0045         s,
0046         ParticleConfig(num=ntracks, pdg=acts.PdgParticle.eMuon, randomizeCharge=True),
0047         EtaConfig(-4.0, 4.0),
0048         rnd=rnd,
0049     )
0050
0051     trkParamExtractor = acts.examples.ParticleTrackParamExtractor(
0052         level=acts.logging.INFO,
0053         inputParticles="particles_generated",
0054         outputTrackParameters="params_particles_generated",
0055     )
0056     s.addAlgorithm(trkParamExtractor)
0057
0058     alg = acts.examples.PropagationAlgorithm(
0059         propagatorImpl=prop,
0060         level=acts.logging.INFO,
0061         sterileLogger=False,
0062         recordMaterialInteractions=True,
0063         inputTrackParameters="params_particles_generated",
0064         outputPropagationSteps="propagation_steps",
0065         outputMaterialTracks="material-tracks",
0066     )
0067     s.addAlgorithm(alg)
0068
0069     s.addWriter(
0070         RootMaterialTrackWriter(
0071             level=acts.logging.INFO,
0072             inputMaterialTracks=alg.config.outputMaterialCollection,
0073             filePath=os.path.join(outputDir, (outputName + ".root")),
0074             storeSurface=True,
0075             storeVolume=True,
0076         )
0077     )
0078
0079     if dumpPropagationSteps:
0080         s.addWriter(
0081             RootPropagationStepsWriter(
0082                 level=acts.logging.INFO,
0083                 collection=alg.config.outputSummaryCollection,
0084                 filePath=outputDir + "/propagation_steps.root",
0085             )
0086         )
0087
0088     return s
0089
0090
0091 if "__main__" == __name__:
0092     p = argparse.ArgumentParser(
0093         description="Script to run material validation on ITk geometry"
0094     )
0095     p.add_argument(
0096         "geo_dir",
0097         help="Input directory containing the ITk standalone geometry. Get in touch if you don't have this.",
0098     )
0099     p.add_argument("--material", type=str, default="", help="Material file")
0100
0101     args = p.parse_args()
0102
0103     geo_example_dir = Path(args.geo_dir)
0104     assert geo_example_dir.exists(), "Detector example input directory missing"
0105     assert os.path.exists(
0106         args.material
0107     ), "Invalid file path/name in --material. Please check your input!"
0108
0109     from acts.examples.itk import buildITkGeometry
0110
0111     detector = buildITkGeometry(geo_example_dir, customMaterialFile=args.material)
0112     trackingGeometry = detector.trackingGeometry()
0113     decorators = detector.contextDecorators()
0114
0115     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * acts.UnitConstants.T))
0116
0117     runMaterialValidation(
0118         trackingGeometry, decorators, field, outputDir=os.getcwd()
0119     ).run()