Scripts/Python/truth_tracking_gsf.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 runTruthTrackingGsf(
0013     trackingGeometry: acts.TrackingGeometry,
0014     field: acts.MagneticFieldProvider,
0015     digiConfigFile: Path,
0016     outputDir: Path,
0017     inputParticlePath: Optional[Path] = None,
0018     decorators=[],
0019     s: acts.examples.Sequencer = None,
0020 ):
0021     from acts.examples.simulation import (
0022         addParticleGun,
0023         ParticleConfig,
0024         EtaConfig,
0025         PhiConfig,
0026         MomentumConfig,
0027         addFatras,
0028         addDigitization,
0029         ParticleSelectorConfig,
0030         addDigiParticleSelection,
0031     )
0032     from acts.examples.reconstruction import (
0033         addSeeding,
0034         SeedingAlgorithm,
0035         addTruthTrackingGsf,
0036     )
0037
0038     s = s or acts.examples.Sequencer(
0039         events=100, numThreads=-1, logLevel=acts.logging.INFO
0040     )
0041
0042     for d in decorators:
0043         s.addContextDecorator(d)
0044
0045     rnd = acts.examples.RandomNumbers(seed=42)
0046     outputDir = Path(outputDir)
0047
0048     if inputParticlePath is None:
0049         addParticleGun(
0050             s,
0051             ParticleConfig(num=1, pdg=acts.PdgParticle.eElectron, randomizeCharge=True),
0052             EtaConfig(-3.0, 3.0, uniform=True),
0053             MomentumConfig(1.0 * u.GeV, 100.0 * u.GeV, transverse=True),
0054             PhiConfig(0.0, 360.0 * u.degree),
0055             vtxGen=acts.examples.GaussianVertexGenerator(
0056                 mean=acts.Vector4(0, 0, 0, 0),
0057                 stddev=acts.Vector4(0, 0, 0, 0),
0058             ),
0059             multiplicity=1,
0060             rnd=rnd,
0061         )
0062     else:
0063         acts.logging.getLogger("GSF Example").info(
0064             "Reading particles from %s", inputParticlePath.resolve()
0065         )
0066         assert inputParticlePath.exists()
0067         s.addReader(
0068             acts.examples.RootParticleReader(
0069                 level=acts.logging.INFO,
0070                 filePath=str(inputParticlePath.resolve()),
0071                 outputParticles="particles_generated",
0072             )
0073         )
0074
0075     addFatras(
0076         s,
0077         trackingGeometry,
0078         field,
0079         rnd=rnd,
0080         enableInteractions=True,
0081     )
0082
0083     addDigitization(
0084         s,
0085         trackingGeometry,
0086         field,
0087         digiConfigFile=digiConfigFile,
0088         rnd=rnd,
0089     )
0090
0091     addDigiParticleSelection(
0092         s,
0093         ParticleSelectorConfig(
0094             pt=(0.9 * u.GeV, None),
0095             measurements=(7, None),
0096             removeNeutral=True,
0097             removeSecondaries=True,
0098         ),
0099     )
0100
0101     addSeeding(
0102         s,
0103         trackingGeometry,
0104         field,
0105         rnd=rnd,
0106         inputParticles="particles_generated",
0107         seedingAlgorithm=SeedingAlgorithm.TruthSmeared,
0108         particleHypothesis=acts.ParticleHypothesis.electron,
0109     )
0110
0111     addTruthTrackingGsf(
0112         s,
0113         trackingGeometry,
0114         field,
0115     )
0116
0117     s.addAlgorithm(
0118         acts.examples.TrackSelectorAlgorithm(
0119             level=acts.logging.INFO,
0120             inputTracks="tracks",
0121             outputTracks="selected-tracks",
0122             selectorConfig=acts.TrackSelector.Config(
0123                 minMeasurements=7,
0124             ),
0125         )
0126     )
0127     s.addWhiteboardAlias("tracks", "selected-tracks")
0128
0129     s.addWriter(
0130         acts.examples.RootTrackStatesWriter(
0131             level=acts.logging.INFO,
0132             inputTracks="tracks",
0133             inputParticles="particles_selected",
0134             inputTrackParticleMatching="track_particle_matching",
0135             inputSimHits="simhits",
0136             inputMeasurementSimHitsMap="measurement_simhits_map",
0137             filePath=str(outputDir / "trackstates_gsf.root"),
0138         )
0139     )
0140
0141     s.addWriter(
0142         acts.examples.RootTrackSummaryWriter(
0143             level=acts.logging.INFO,
0144             inputTracks="tracks",
0145             inputParticles="particles_selected",
0146             inputTrackParticleMatching="track_particle_matching",
0147             filePath=str(outputDir / "tracksummary_gsf.root"),
0148             writeGsfSpecific=True,
0149         )
0150     )
0151
0152     s.addWriter(
0153         acts.examples.TrackFitterPerformanceWriter(
0154             level=acts.logging.INFO,
0155             inputTracks="tracks",
0156             inputParticles="particles_selected",
0157             inputTrackParticleMatching="track_particle_matching",
0158             filePath=str(outputDir / "performance_gsf.root"),
0159         )
0160     )
0161
0162     return s
0163
0164
0165 if "__main__" == __name__:
0166     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0167
0168     # ODD
0169     from acts.examples.odd import getOpenDataDetector
0170
0171     detector = getOpenDataDetector()
0172     trackingGeometry = detector.trackingGeometry()
0173     digiConfigFile = srcdir / "Examples/Configs/odd-digi-smearing-config.json"
0174
0175     ## GenericDetector
0176     # detector = acts.examples.GenericDetector()
0177     # trackingGeometry = detector.trackingGeometry()
0178     # digiConfigFile = (
0179     #     srcdir
0180     #     / "Examples/Configs/generic-digi-smearing-config.json"
0181     # )
0182
0183     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0184
0185     runTruthTrackingGsf(
0186         trackingGeometry=trackingGeometry,
0187         field=field,
0188         digiConfigFile=digiConfigFile,
0189         outputDir=Path.cwd(),
0190     ).run()