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