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