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         addFatras,
0031         addDigitization,
0032         ParticleSelectorConfig,
0033         addDigiParticleSelection,
0034     )
0035
0036     from acts.examples.reconstruction import (
0037         addSeeding,
0038         TrackSmearingSigmas,
0039         SeedFinderConfigArg,
0040         SeedFinderOptionsArg,
0041         SeedingAlgorithm,
0042         TruthEstimatedSeedingAlgorithmConfigArg,
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.logging.getLogger("CKFExample").info(
0068             "Reading particles from %s", 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 eveything 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         TruthEstimatedSeedingAlgorithmConfigArg(deltaR=(10.0 * u.mm, None)),
0133         seedingAlgorithm=(
0134             SeedingAlgorithm.TruthSmeared
0135             if truthSmearedSeeded
0136             else (
0137                 SeedingAlgorithm.TruthEstimated
0138                 if truthEstimatedSeeded
0139                 else SeedingAlgorithm.Default
0140             )
0141         ),
0142         initialSigmas=[
0143             1 * u.mm,
0144             1 * u.mm,
0145             1 * u.degree,
0146             1 * u.degree,
0147             0 * u.e / u.GeV,
0148             1 * u.ns,
0149         ],
0150         initialSigmaQoverPt=0.1 * u.e / u.GeV,
0151         initialSigmaPtRel=0.1,
0152         initialVarInflation=[1.0] * 6,
0153         geoSelectionConfigFile=geometrySelection,
0154         outputDirRoot=outputDir,
0155         rnd=rnd,  # only used by SeedingAlgorithm.TruthSmeared
0156     )
0157
0158     addCKFTracks(
0159         s,
0160         trackingGeometry,
0161         field,
0162         TrackSelectorConfig(
0163             pt=(500 * u.MeV, None),
0164             absEta=(None, 3.0),
0165             loc0=(-4.0 * u.mm, 4.0 * u.mm),
0166             nMeasurementsMin=7,
0167             maxHoles=2,
0168             maxOutliers=2,
0169         ),
0170         CkfConfig(
0171             chi2CutOffMeasurement=15.0,
0172             chi2CutOffOutlier=25.0,
0173             numMeasurementsCutOff=10,
0174             seedDeduplication=True if not truthSmearedSeeded else False,
0175             stayOnSeed=True if not truthSmearedSeeded else False,
0176         ),
0177         outputDirRoot=outputDir,
0178         outputDirCsv=outputDir / "csv" if outputCsv else None,
0179         writeTrackStates=True,
0180     )
0181
0182     return s
0183
0184
0185 if "__main__" == __name__:
0186     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0187
0188     detector = GenericDetector()
0189     trackingGeometry = detector.trackingGeometry()
0190     decorators = detector.contextDecorators()
0191
0192     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0193
0194     inputParticlePath = Path("particles.root")
0195     if not inputParticlePath.exists():
0196         inputParticlePath = None
0197
0198     runCKFTracks(
0199         trackingGeometry,
0200         decorators,
0201         field=field,
0202         geometrySelection=srcdir / "Examples/Configs/generic-seeding-config.json",
0203         digiConfigFile=srcdir / "Examples/Configs/generic-digi-smearing-config.json",
0204         truthSmearedSeeded=False,
0205         truthEstimatedSeeded=False,
0206         inputParticlePath=inputParticlePath,
0207         outputDir=Path.cwd(),
0208         outputCsv=True,
0209     ).run()