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