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         TrackSmearingSigmas,
0036         addTruthTrackingGsf,
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.eElectron, 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("GSF 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_generated",
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     addDigiParticleSelection(
0093         s,
0094         ParticleSelectorConfig(
0095             pt=(0.9 * u.GeV, None),
0096             measurements=(7, None),
0097             removeNeutral=True,
0098             removeSecondaries=True,
0099         ),
0100     )
0101
0102     addSeeding(
0103         s,
0104         trackingGeometry,
0105         field,
0106         rnd=rnd,
0107         inputParticles="particles_generated",
0108         seedingAlgorithm=SeedingAlgorithm.TruthSmeared,
0109         trackSmearingSigmas=TrackSmearingSigmas(
0110             # zero eveything so the GSF has a chance to find the measurements
0111             loc0=0,
0112             loc0PtA=0,
0113             loc0PtB=0,
0114             loc1=0,
0115             loc1PtA=0,
0116             loc1PtB=0,
0117             time=0,
0118             phi=0,
0119             theta=0,
0120             ptRel=0,
0121         ),
0122         particleHypothesis=acts.ParticleHypothesis.electron,
0123         initialSigmas=[
0124             1 * u.mm,
0125             1 * u.mm,
0126             1 * u.degree,
0127             1 * u.degree,
0128             0 / u.GeV,
0129             1 * u.ns,
0130         ],
0131         initialSigmaQoverPt=0.1 / u.GeV,
0132         initialSigmaPtRel=0.1,
0133         initialVarInflation=[1e0, 1e0, 1e0, 1e0, 1e0, 1e0],
0134     )
0135
0136     addTruthTrackingGsf(
0137         s,
0138         trackingGeometry,
0139         field,
0140     )
0141
0142     s.addAlgorithm(
0143         acts.examples.TrackSelectorAlgorithm(
0144             level=acts.logging.INFO,
0145             inputTracks="tracks",
0146             outputTracks="selected-tracks",
0147             selectorConfig=acts.TrackSelector.Config(
0148                 minMeasurements=7,
0149             ),
0150         )
0151     )
0152     s.addWhiteboardAlias("tracks", "selected-tracks")
0153
0154     s.addWriter(
0155         acts.examples.RootTrackStatesWriter(
0156             level=acts.logging.INFO,
0157             inputTracks="tracks",
0158             inputParticles="particles_selected",
0159             inputTrackParticleMatching="track_particle_matching",
0160             inputSimHits="simhits",
0161             inputMeasurementSimHitsMap="measurement_simhits_map",
0162             filePath=str(outputDir / "trackstates_gsf.root"),
0163         )
0164     )
0165
0166     s.addWriter(
0167         acts.examples.RootTrackSummaryWriter(
0168             level=acts.logging.INFO,
0169             inputTracks="tracks",
0170             inputParticles="particles_selected",
0171             inputTrackParticleMatching="track_particle_matching",
0172             filePath=str(outputDir / "tracksummary_gsf.root"),
0173             writeGsfSpecific=True,
0174         )
0175     )
0176
0177     s.addWriter(
0178         acts.examples.TrackFitterPerformanceWriter(
0179             level=acts.logging.INFO,
0180             inputTracks="tracks",
0181             inputParticles="particles_selected",
0182             inputTrackParticleMatching="track_particle_matching",
0183             filePath=str(outputDir / "performance_gsf.root"),
0184         )
0185     )
0186
0187     return s
0188
0189
0190 if "__main__" == __name__:
0191     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0192
0193     # ODD
0194     from acts.examples.odd import getOpenDataDetector
0195
0196     detector = getOpenDataDetector()
0197     trackingGeometry = detector.trackingGeometry()
0198     digiConfigFile = srcdir / "Examples/Configs/odd-digi-smearing-config.json"
0199
0200     ## GenericDetector
0201     # detector = acts.examples.GenericDetector()
0202     # trackingGeometry = detector.trackingGeometry()
0203     # digiConfigFile = (
0204     #     srcdir
0205     #     / "Examples/Configs/generic-digi-smearing-config.json"
0206     # )
0207
0208     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0209
0210     runTruthTrackingGsf(
0211         trackingGeometry=trackingGeometry,
0212         field=field,
0213         digiConfigFile=digiConfigFile,
0214         outputDir=Path.cwd(),
0215     ).run()