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
0007 from acts.examples.odd import getOpenDataDetector
0008 from acts.examples.simulation import (
0009     addParticleGun,
0010     EtaConfig,
0011     ParticleConfig,
0012     MomentumConfig,
0013 )
0014
0015 from acts.examples.root import RootPropagationSummaryWriter, RootPropagationStepsWriter
0016
0017 u = acts.UnitConstants
0018
0019
0020 def runPropagation(
0021     trackingGeometry, field, outputDir, s=None, decorators=[], sterileLogger=True
0022 ):
0023     s = s or acts.examples.Sequencer(events=100, numThreads=1)
0024
0025     for d in decorators:
0026         s.addContextDecorator(d)
0027
0028     rnd = acts.examples.RandomNumbers(seed=42)
0029
0030     addParticleGun(
0031         s,
0032         ParticleConfig(num=1000, pdg=acts.PdgParticle.eMuon, randomizeCharge=True),
0033         EtaConfig(-4.0, 4.0),
0034         MomentumConfig(1 * u.GeV, 100 * u.GeV, transverse=True),
0035         rnd=rnd,
0036     )
0037
0038     trkParamExtractor = acts.examples.ParticleTrackParamExtractor(
0039         level=acts.logging.WARNING,
0040         inputParticles="particles_generated",
0041         outputTrackParameters="params_particles_generated",
0042     )
0043     s.addAlgorithm(trkParamExtractor)
0044
0045     nav = acts.Navigator(trackingGeometry=trackingGeometry)
0046
0047     stepper = acts.EigenStepper(field)
0048     # stepper = acts.AtlasStepper(field)
0049     # stepper = acts.StraightLineStepper()
0050
0051     propagator = acts.examples.ConcretePropagator(acts.Propagator(stepper, nav))
0052
0053     propagationAlgorithm = acts.examples.PropagationAlgorithm(
0054         propagatorImpl=propagator,
0055         level=acts.logging.INFO,
0056         sterileLogger=sterileLogger,
0057         inputTrackParameters="params_particles_generated",
0058         outputSummaryCollection="propagation_summary",
0059     )
0060     s.addAlgorithm(propagationAlgorithm)
0061
0062     s.addWriter(
0063         RootPropagationSummaryWriter(
0064             level=acts.logging.INFO,
0065             inputSummaryCollection="propagation_summary",
0066             filePath=outputDir + "/propagation_summary.root",
0067         )
0068     )
0069
0070     if sterileLogger is False:
0071         s.addWriter(
0072             RootPropagationStepsWriter(
0073                 level=acts.logging.INFO,
0074                 collection="propagation_summary",
0075                 filePath=outputDir + "/propagation_steps.root",
0076             )
0077         )
0078
0079     return s
0080
0081
0082 if "__main__" == __name__:
0083     # from acts.examples.alignment import AlignmentDecorator
0084     matDeco = None
0085     contextDecorators = []
0086     # matDeco = acts.IMaterialDecorator.fromFile("material.json")
0087     # matDeco = acts.IMaterialDecorator.fromFile("material.root")
0088
0089     ## Generic detector: Default
0090     detector = GenericDetector(materialDecorator=matDeco)
0091     trackingGeometry = detector.trackingGeometry()
0092
0093     ## Alternative: Aligned Generic detector
0094     # detector = AlignedGenericDetector(materialDecorator=matDeco)
0095
0096     ## Alternative: DD4hep detector
0097     # detector = getOpenDataDetector()
0098     # trackingGeometry = detector.trackingGeometry()
0099
0100     ## Alternative: Misaligned DD4hep detector
0101     # detector = getOpenDataDetector(misaligned=True)
0102     # trackingGeometry = detector.trackingGeometry()
0103     # structureSelector = StructureSelector(trackingGeometry)
0104     # pixelBarrelID = acts.GeometryIdentifier(volume=17)
0105     # pixelBarrelTransforms = structureSelector.selectedTransforms(
0106     #     acts.GeometryContext.dangerouslyDefaultConstruct(), pixelBarrelID
0107     # )
0108     # alignDecoConfig = AlignmentDecorator.Config()
0109     # alignDecoConfig.nominalStore = acts.examples.GeoIdAlignmentStore(
0110     #     pixelBarrelTransforms
0111     # )
0112
0113     # gRot = acts.examples.AlignmentGeneratorGlobalRotation()
0114     # gRot.axis = acts.Vector3(0.0, 0.0, 1.0)
0115     # gRot.angle = 0.05
0116
0117     # lShift = acts.examples.AlignmentGeneratorLocalShift()
0118     # lShift.axisDirection = acts.AxisDirection.AxisZ
0119     # lShift.shift = 3.0
0120
0121     # alignDecoConfig.iovGenerators = [((0, 25), lShift), ((25, 50), gRot)]
0122     # alignDecoConfig.garbageCollection = True
0123     # alignDecoConfig.gcInterval = 20
0124
0125     # alignDeco = AlignmentDecorator(alignDecoConfig, acts.logging.VERBOSE)
0126     # contextDecorators = [alignDeco]
0127
0128     ## Magnetic field setup: Default: constant 2T longitudinal field
0129     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * acts.UnitConstants.T))
0130
0131     ## Alternative: no B field
0132     # field = acts.NullBField()
0133
0134     ## Alternative: Analytical solenoid B field, discretized in an interpolated field map
0135     # solenoid = acts.SolenoidBField(
0136     #     radius = 1200*u.mm,
0137     #     length = 6000*u.mm,
0138     #     bMagCenter = 2*u.T,
0139     #     nCoils = 1194
0140     # )
0141     # field = acts.solenoidFieldMap(
0142     #     rlim=(0, 1200*u.mm),
0143     #     zlim=(-5000*u.mm, 5000*u.mm),
0144     #     nbins=(50, 50),
0145     #     field=solenoid
0146     # )
0147
0148     os.makedirs(os.getcwd() + "/propagation", exist_ok=True)
0149
0150     runPropagation(
0151         trackingGeometry=trackingGeometry,
0152         field=field,
0153         outputDir=os.getcwd() + "/propagation",
0154         s=None,
0155         decorators=contextDecorators,
0156         sterileLogger=True,
0157     ).run()