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 from acts.examples.root import (
0009     RootParticleReader,
0010     RootSimHitReader,
0011 )
0012
0013 u = acts.UnitConstants
0014
0015
0016 def runTruthTrackingGsf(
0017     trackingGeometry: acts.TrackingGeometry,
0018     field: acts.MagneticFieldProvider,
0019     digiConfigFile: Path,
0020     outputDir: Path,
0021     inputParticlePath: Optional[Path] = None,
0022     inputSimHitsPath: Optional[Path] = None,
0023     decorators=[],
0024     s: acts.examples.Sequencer = None,
0025 ):
0026     from acts.examples.simulation import (
0027         addParticleGun,
0028         ParticleConfig,
0029         EtaConfig,
0030         PhiConfig,
0031         MomentumConfig,
0032         addFatras,
0033         addDigitization,
0034         ParticleSelectorConfig,
0035         addDigiParticleSelection,
0036     )
0037     from acts.examples.reconstruction import (
0038         addSeeding,
0039         SeedingAlgorithm,
0040         TrackSmearingSigmas,
0041         addTruthTrackingGsf,
0042     )
0043     from acts.examples.root import (
0044         RootParticleReader,
0045         RootTrackStatesWriter,
0046         RootTrackSummaryWriter,
0047         RootTrackFitterPerformanceWriter,
0048     )
0049
0050     s = s or acts.examples.Sequencer(
0051         events=100, numThreads=-1, logLevel=acts.logging.INFO
0052     )
0053
0054     for d in decorators:
0055         s.addContextDecorator(d)
0056
0057     rnd = acts.examples.RandomNumbers(seed=42)
0058     outputDir = Path(outputDir)
0059     logger = acts.logging.getLogger("GSF Example")
0060
0061     if inputParticlePath is None:
0062         addParticleGun(
0063             s,
0064             ParticleConfig(num=1, pdg=acts.PdgParticle.eElectron, randomizeCharge=True),
0065             EtaConfig(-3.0, 3.0, uniform=True),
0066             MomentumConfig(1.0 * u.GeV, 100.0 * u.GeV, transverse=True),
0067             PhiConfig(0.0, 360.0 * u.degree),
0068             vtxGen=acts.examples.GaussianVertexGenerator(
0069                 mean=acts.Vector4(0, 0, 0, 0),
0070                 stddev=acts.Vector4(0, 0, 0, 0),
0071             ),
0072             multiplicity=1,
0073             rnd=rnd,
0074         )
0075     else:
0076         logger.info("Reading particles from %s", inputParticlePath.resolve())
0077         assert inputParticlePath.exists()
0078         s.addReader(
0079             RootParticleReader(
0080                 level=acts.logging.INFO,
0081                 filePath=str(inputParticlePath.resolve()),
0082                 outputParticles="particles_generated",
0083             )
0084         )
0085         s.addWhiteboardAlias("particles", "particles_generated")
0086
0087     if inputSimHitsPath is None:
0088         addFatras(
0089             s,
0090             trackingGeometry,
0091             field,
0092             rnd=rnd,
0093             enableInteractions=True,
0094         )
0095     else:
0096         logger.info("Reading hits from %s", inputSimHitsPath.resolve())
0097         s.addReader(
0098             RootSimHitReader(
0099                 level=acts.logging.INFO,
0100                 filePath=str(inputSimHitsPath.resolve()),
0101                 outputSimHits="simhits",
0102             )
0103         )
0104         s.addWhiteboardAlias("particles_simulated_selected", "particles_generated")
0105
0106     addDigitization(
0107         s,
0108         trackingGeometry,
0109         field,
0110         digiConfigFile=digiConfigFile,
0111         rnd=rnd,
0112     )
0113
0114     addDigiParticleSelection(
0115         s,
0116         ParticleSelectorConfig(
0117             pt=(0.9 * u.GeV, None),
0118             measurements=(7, None),
0119             removeNeutral=True,
0120             removeSecondaries=True,
0121         ),
0122     )
0123
0124     addSeeding(
0125         s,
0126         trackingGeometry,
0127         field,
0128         rnd=rnd,
0129         inputParticles="particles_generated",
0130         seedingAlgorithm=SeedingAlgorithm.TruthSmeared,
0131         trackSmearingSigmas=TrackSmearingSigmas(
0132             # zero eveything so the GSF has a chance to find the measurements
0133             loc0=0,
0134             loc0PtA=0,
0135             loc0PtB=0,
0136             loc1=0,
0137             loc1PtA=0,
0138             loc1PtB=0,
0139             time=0,
0140             phi=0,
0141             theta=0,
0142             ptRel=0,
0143         ),
0144         particleHypothesis=acts.ParticleHypothesis.electron,
0145         initialSigmas=[
0146             1 * u.mm,
0147             1 * u.mm,
0148             1 * u.degree,
0149             1 * u.degree,
0150             0 / u.GeV,
0151             1 * u.ns,
0152         ],
0153         initialSigmaQoverPt=0.1 / u.GeV,
0154         initialSigmaPtRel=0.1,
0155         initialVarInflation=[1e0, 1e0, 1e0, 1e0, 1e0, 1e0],
0156     )
0157
0158     addTruthTrackingGsf(
0159         s,
0160         trackingGeometry,
0161         field,
0162     )
0163
0164     s.addAlgorithm(
0165         acts.examples.TrackSelectorAlgorithm(
0166             level=acts.logging.INFO,
0167             inputTracks="tracks",
0168             outputTracks="selected-tracks",
0169             selectorConfig=acts.TrackSelector.Config(
0170                 minMeasurements=7,
0171             ),
0172         )
0173     )
0174     s.addWhiteboardAlias("tracks", "selected-tracks")
0175
0176     s.addWriter(
0177         RootTrackStatesWriter(
0178             level=acts.logging.INFO,
0179             inputTracks="tracks",
0180             inputParticles="particles_selected",
0181             inputTrackParticleMatching="track_particle_matching",
0182             inputSimHits="simhits",
0183             inputMeasurementSimHitsMap="measurement_simhits_map",
0184             filePath=str(outputDir / "trackstates_gsf.root"),
0185         )
0186     )
0187
0188     s.addWriter(
0189         RootTrackSummaryWriter(
0190             level=acts.logging.INFO,
0191             inputTracks="tracks",
0192             inputParticles="particles_selected",
0193             inputTrackParticleMatching="track_particle_matching",
0194             filePath=str(outputDir / "tracksummary_gsf.root"),
0195             writeGsfSpecific=True,
0196         )
0197     )
0198
0199     s.addWriter(
0200         RootTrackFitterPerformanceWriter(
0201             level=acts.logging.INFO,
0202             inputTracks="tracks",
0203             inputParticles="particles_selected",
0204             inputTrackParticleMatching="track_particle_matching",
0205             filePath=str(outputDir / "performance_gsf.root"),
0206         )
0207     )
0208
0209     return s
0210
0211
0212 if "__main__" == __name__:
0213     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0214
0215     # ODD
0216     from acts.examples.odd import getOpenDataDetector
0217
0218     detector = getOpenDataDetector()
0219     trackingGeometry = detector.trackingGeometry()
0220     digiConfigFile = srcdir / "Examples/Configs/odd-digi-smearing-config.json"
0221
0222     ## GenericDetector
0223     # detector = acts.examples.GenericDetector()
0224     # trackingGeometry = detector.trackingGeometry()
0225     # digiConfigFile = (
0226     #     srcdir
0227     #     / "Examples/Configs/generic-digi-smearing-config.json"
0228     # )
0229
0230     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0231
0232     runTruthTrackingGsf(
0233         trackingGeometry=trackingGeometry,
0234         field=field,
0235         digiConfigFile=digiConfigFile,
0236         outputDir=Path.cwd(),
0237     ).run()