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, StructureSelector, AlignmentDecorator
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 is False:
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     contextDecorators = []
0083     # matDeco = acts.IMaterialDecorator.fromFile("material.json")
0084     # matDeco = acts.IMaterialDecorator.fromFile("material.root")
0085
0086     ## Generic detector: Default
0087     detector = GenericDetector(materialDecorator=matDeco)
0088     trackingGeometry = detector.trackingGeometry()
0089
0090     ## Alternative: Aligned Generic detector
0091     # detector = AlignedGenericDetector(materialDecorator=matDeco)
0092
0093     ## Alternative: DD4hep detector
0094     # detector = getOpenDataDetector()
0095     # trackingGeometry = detector.trackingGeometry()
0096
0097     ## Alternative: Misaligned DD4hep detector
0098     # detector = getOpenDataDetector(misaligned=True)
0099     # trackingGeometry = detector.trackingGeometry()
0100     # structureSelector = StructureSelector(trackingGeometry)
0101     # pixelBarrelID = acts.GeometryIdentifier(volume=17)
0102     # pixelBarrelTransforms = structureSelector.selectedTransforms(
0103     #     acts.GeometryContext(), pixelBarrelID
0104     # )
0105     # alignDecoConfig = AlignmentDecorator.Config()
0106     # alignDecoConfig.nominalStore = acts.examples.GeoIdAlignmentStore(
0107     #     pixelBarrelTransforms
0108     # )
0109
0110     # gRot = acts.examples.AlignmentGeneratorGlobalRotation()
0111     # gRot.axis = acts.Vector3(0.0, 0.0, 1.0)
0112     # gRot.angle = 0.05
0113
0114     # lShift = acts.examples.AlignmentGeneratorLocalShift()
0115     # lShift.axisDirection = acts.AxisDirection.AxisZ
0116     # lShift.shift = 3.0
0117
0118     # alignDecoConfig.iovGenerators = [((0, 25), lShift), ((25, 50), gRot)]
0119     # alignDecoConfig.garbageCollection = True
0120     # alignDecoConfig.gcInterval = 20
0121
0122     # alignDeco = AlignmentDecorator(alignDecoConfig, acts.logging.VERBOSE)
0123     # contextDecorators = [alignDeco]
0124
0125     ## Magnetic field setup: Default: constant 2T longitudinal field
0126     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * acts.UnitConstants.T))
0127
0128     ## Alternative: no B field
0129     # field = acts.NullBField()
0130
0131     ## Alternative: Analytical solenoid B field, discretized in an interpolated field map
0132     # solenoid = acts.SolenoidBField(
0133     #     radius = 1200*u.mm,
0134     #     length = 6000*u.mm,
0135     #     bMagCenter = 2*u.T,
0136     #     nCoils = 1194
0137     # )
0138     # field = acts.solenoidFieldMap(
0139     #     rlim=(0, 1200*u.mm),
0140     #     zlim=(-5000*u.mm, 5000*u.mm),
0141     #     nbins=(50, 50),
0142     #     field=solenoid
0143     # )
0144
0145     os.makedirs(os.getcwd() + "/propagation", exist_ok=True)
0146
0147     runPropagation(
0148         trackingGeometry=trackingGeometry,
0149         field=field,
0150         outputDir=os.getcwd() + "/propagation",
0151         s=None,
0152         decorators=contextDecorators,
0153         sterileLogger=True,
0154     ).run()