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     )
0030     from acts.examples.reconstruction import (
0031         addSeeding,
0032         SeedingAlgorithm,
0033         TruthSeedRanges,
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     )
0081
0082     addDigitization(
0083         s,
0084         trackingGeometry,
0085         field,
0086         digiConfigFile=digiConfigFile,
0087         rnd=rnd,
0088     )
0089
0090     addSeeding(
0091         s,
0092         trackingGeometry,
0093         field,
0094         rnd=rnd,
0095         inputParticles="particles_input",
0096         seedingAlgorithm=SeedingAlgorithm.TruthSmeared,
0097         particleHypothesis=acts.ParticleHypothesis.electron,
0098         truthSeedRanges=TruthSeedRanges(
0099             pt=(1 * u.GeV, None),
0100             nHits=(7, None),
0101         ),
0102     )
0103
0104     addTruthTrackingGsf(
0105         s,
0106         trackingGeometry,
0107         field,
0108     )
0109
0110     s.addAlgorithm(
0111         acts.examples.TrackSelectorAlgorithm(
0112             level=acts.logging.INFO,
0113             inputTracks="tracks",
0114             outputTracks="selected-tracks",
0115             selectorConfig=acts.TrackSelector.Config(
0116                 minMeasurements=7,
0117             ),
0118         )
0119     )
0120     s.addWhiteboardAlias("tracks", "selected-tracks")
0121
0122     s.addWriter(
0123         acts.examples.RootTrackStatesWriter(
0124             level=acts.logging.INFO,
0125             inputTracks="tracks",
0126             inputParticles="truth_seeds_selected",
0127             inputTrackParticleMatching="track_particle_matching",
0128             inputSimHits="simhits",
0129             inputMeasurementSimHitsMap="measurement_simhits_map",
0130             filePath=str(outputDir / "trackstates_gsf.root"),
0131         )
0132     )
0133
0134     s.addWriter(
0135         acts.examples.RootTrackSummaryWriter(
0136             level=acts.logging.INFO,
0137             inputTracks="tracks",
0138             inputParticles="truth_seeds_selected",
0139             inputTrackParticleMatching="track_particle_matching",
0140             filePath=str(outputDir / "tracksummary_gsf.root"),
0141             writeGsfSpecific=True,
0142         )
0143     )
0144
0145     s.addWriter(
0146         acts.examples.TrackFitterPerformanceWriter(
0147             level=acts.logging.INFO,
0148             inputTracks="tracks",
0149             inputParticles="truth_seeds_selected",
0150             inputTrackParticleMatching="track_particle_matching",
0151             filePath=str(outputDir / "performance_gsf.root"),
0152         )
0153     )
0154
0155     return s
0156
0157
0158 if "__main__" == __name__:
0159     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0160
0161     detector, trackingGeometry, decorators = acts.examples.GenericDetector.create()
0162
0163     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0164
0165     inputParticlePath = Path("particles.root")
0166     if not inputParticlePath.exists():
0167         inputParticlePath = None
0168
0169     runTruthTrackingGsf(
0170         trackingGeometry=trackingGeometry,
0171         field=field,
0172         digiConfigFile=srcdir
0173         / "Examples/Algorithms/Digitization/share/default-smearing-config-generic.json",
0174         inputParticlePath=inputParticlePath,
0175         outputDir=Path.cwd(),
0176     ).run()