Scripts/Python/truth_tracking_gx2f.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 runTruthTrackingGx2f(
0013     trackingGeometry: acts.TrackingGeometry,
0014     field: acts.MagneticFieldProvider,
0015     digiConfigFile: Path,
0016     outputDir: Path,
0017     inputParticlePath: Optional[Path] = None,
0018     s: acts.examples.Sequencer = None,
0019 ):
0020     from acts.examples.simulation import (
0021         addParticleGun,
0022         ParticleConfig,
0023         EtaConfig,
0024         PhiConfig,
0025         MomentumConfig,
0026         addFatras,
0027         addDigitization,
0028         ParticleSelectorConfig,
0029         addDigiParticleSelection,
0030     )
0031     from acts.examples.reconstruction import (
0032         addSeeding,
0033         SeedingAlgorithm,
0034         addGx2fTracks,
0035     )
0036
0037     s = s or acts.examples.Sequencer(
0038         events=100, numThreads=-1, logLevel=acts.logging.INFO
0039     )
0040
0041     rnd = acts.examples.RandomNumbers(seed=42)
0042     outputDir = Path(outputDir)
0043
0044     if inputParticlePath is None:
0045         addParticleGun(
0046             s,
0047             ParticleConfig(num=1, pdg=acts.PdgParticle.eMuon, randomizeCharge=True),
0048             EtaConfig(-3.0, 3.0, uniform=True),
0049             MomentumConfig(1.0 * u.GeV, 100.0 * u.GeV, transverse=True),
0050             PhiConfig(0.0, 360.0 * u.degree),
0051             vtxGen=acts.examples.GaussianVertexGenerator(
0052                 mean=acts.Vector4(0, 0, 0, 0),
0053                 stddev=acts.Vector4(0, 0, 0, 0),
0054             ),
0055             multiplicity=1,
0056             rnd=rnd,
0057             outputDirRoot=outputDir,
0058         )
0059     else:
0060         acts.logging.getLogger("Truth tracking example").info(
0061             "Reading particles from %s", inputParticlePath.resolve()
0062         )
0063         assert inputParticlePath.exists()
0064         s.addReader(
0065             acts.examples.RootParticleReader(
0066                 level=acts.logging.INFO,
0067                 filePath=str(inputParticlePath.resolve()),
0068                 outputParticles="particles_generated",
0069             )
0070         )
0071
0072     addFatras(
0073         s,
0074         trackingGeometry,
0075         field,
0076         rnd=rnd,
0077         enableInteractions=True,
0078     )
0079
0080     addDigitization(
0081         s,
0082         trackingGeometry,
0083         field,
0084         digiConfigFile=digiConfigFile,
0085         rnd=rnd,
0086     )
0087
0088     addDigiParticleSelection(
0089         s,
0090         ParticleSelectorConfig(
0091             pt=(0.9 * u.GeV, None),
0092             measurements=(7, None),
0093             removeNeutral=True,
0094             removeSecondaries=True,
0095         ),
0096     )
0097
0098     addSeeding(
0099         s,
0100         trackingGeometry,
0101         field,
0102         rnd=rnd,
0103         inputParticles="particles_generated",
0104         seedingAlgorithm=SeedingAlgorithm.TruthSmeared,
0105         particleHypothesis=acts.ParticleHypothesis.muon,
0106     )
0107
0108     addGx2fTracks(
0109         s,
0110         trackingGeometry,
0111         field,
0112         nUpdateMax=17,
0113         relChi2changeCutOff=1e-7,
0114         multipleScattering=True,
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_gx2f.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_gx2f.root"),
0148             writeGx2fSpecific=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_gx2f.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     runTruthTrackingGx2f(
0186         trackingGeometry=trackingGeometry,
0187         field=field,
0188         digiConfigFile=digiConfigFile,
0189         outputDir=Path.cwd(),
0190     ).run()