Scripts/Python/ckf_tracks.py

0001 #!/usr/bin/env python3
0002
0003 from pathlib import Path
0004 from typing import Optional
0005
0006 import acts
0007 from acts import UnitConstants as u
0008 from acts.examples import GenericDetector, RootParticleReader
0009
0010
0011 def runCKFTracks(
0012     trackingGeometry,
0013     decorators,
0014     geometrySelection: Path,
0015     digiConfigFile: Path,
0016     field,
0017     outputDir: Path,
0018     outputCsv=True,
0019     truthSmearedSeeded=False,
0020     truthEstimatedSeeded=False,
0021     inputParticlePath: Optional[Path] = None,
0022     s=None,
0023 ):
0024     from acts.examples.simulation import (
0025         addParticleGun,
0026         MomentumConfig,
0027         EtaConfig,
0028         PhiConfig,
0029         ParticleConfig,
0030         ParticleSelectorConfig,
0031         addFatras,
0032         addDigitization,
0033     )
0034
0035     from acts.examples.reconstruction import (
0036         addSeeding,
0037         TrackSmearingSigmas,
0038         SeedFinderConfigArg,
0039         SeedFinderOptionsArg,
0040         SeedingAlgorithm,
0041         TruthEstimatedSeedingAlgorithmConfigArg,
0042         addCKFTracks,
0043         TrackSelectorConfig,
0044         CkfConfig,
0045     )
0046
0047     s = s or acts.examples.Sequencer(
0048         events=100, numThreads=-1, logLevel=acts.logging.INFO
0049     )
0050     for d in decorators:
0051         s.addContextDecorator(d)
0052     rnd = acts.examples.RandomNumbers(seed=42)
0053     outputDir = Path(outputDir)
0054
0055     if inputParticlePath is None:
0056         addParticleGun(
0057             s,
0058             MomentumConfig(1 * u.GeV, 10 * u.GeV, transverse=True),
0059             EtaConfig(-2.0, 2.0, uniform=True),
0060             PhiConfig(0.0, 360.0 * u.degree),
0061             ParticleConfig(4, acts.PdgParticle.eMuon, randomizeCharge=True),
0062             multiplicity=2,
0063             rnd=rnd,
0064         )
0065     else:
0066         acts.logging.getLogger("CKFExample").info(
0067             "Reading particles from %s", inputParticlePath.resolve()
0068         )
0069         assert inputParticlePath.exists()
0070         s.addReader(
0071             RootParticleReader(
0072                 level=acts.logging.INFO,
0073                 filePath=str(inputParticlePath.resolve()),
0074                 outputParticles="particles_input",
0075             )
0076         )
0077
0078     addFatras(
0079         s,
0080         trackingGeometry,
0081         field,
0082         rnd=rnd,
0083         postSelectParticles=ParticleSelectorConfig(
0084             pt=(0.5 * u.GeV, None),
0085             hits=(9, None),
0086             removeNeutral=True,
0087         ),
0088     )
0089
0090     addDigitization(
0091         s,
0092         trackingGeometry,
0093         field,
0094         digiConfigFile=digiConfigFile,
0095         rnd=rnd,
0096     )
0097
0098     addSeeding(
0099         s,
0100         trackingGeometry,
0101         field,
0102         TrackSmearingSigmas(  # only used by SeedingAlgorithm.TruthSmeared
0103             # zero eveything so the CKF has a chance to find the measurements
0104             loc0=0,
0105             loc0PtA=0,
0106             loc0PtB=0,
0107             loc1=0,
0108             loc1PtA=0,
0109             loc1PtB=0,
0110             time=0,
0111             phi=0,
0112             theta=0,
0113             ptRel=0,
0114         ),
0115         SeedFinderConfigArg(
0116             r=(None, 200 * u.mm),  # rMin=default, 33mm
0117             deltaR=(1 * u.mm, 60 * u.mm),
0118             collisionRegion=(-250 * u.mm, 250 * u.mm),
0119             z=(-2000 * u.mm, 2000 * u.mm),
0120             maxSeedsPerSpM=1,
0121             sigmaScattering=5,
0122             radLengthPerSeed=0.1,
0123             minPt=500 * u.MeV,
0124             impactMax=3 * u.mm,
0125         ),
0126         SeedFinderOptionsArg(bFieldInZ=2 * u.T, beamPos=(0.0, 0.0)),
0127         TruthEstimatedSeedingAlgorithmConfigArg(deltaR=(10.0 * u.mm, None)),
0128         seedingAlgorithm=(
0129             SeedingAlgorithm.TruthSmeared
0130             if truthSmearedSeeded
0131             else (
0132                 SeedingAlgorithm.TruthEstimated
0133                 if truthEstimatedSeeded
0134                 else SeedingAlgorithm.Default
0135             )
0136         ),
0137         initialSigmas=[
0138             1 * u.mm,
0139             1 * u.mm,
0140             1 * u.degree,
0141             1 * u.degree,
0142             0.1 * u.e / u.GeV,
0143             1 * u.ns,
0144         ],
0145         initialSigmaPtRel=0.01,
0146         initialVarInflation=[1.0] * 6,
0147         geoSelectionConfigFile=geometrySelection,
0148         outputDirRoot=outputDir,
0149         rnd=rnd,  # only used by SeedingAlgorithm.TruthSmeared
0150     )
0151
0152     addCKFTracks(
0153         s,
0154         trackingGeometry,
0155         field,
0156         TrackSelectorConfig(
0157             pt=(500 * u.MeV, None),
0158             absEta=(None, 3.0),
0159             loc0=(-4.0 * u.mm, 4.0 * u.mm),
0160             nMeasurementsMin=7,
0161             maxHoles=2,
0162             maxOutliers=2,
0163         ),
0164         CkfConfig(
0165             chi2CutOffMeasurement=15.0,
0166             chi2CutOffOutlier=25.0,
0167             numMeasurementsCutOff=10,
0168             seedDeduplication=True if not truthSmearedSeeded else False,
0169             stayOnSeed=True if not truthSmearedSeeded else False,
0170         ),
0171         outputDirRoot=outputDir,
0172         outputDirCsv=outputDir / "csv" if outputCsv else None,
0173         writeTrackStates=True,
0174     )
0175
0176     return s
0177
0178
0179 if "__main__" == __name__:
0180     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0181
0182     detector = GenericDetector()
0183     trackingGeometry = detector.trackingGeometry()
0184     decorators = detector.contextDecorators()
0185
0186     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0187
0188     inputParticlePath = Path("particles.root")
0189     if not inputParticlePath.exists():
0190         inputParticlePath = None
0191
0192     runCKFTracks(
0193         trackingGeometry,
0194         decorators,
0195         field=field,
0196         geometrySelection=srcdir
0197         / "Examples/Algorithms/TrackFinding/share/geoSelection-genericDetector.json",
0198         digiConfigFile=srcdir
0199         / "Examples/Algorithms/Digitization/share/default-smearing-config-generic.json",
0200         truthSmearedSeeded=False,
0201         truthEstimatedSeeded=False,
0202         inputParticlePath=inputParticlePath,
0203         outputDir=Path.cwd(),
0204         outputCsv=True,
0205     ).run()