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