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