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