Scripts/Python/propagation.py

0001 #!/usr/bin/env python3
0002
0003 import os
0004 import acts
0005 import acts.examples
0006 from acts.examples import GenericDetector, AlignedGenericDetector
0007 from acts.examples.odd import getOpenDataDetector
0008 from acts.examples.simulation import (
0009     addParticleGun,
0010     EtaConfig,
0011     ParticleConfig,
0012     MomentumConfig,
0013 )
0014
0015 u = acts.UnitConstants
0016
0017
0018 def runPropagation(
0019     trackingGeometry, field, outputDir, s=None, decorators=[], sterileLogger=True
0020 ):
0021     s = s or acts.examples.Sequencer(events=100, numThreads=1)
0022
0023     for d in decorators:
0024         s.addContextDecorator(d)
0025
0026     rnd = acts.examples.RandomNumbers(seed=42)
0027
0028     addParticleGun(
0029         s,
0030         ParticleConfig(num=1000, pdg=acts.PdgParticle.eMuon, randomizeCharge=True),
0031         EtaConfig(-4.0, 4.0),
0032         MomentumConfig(1 * u.GeV, 100 * u.GeV, transverse=True),
0033         rnd=rnd,
0034     )
0035
0036     trkParamExtractor = acts.examples.ParticleTrackParamExtractor(
0037         level=acts.logging.WARNING,
0038         inputParticles="particles_generated",
0039         outputTrackParameters="params_particles_generated",
0040     )
0041     s.addAlgorithm(trkParamExtractor)
0042
0043     nav = acts.Navigator(trackingGeometry=trackingGeometry)
0044
0045     stepper = acts.EigenStepper(field)
0046     # stepper = acts.AtlasStepper(field)
0047     # stepper = acts.StraightLineStepper()
0048
0049     propagator = acts.examples.ConcretePropagator(acts.Propagator(stepper, nav))
0050
0051     propagationAlgorithm = acts.examples.PropagationAlgorithm(
0052         propagatorImpl=propagator,
0053         level=acts.logging.INFO,
0054         sterileLogger=sterileLogger,
0055         inputTrackParameters="params_particles_generated",
0056         outputSummaryCollection="propagation_summary",
0057     )
0058     s.addAlgorithm(propagationAlgorithm)
0059
0060     s.addWriter(
0061         acts.examples.RootPropagationSummaryWriter(
0062             level=acts.logging.INFO,
0063             inputSummaryCollection="propagation_summary",
0064             filePath=outputDir + "/propagation_summary.root",
0065         )
0066     )
0067
0068     if sterileLogger:
0069         s.addWriter(
0070             acts.examples.RootPropagationStepsWriter(
0071                 level=acts.logging.INFO,
0072                 collection="propagation_summary",
0073                 filePath=outputDir + "/propagation_steps.root",
0074             )
0075         )
0076
0077     return s
0078
0079
0080 if "__main__" == __name__:
0081     matDeco = None
0082     # matDeco = acts.IMaterialDecorator.fromFile("material.json")
0083     # matDeco = acts.IMaterialDecorator.fromFile("material.root")
0084
0085     ## Generic detector: Default
0086     detector = GenericDetector(materialDecorator=matDeco)
0087
0088     ## Alternative: Aligned Generic detector
0089     # detector = AlignedGenericDetector(materialDecorator=matDeco)
0090
0091     ## Alternative: DD4hep detector
0092     # detector = getOpenDataDetector()
0093     # trackingGeometry = detector.trackingGeometry()
0094
0095     trackingGeometry = detector.trackingGeometry()
0096     contextDecorators = detector.contextDecorators()
0097
0098     ## Magnetic field setup: Default: constant 2T longitudinal field
0099     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * acts.UnitConstants.T))
0100
0101     ## Alternative: no B field
0102     # field = acts.NullBField()
0103
0104     ## Alternative: Analytical solenoid B field, discretized in an interpolated field map
0105     # solenoid = acts.SolenoidBField(
0106     #     radius = 1200*u.mm,
0107     #     length = 6000*u.mm,
0108     #     bMagCenter = 2*u.T,
0109     #     nCoils = 1194
0110     # )
0111     # field = acts.solenoidFieldMap(
0112     #     rlim=(0, 1200*u.mm),
0113     #     zlim=(-5000*u.mm, 5000*u.mm),
0114     #     nbins=(50, 50),
0115     #     field=solenoid
0116     # )
0117
0118     os.makedirs(os.getcwd() + "/propagation", exist_ok=True)
0119
0120     runPropagation(
0121         trackingGeometry,
0122         field,
0123         os.getcwd() + "/propagation",
0124         decorators=contextDecorators,
0125         sterileLogger=True,
0126     ).run()