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