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     s: acts.examples.Sequencer = None,
0022 ):
0023     s = runTruthTrackingKalman(
0024         trackingGeometry,
0025         field,
0026         digiConfigFile=digiConfigFile,
0027         outputDir=outputDir,
0028         s=s,
0029     )
0030
0031     kalmanOptions = {
0032         "multipleScattering": multipleScattering,
0033         "energyLoss": energyLoss,
0034         "reverseFilteringMomThreshold": reverseFilteringMomThreshold,
0035         "freeToBoundCorrection": acts.examples.FreeToBoundCorrection(False),
0036         "level": acts.logging.INFO,
0037         "chi2Cut": float("inf"),
0038     }
0039
0040     s.addAlgorithm(
0041         acts.examples.RefittingAlgorithm(
0042             level=acts.logging.INFO,
0043             inputTracks="kf_tracks",
0044             outputTracks="kf_refit_tracks",
0045             fit=acts.examples.makeKalmanFitterFunction(
0046                 trackingGeometry, field, **kalmanOptions
0047             ),
0048         )
0049     )
0050
0051     s.addAlgorithm(
0052         acts.examples.TrackTruthMatcher(
0053             level=acts.logging.INFO,
0054             inputTracks="kf_refit_tracks",
0055             inputParticles="particles_selected",
0056             inputMeasurementParticlesMap="measurement_particles_map",
0057             outputTrackParticleMatching="refit_track_particle_matching",
0058             outputParticleTrackMatching="refit_particle_track_matching",
0059         )
0060     )
0061
0062     s.addWriter(
0063         acts.examples.RootTrackStatesWriter(
0064             level=acts.logging.INFO,
0065             inputTracks="kf_refit_tracks",
0066             inputParticles="particles_selected",
0067             inputTrackParticleMatching="refit_track_particle_matching",
0068             inputSimHits="simhits",
0069             inputMeasurementSimHitsMap="measurement_simhits_map",
0070             filePath=str(outputDir / "trackstates_kf_refit.root"),
0071         )
0072     )
0073
0074     s.addWriter(
0075         acts.examples.RootTrackSummaryWriter(
0076             level=acts.logging.INFO,
0077             inputTracks="kf_refit_tracks",
0078             inputParticles="particles_selected",
0079             inputTrackParticleMatching="refit_track_particle_matching",
0080             filePath=str(outputDir / "tracksummary_kf_refit.root"),
0081         )
0082     )
0083
0084     s.addWriter(
0085         acts.examples.TrackFitterPerformanceWriter(
0086             level=acts.logging.INFO,
0087             inputTracks="kf_refit_tracks",
0088             inputParticles="particles_selected",
0089             inputTrackParticleMatching="refit_track_particle_matching",
0090             filePath=str(outputDir / "performance_kf_refit.root"),
0091         )
0092     )
0093
0094     return s
0095
0096
0097 if __name__ == "__main__":
0098     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0099     outputDir = Path.cwd()
0100
0101     # ODD
0102     from acts.examples.odd import getOpenDataDetector
0103
0104     detector = getOpenDataDetector()
0105     trackingGeometry = detector.trackingGeometry()
0106     decorators = detector.contextDecorators()
0107     digiConfigFile = srcdir / "Examples/Configs/odd-digi-smearing-config.json"
0108
0109     ## GenericDetector
0110     # detector = acts.examples.GenericDetector()
0111     # trackingGeometry = detector.trackingGeometry()
0112     # digiConfigFile = (
0113     #     srcdir
0114     #     / "Examples/Configs/generic-digi-smearing-config.json"
0115     # )
0116
0117     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0118
0119     runRefittingKf(
0120         trackingGeometry=trackingGeometry,
0121         field=field,
0122         digiConfigFile=digiConfigFile,
0123         outputDir=Path.cwd(),
0124     ).run()