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