Scripts/Python/truth_tracking_kalman.py

0001 #!/usr/bin/env python3
0002
0003 from pathlib import Path
0004 from typing import Optional
0005
0006 import acts
0007 import acts.examples
0008
0009 u = acts.UnitConstants
0010
0011
0012 def runTruthTrackingKalman(
0013     trackingGeometry: acts.TrackingGeometry,
0014     field: acts.MagneticFieldProvider,
0015     digiConfigFile: Path,
0016     outputDir: Path,
0017     inputParticlePath: Optional[Path] = None,
0018     decorators=[],
0019     directNavigation=False,
0020     reverseFilteringMomThreshold=0 * u.GeV,
0021     s: acts.examples.Sequencer = None,
0022 ):
0023     from acts.examples.simulation import (
0024         addParticleGun,
0025         ParticleConfig,
0026         EtaConfig,
0027         PhiConfig,
0028         MomentumConfig,
0029         addFatras,
0030         addDigitization,
0031     )
0032     from acts.examples.reconstruction import (
0033         addSeeding,
0034         SeedingAlgorithm,
0035         TruthSeedRanges,
0036         addKalmanTracks,
0037     )
0038
0039     s = s or acts.examples.Sequencer(
0040         events=100, numThreads=-1, logLevel=acts.logging.INFO
0041     )
0042
0043     for d in decorators:
0044         s.addContextDecorator(d)
0045
0046     rnd = acts.examples.RandomNumbers(seed=42)
0047     outputDir = Path(outputDir)
0048
0049     if inputParticlePath is None:
0050         addParticleGun(
0051             s,
0052             ParticleConfig(num=1, pdg=acts.PdgParticle.eMuon, randomizeCharge=True),
0053             EtaConfig(-3.0, 3.0, uniform=True),
0054             MomentumConfig(1.0 * u.GeV, 100.0 * u.GeV, transverse=True),
0055             PhiConfig(0.0, 360.0 * u.degree),
0056             vtxGen=acts.examples.GaussianVertexGenerator(
0057                 mean=acts.Vector4(0, 0, 0, 0),
0058                 stddev=acts.Vector4(0, 0, 0, 0),
0059             ),
0060             multiplicity=1,
0061             rnd=rnd,
0062         )
0063     else:
0064         acts.logging.getLogger("Truth tracking example").info(
0065             "Reading particles from %s", inputParticlePath.resolve()
0066         )
0067         assert inputParticlePath.exists()
0068         s.addReader(
0069             acts.examples.RootParticleReader(
0070                 level=acts.logging.INFO,
0071                 filePath=str(inputParticlePath.resolve()),
0072                 outputParticles="particles_input",
0073             )
0074         )
0075
0076     addFatras(
0077         s,
0078         trackingGeometry,
0079         field,
0080         rnd=rnd,
0081         enableInteractions=True,
0082     )
0083
0084     addDigitization(
0085         s,
0086         trackingGeometry,
0087         field,
0088         digiConfigFile=digiConfigFile,
0089         rnd=rnd,
0090     )
0091
0092     addSeeding(
0093         s,
0094         trackingGeometry,
0095         field,
0096         rnd=rnd,
0097         inputParticles="particles_input",
0098         seedingAlgorithm=SeedingAlgorithm.TruthSmeared,
0099         particleHypothesis=acts.ParticleHypothesis.muon,
0100         truthSeedRanges=TruthSeedRanges(
0101             pt=(1 * u.GeV, None),
0102             nHits=(7, None),
0103         ),
0104     )
0105
0106     addKalmanTracks(
0107         s,
0108         trackingGeometry,
0109         field,
0110         directNavigation,
0111         reverseFilteringMomThreshold,
0112     )
0113
0114     s.addAlgorithm(
0115         acts.examples.TrackSelectorAlgorithm(
0116             level=acts.logging.INFO,
0117             inputTracks="tracks",
0118             outputTracks="selected-tracks",
0119             selectorConfig=acts.TrackSelector.Config(
0120                 minMeasurements=7,
0121             ),
0122         )
0123     )
0124     s.addWhiteboardAlias("tracks", "selected-tracks")
0125
0126     s.addWriter(
0127         acts.examples.RootTrackStatesWriter(
0128             level=acts.logging.INFO,
0129             inputTracks="tracks",
0130             inputParticles="truth_seeds_selected",
0131             inputTrackParticleMatching="track_particle_matching",
0132             inputSimHits="simhits",
0133             inputMeasurementSimHitsMap="measurement_simhits_map",
0134             filePath=str(outputDir / "trackstates_fitter.root"),
0135         )
0136     )
0137
0138     s.addWriter(
0139         acts.examples.RootTrackSummaryWriter(
0140             level=acts.logging.INFO,
0141             inputTracks="tracks",
0142             inputParticles="truth_seeds_selected",
0143             inputTrackParticleMatching="track_particle_matching",
0144             filePath=str(outputDir / "tracksummary_fitter.root"),
0145         )
0146     )
0147
0148     s.addWriter(
0149         acts.examples.TrackFitterPerformanceWriter(
0150             level=acts.logging.INFO,
0151             inputTracks="tracks",
0152             inputParticles="truth_seeds_selected",
0153             inputTrackParticleMatching="track_particle_matching",
0154             filePath=str(outputDir / "performance_track_fitter.root"),
0155         )
0156     )
0157
0158     return s
0159
0160
0161 if "__main__" == __name__:
0162     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0163
0164     # detector, trackingGeometry, _ = getOpenDataDetector()
0165     detector, trackingGeometry, decorators = acts.examples.GenericDetector.create()
0166
0167     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0168
0169     runTruthTrackingKalman(
0170         trackingGeometry,
0171         field,
0172         digiConfigFile=srcdir
0173         / "Examples/Algorithms/Digitization/share/default-smearing-config-generic.json",
0174         # "thirdparty/OpenDataDetector/config/odd-digi-smearing-config.json",
0175         outputDir=Path.cwd(),
0176     ).run()