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