Scripts/Python/truth_tracking_kalman_refitting.py

0001 #!/usr/bin/env python3
0002
0003 from pathlib import Path
0004
0005 import acts
0006 import acts.examples
0007
0008 from truth_tracking_kalman import runTruthTrackingKalman
0009
0010 u = acts.UnitConstants
0011
0012
0013 def runRefittingKf(
0014     trackingGeometry: acts.TrackingGeometry,
0015     field: acts.MagneticFieldProvider,
0016     digiConfigFile: Path,
0017     outputDir: Path,
0018     multipleScattering: bool = True,
0019     energyLoss: bool = True,
0020     reverseFilteringMomThreshold=0 * u.GeV,
0021     reverseFilteringCovarianceScaling=1.0,
0022     s: acts.examples.Sequencer = None,
0023 ):
0024     s = runTruthTrackingKalman(
0025         trackingGeometry,
0026         field,
0027         digiConfigFile=digiConfigFile,
0028         outputDir=outputDir,
0029         reverseFilteringMomThreshold=reverseFilteringMomThreshold,
0030         reverseFilteringCovarianceScaling=reverseFilteringCovarianceScaling,
0031         s=s,
0032     )
0033
0034     kalmanOptions = {
0035         "multipleScattering": multipleScattering,
0036         "energyLoss": energyLoss,
0037         "reverseFilteringMomThreshold": reverseFilteringMomThreshold,
0038         "reverseFilteringCovarianceScaling": reverseFilteringCovarianceScaling,
0039         "freeToBoundCorrection": acts.examples.FreeToBoundCorrection(False),
0040         "level": acts.logging.INFO,
0041         "chi2Cut": float("inf"),
0042     }
0043
0044     s.addAlgorithm(
0045         acts.examples.RefittingAlgorithm(
0046             level=acts.logging.INFO,
0047             inputTracks="kf_tracks",
0048             outputTracks="kf_refit_tracks",
0049             initialVarInflation=6 * [100.0],
0050             fit=acts.examples.makeKalmanFitterFunction(
0051                 trackingGeometry, field, **kalmanOptions
0052             ),
0053         )
0054     )
0055
0056     s.addAlgorithm(
0057         acts.examples.TrackTruthMatcher(
0058             level=acts.logging.INFO,
0059             inputTracks="kf_refit_tracks",
0060             inputParticles="particles_selected",
0061             inputMeasurementParticlesMap="measurement_particles_map",
0062             outputTrackParticleMatching="refit_track_particle_matching",
0063             outputParticleTrackMatching="refit_particle_track_matching",
0064         )
0065     )
0066
0067     s.addWriter(
0068         acts.examples.RootTrackStatesWriter(
0069             level=acts.logging.INFO,
0070             inputTracks="kf_refit_tracks",
0071             inputParticles="particles_selected",
0072             inputTrackParticleMatching="refit_track_particle_matching",
0073             inputSimHits="simhits",
0074             inputMeasurementSimHitsMap="measurement_simhits_map",
0075             filePath=str(outputDir / "trackstates_kf_refit.root"),
0076         )
0077     )
0078
0079     s.addWriter(
0080         acts.examples.RootTrackSummaryWriter(
0081             level=acts.logging.INFO,
0082             inputTracks="kf_refit_tracks",
0083             inputParticles="particles_selected",
0084             inputTrackParticleMatching="refit_track_particle_matching",
0085             filePath=str(outputDir / "tracksummary_kf_refit.root"),
0086         )
0087     )
0088
0089     s.addWriter(
0090         acts.examples.TrackFitterPerformanceWriter(
0091             level=acts.logging.INFO,
0092             inputTracks="kf_refit_tracks",
0093             inputParticles="particles_selected",
0094             inputTrackParticleMatching="refit_track_particle_matching",
0095             filePath=str(outputDir / "performance_kf_refit.root"),
0096         )
0097     )
0098
0099     return s
0100
0101
0102 if __name__ == "__main__":
0103     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0104     outputDir = Path.cwd()
0105
0106     # ODD
0107     from acts.examples.odd import getOpenDataDetector
0108
0109     detector = getOpenDataDetector()
0110     trackingGeometry = detector.trackingGeometry()
0111     decorators = detector.contextDecorators()
0112     digiConfigFile = srcdir / "Examples/Configs/odd-digi-smearing-config.json"
0113
0114     ## GenericDetector
0115     # detector = acts.examples.GenericDetector()
0116     # trackingGeometry = detector.trackingGeometry()
0117     # digiConfigFile = (
0118     #     srcdir
0119     #     / "Examples/Configs/generic-digi-smearing-config.json"
0120     # )
0121
0122     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0123
0124     runRefittingKf(
0125         trackingGeometry=trackingGeometry,
0126         field=field,
0127         digiConfigFile=digiConfigFile,
0128         outputDir=Path.cwd(),
0129     ).run()