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     reverseFilteringMomThreshold=0 * u.GeV,
0021     s: acts.examples.Sequencer = None,
0022 ):
0023     from acts.examples.simulation import (
0024         addParticleGun,
0025         ParticleConfig,
0026         EtaConfig,
0027         PhiConfig,
0028         MomentumConfig,
0029         addFatras,
0030         addDigitization,
0031         ParticleSelectorConfig,
0032         addDigiParticleSelection,
0033     )
0034     from acts.examples.reconstruction import (
0035         addSeeding,
0036         SeedingAlgorithm,
0037         addKalmanTracks,
0038     )
0039
0040     s = s or acts.examples.Sequencer(
0041         events=100, numThreads=-1, logLevel=acts.logging.INFO
0042     )
0043
0044     for d in decorators:
0045         s.addContextDecorator(d)
0046
0047     rnd = acts.examples.RandomNumbers(seed=42)
0048     outputDir = Path(outputDir)
0049
0050     logger = acts.logging.getLogger("Truth tracking example")
0051
0052     if inputParticlePath is None:
0053         addParticleGun(
0054             s,
0055             ParticleConfig(num=1, pdg=acts.PdgParticle.eMuon, randomizeCharge=True),
0056             EtaConfig(-3.0, 3.0, uniform=True),
0057             MomentumConfig(1.0 * u.GeV, 100.0 * u.GeV, transverse=True),
0058             PhiConfig(0.0, 360.0 * u.degree),
0059             vtxGen=acts.examples.GaussianVertexGenerator(
0060                 mean=acts.Vector4(0, 0, 0, 0),
0061                 stddev=acts.Vector4(0, 0, 0, 0),
0062             ),
0063             multiplicity=1,
0064             rnd=rnd,
0065         )
0066     else:
0067         logger.info("Reading particles from %s", inputParticlePath.resolve())
0068         assert inputParticlePath.exists()
0069         s.addReader(
0070             acts.examples.RootParticleReader(
0071                 level=acts.logging.INFO,
0072                 filePath=str(inputParticlePath.resolve()),
0073                 outputParticles="particles_generated",
0074             )
0075         )
0076         s.addWhiteboardAlias("particles", "particles_generated")
0077
0078     if inputHitsPath is None:
0079         addFatras(
0080             s,
0081             trackingGeometry,
0082             field,
0083             rnd=rnd,
0084             enableInteractions=True,
0085         )
0086     else:
0087         logger.info("Reading hits from %s", inputHitsPath.resolve())
0088         assert inputHitsPath.exists()
0089         s.addReader(
0090             acts.examples.RootSimHitReader(
0091                 level=acts.logging.INFO,
0092                 filePath=str(inputHitsPath.resolve()),
0093                 outputSimHits="simhits",
0094             )
0095         )
0096
0097     addDigitization(
0098         s,
0099         trackingGeometry,
0100         field,
0101         digiConfigFile=digiConfigFile,
0102         rnd=rnd,
0103     )
0104
0105     addDigiParticleSelection(
0106         s,
0107         ParticleSelectorConfig(
0108             pt=(0.9 * u.GeV, None),
0109             measurements=(7, None),
0110             removeNeutral=True,
0111             removeSecondaries=True,
0112         ),
0113     )
0114
0115     addSeeding(
0116         s,
0117         trackingGeometry,
0118         field,
0119         rnd=rnd,
0120         inputParticles="particles_generated",
0121         seedingAlgorithm=SeedingAlgorithm.TruthSmeared,
0122         particleHypothesis=acts.ParticleHypothesis.muon,
0123     )
0124
0125     addKalmanTracks(
0126         s,
0127         trackingGeometry,
0128         field,
0129         reverseFilteringMomThreshold,
0130     )
0131
0132     s.addAlgorithm(
0133         acts.examples.TrackSelectorAlgorithm(
0134             level=acts.logging.INFO,
0135             inputTracks="tracks",
0136             outputTracks="selected-tracks",
0137             selectorConfig=acts.TrackSelector.Config(
0138                 minMeasurements=7,
0139             ),
0140         )
0141     )
0142     s.addWhiteboardAlias("tracks", "selected-tracks")
0143
0144     s.addWriter(
0145         acts.examples.RootTrackStatesWriter(
0146             level=acts.logging.INFO,
0147             inputTracks="tracks",
0148             inputParticles="particles_selected",
0149             inputTrackParticleMatching="track_particle_matching",
0150             inputSimHits="simhits",
0151             inputMeasurementSimHitsMap="measurement_simhits_map",
0152             filePath=str(outputDir / "trackstates_kf.root"),
0153         )
0154     )
0155
0156     s.addWriter(
0157         acts.examples.RootTrackSummaryWriter(
0158             level=acts.logging.INFO,
0159             inputTracks="tracks",
0160             inputParticles="particles_selected",
0161             inputTrackParticleMatching="track_particle_matching",
0162             filePath=str(outputDir / "tracksummary_kf.root"),
0163         )
0164     )
0165
0166     s.addWriter(
0167         acts.examples.TrackFitterPerformanceWriter(
0168             level=acts.logging.INFO,
0169             inputTracks="tracks",
0170             inputParticles="particles_selected",
0171             inputTrackParticleMatching="track_particle_matching",
0172             filePath=str(outputDir / "performance_kf.root"),
0173         )
0174     )
0175
0176     return s
0177
0178
0179 if "__main__" == __name__:
0180     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0181
0182     # ODD
0183     from acts.examples.odd import getOpenDataDetector
0184
0185     detector = getOpenDataDetector()
0186     trackingGeometry = detector.trackingGeometry()
0187     digiConfigFile = srcdir / "Examples/Configs/odd-digi-smearing-config.json"
0188
0189     ## GenericDetector
0190     # detector = acts.examples.GenericDetector()
0191     # trackingGeometry = detector.trackingGeometry()
0192     # digiConfigFile = (
0193     #     srcdir
0194     #     / "Examples/Configs/generic-digi-smearing-config.json"
0195     # )
0196
0197     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0198
0199     runTruthTrackingKalman(
0200         trackingGeometry=trackingGeometry,
0201         field=field,
0202         digiConfigFile=digiConfigFile,
0203         outputDir=Path.cwd(),
0204     ).run()