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.alignment import AlignmentDecorator
0008 from acts.examples.odd import getOpenDataDetector
0009 from acts.examples.simulation import (
0010     addParticleGun,
0011     EtaConfig,
0012     ParticleConfig,
0013     MomentumConfig,
0014 )
0015
0016 from acts.examples.root import RootPropagationSummaryWriter, RootPropagationStepsWriter
0017
0018 u = acts.UnitConstants
0019
0020
0021 def runPropagation(
0022     trackingGeometry, field, outputDir, s=None, decorators=[], sterileLogger=True
0023 ):
0024     s = s or acts.examples.Sequencer(events=100, numThreads=1)
0025
0026     for d in decorators:
0027         s.addContextDecorator(d)
0028
0029     rnd = acts.examples.RandomNumbers(seed=42)
0030
0031     addParticleGun(
0032         s,
0033         ParticleConfig(num=1000, pdg=acts.PdgParticle.eMuon, randomizeCharge=True),
0034         EtaConfig(-4.0, 4.0),
0035         MomentumConfig(1 * u.GeV, 100 * u.GeV, transverse=True),
0036         rnd=rnd,
0037     )
0038
0039     trkParamExtractor = acts.examples.ParticleTrackParamExtractor(
0040         level=acts.logging.WARNING,
0041         inputParticles="particles_generated",
0042         outputTrackParameters="params_particles_generated",
0043     )
0044     s.addAlgorithm(trkParamExtractor)
0045
0046     nav = acts.Navigator(trackingGeometry=trackingGeometry)
0047
0048     stepper = acts.EigenStepper(field)
0049     # stepper = acts.AtlasStepper(field)
0050     # stepper = acts.StraightLineStepper()
0051
0052     propagator = acts.examples.ConcretePropagator(acts.Propagator(stepper, nav))
0053
0054     propagationAlgorithm = acts.examples.PropagationAlgorithm(
0055         propagatorImpl=propagator,
0056         level=acts.logging.INFO,
0057         sterileLogger=sterileLogger,
0058         inputTrackParameters="params_particles_generated",
0059         outputSummaryCollection="propagation_summary",
0060     )
0061     s.addAlgorithm(propagationAlgorithm)
0062
0063     s.addWriter(
0064         RootPropagationSummaryWriter(
0065             level=acts.logging.INFO,
0066             inputSummaryCollection="propagation_summary",
0067             filePath=outputDir + "/propagation_summary.root",
0068         )
0069     )
0070
0071     if sterileLogger is False:
0072         s.addWriter(
0073             RootPropagationStepsWriter(
0074                 level=acts.logging.INFO,
0075                 collection="propagation_summary",
0076                 filePath=outputDir + "/propagation_steps.root",
0077             )
0078         )
0079
0080     return s
0081
0082
0083 if "__main__" == __name__:
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(), 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()