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