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