Scripts/Python/truth_tracking_kalman.py

0001 #!/usr/bin/env python3
0002
0003 from pathlib import Path
0004 from typing import Optional
0005
0006 import acts
0007 import acts.examples
0008
0009 u = acts.UnitConstants
0010
0011
0012 def runTruthTrackingKalman(
0013     trackingGeometry: acts.TrackingGeometry,
0014     field: acts.MagneticFieldProvider,
0015     digiConfigFile: Path,
0016     outputDir: Path,
0017     inputParticlePath: Optional[Path] = None,
0018     inputHitsPath: Optional[Path] = None,
0019     decorators=[],
0020     generatedParticleType: acts.PdgParticle = acts.PdgParticle.eMuon,
0021     reverseFilteringMomThreshold=0 * u.GeV,
0022     reverseFilteringCovarianceScaling=100.0,
0023     s: acts.examples.Sequencer = None,
0024 ):
0025     from acts.examples.simulation import (
0026         addParticleGun,
0027         ParticleConfig,
0028         EtaConfig,
0029         PhiConfig,
0030         MomentumConfig,
0031         addFatras,
0032         addDigitization,
0033         ParticleSelectorConfig,
0034         addDigiParticleSelection,
0035     )
0036
0037     from acts.examples.root import (
0038         RootParticleReader,
0039         RootSimHitReader,
0040         RootTrackStatesWriter,
0041         RootTrackSummaryWriter,
0042         RootTrackFitterPerformanceWriter,
0043     )
0044
0045     from acts.examples.reconstruction import (
0046         addSeeding,
0047         SeedingAlgorithm,
0048         TrackSmearingSigmas,
0049         addKalmanTracks,
0050     )
0051
0052     s = s or acts.examples.Sequencer(
0053         events=100, numThreads=-1, logLevel=acts.logging.INFO
0054     )
0055
0056     for d in decorators:
0057         s.addContextDecorator(d)
0058
0059     rnd = acts.examples.RandomNumbers(seed=42)
0060     outputDir = Path(outputDir)
0061
0062     logger = acts.logging.getLogger("Truth tracking example")
0063
0064     if inputParticlePath is None:
0065         addParticleGun(
0066             s,
0067             ParticleConfig(num=1, pdg=generatedParticleType, randomizeCharge=True),
0068             EtaConfig(-3.0, 3.0, uniform=True),
0069             MomentumConfig(1.0 * u.GeV, 100.0 * u.GeV, transverse=True),
0070             PhiConfig(0.0, 360.0 * u.degree),
0071             vtxGen=acts.examples.GaussianVertexGenerator(
0072                 mean=acts.Vector4(0, 0, 0, 0),
0073                 stddev=acts.Vector4(0, 0, 0, 0),
0074             ),
0075             multiplicity=1,
0076             rnd=rnd,
0077         )
0078     else:
0079         logger.info("Reading particles from %s", inputParticlePath.resolve())
0080         assert inputParticlePath.exists()
0081         s.addReader(
0082             RootParticleReader(
0083                 level=acts.logging.INFO,
0084                 filePath=str(inputParticlePath.resolve()),
0085                 outputParticles="particles_generated",
0086             )
0087         )
0088         s.addWhiteboardAlias("particles", "particles_generated")
0089
0090     if inputHitsPath is None:
0091         addFatras(
0092             s,
0093             trackingGeometry,
0094             field,
0095             rnd=rnd,
0096             enableInteractions=True,
0097         )
0098     else:
0099         logger.info("Reading hits from %s", inputHitsPath.resolve())
0100         assert inputHitsPath.exists()
0101         s.addReader(
0102             RootSimHitReader(
0103                 level=acts.logging.INFO,
0104                 filePath=str(inputHitsPath.resolve()),
0105                 outputSimHits="simhits",
0106             )
0107         )
0108         s.addWhiteboardAlias("particles_simulated_selected", "particles_generated")
0109
0110     addDigitization(
0111         s,
0112         trackingGeometry,
0113         field,
0114         digiConfigFile=digiConfigFile,
0115         rnd=rnd,
0116     )
0117
0118     addDigiParticleSelection(
0119         s,
0120         ParticleSelectorConfig(
0121             pt=(0.9 * u.GeV, None),
0122             measurements=(7, None),
0123             removeNeutral=True,
0124             removeSecondaries=True,
0125         ),
0126     )
0127
0128     addSeeding(
0129         s,
0130         trackingGeometry,
0131         field,
0132         rnd=rnd,
0133         inputParticles="particles_generated",
0134         seedingAlgorithm=SeedingAlgorithm.TruthSmeared,
0135         trackSmearingSigmas=TrackSmearingSigmas(
0136             # zero everything so the KF has a chance to find the measurements
0137             loc0=0,
0138             loc0PtA=0,
0139             loc0PtB=0,
0140             loc1=0,
0141             loc1PtA=0,
0142             loc1PtB=0,
0143             time=0,
0144             phi=0,
0145             theta=0,
0146             ptRel=0,
0147         ),
0148         particleHypothesis=acts.ParticleHypothesis.muon,
0149         initialSigmas=[
0150             1 * u.mm,
0151             1 * u.mm,
0152             1 * u.degree,
0153             1 * u.degree,
0154             0 / u.GeV,
0155             1 * u.ns,
0156         ],
0157         initialSigmaQoverPt=0.1 / u.GeV,
0158         initialSigmaPtRel=0.1,
0159         initialVarInflation=[1e0, 1e0, 1e0, 1e0, 1e0, 1e0],
0160     )
0161
0162     addKalmanTracks(
0163         s,
0164         trackingGeometry,
0165         field,
0166         reverseFilteringMomThreshold,
0167         reverseFilteringCovarianceScaling,
0168     )
0169
0170     s.addAlgorithm(
0171         acts.examples.TrackSelectorAlgorithm(
0172             level=acts.logging.INFO,
0173             inputTracks="tracks",
0174             outputTracks="selected-tracks",
0175             selectorConfig=acts.TrackSelector.Config(
0176                 minMeasurements=7,
0177             ),
0178         )
0179     )
0180     s.addWhiteboardAlias("tracks", "selected-tracks")
0181
0182     s.addWriter(
0183         RootTrackStatesWriter(
0184             level=acts.logging.INFO,
0185             inputTracks="tracks",
0186             inputParticles="particles_selected",
0187             inputTrackParticleMatching="track_particle_matching",
0188             inputSimHits="simhits",
0189             inputMeasurementSimHitsMap="measurement_simhits_map",
0190             filePath=str(outputDir / "trackstates_kf.root"),
0191         )
0192     )
0193
0194     s.addWriter(
0195         RootTrackSummaryWriter(
0196             level=acts.logging.INFO,
0197             inputTracks="tracks",
0198             inputParticles="particles_selected",
0199             inputTrackParticleMatching="track_particle_matching",
0200             filePath=str(outputDir / "tracksummary_kf.root"),
0201         )
0202     )
0203
0204     s.addWriter(
0205         RootTrackFitterPerformanceWriter(
0206             level=acts.logging.INFO,
0207             inputTracks="tracks",
0208             inputParticles="particles_selected",
0209             inputTrackParticleMatching="track_particle_matching",
0210             filePath=str(outputDir / "performance_kf.root"),
0211         )
0212     )
0213
0214     return s
0215
0216
0217 if "__main__" == __name__:
0218     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0219
0220     # ODD
0221     from acts.examples.odd import getOpenDataDetector
0222
0223     detector = getOpenDataDetector()
0224     trackingGeometry = detector.trackingGeometry()
0225     digiConfigFile = srcdir / "Examples/Configs/odd-digi-smearing-config.json"
0226
0227     ## GenericDetector
0228     # detector = acts.examples.GenericDetector()
0229     # trackingGeometry = detector.trackingGeometry()
0230     # digiConfigFile = (
0231     #     srcdir
0232     #     / "Examples/Configs/generic-digi-smearing-config.json"
0233     # )
0234
0235     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0236
0237     runTruthTrackingKalman(
0238         trackingGeometry=trackingGeometry,
0239         field=field,
0240         digiConfigFile=digiConfigFile,
0241         outputDir=Path.cwd(),
0242     ).run()