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         EtaConfig,
0027         PhiConfig,
0028         ParticleConfig,
0029         addFatras,
0030         addDigitization,
0031     )
0032
0033     from acts.examples.reconstruction import (
0034         addSeeding,
0035         TruthSeedRanges,
0036         ParticleSmearingSigmas,
0037         SeedFinderConfigArg,
0038         SeedFinderOptionsArg,
0039         SeedingAlgorithm,
0040         TruthEstimatedSeedingAlgorithmConfigArg,
0041         addCKFTracks,
0042     )
0043
0044     s = s or acts.examples.Sequencer(
0045         events=100, numThreads=-1, logLevel=acts.logging.INFO
0046     )
0047     for d in decorators:
0048         s.addContextDecorator(d)
0049     rnd = acts.examples.RandomNumbers(seed=42)
0050     outputDir = Path(outputDir)
0051
0052     if inputParticlePath is None:
0053         addParticleGun(
0054             s,
0055             EtaConfig(-2.0, 2.0),
0056             ParticleConfig(4, acts.PdgParticle.eMuon, True),
0057             PhiConfig(0.0, 360.0 * u.degree),
0058             multiplicity=2,
0059             rnd=rnd,
0060         )
0061     else:
0062         acts.logging.getLogger("CKFExample").info(
0063             "Reading particles from %s", inputParticlePath.resolve()
0064         )
0065         assert inputParticlePath.exists()
0066         s.addReader(
0067             RootParticleReader(
0068                 level=acts.logging.INFO,
0069                 filePath=str(inputParticlePath.resolve()),
0070                 outputParticles="particles_input",
0071             )
0072         )
0073
0074     addFatras(
0075         s,
0076         trackingGeometry,
0077         field,
0078         rnd=rnd,
0079     )
0080
0081     addDigitization(
0082         s,
0083         trackingGeometry,
0084         field,
0085         digiConfigFile=digiConfigFile,
0086         rnd=rnd,
0087     )
0088
0089     addSeeding(
0090         s,
0091         trackingGeometry,
0092         field,
0093         TruthSeedRanges(pt=(500.0 * u.MeV, None), nHits=(9, None)),
0094         ParticleSmearingSigmas(pRel=0.01),  # only used by SeedingAlgorithm.TruthSmeared
0095         SeedFinderConfigArg(
0096             r=(None, 200 * u.mm),  # rMin=default, 33mm
0097             deltaR=(1 * u.mm, 60 * u.mm),
0098             collisionRegion=(-250 * u.mm, 250 * u.mm),
0099             z=(-2000 * u.mm, 2000 * u.mm),
0100             maxSeedsPerSpM=1,
0101             sigmaScattering=5,
0102             radLengthPerSeed=0.1,
0103             minPt=500 * u.MeV,
0104             impactMax=3 * u.mm,
0105         ),
0106         SeedFinderOptionsArg(bFieldInZ=2 * u.T, beamPos=(0.0, 0.0)),
0107         TruthEstimatedSeedingAlgorithmConfigArg(deltaR=(10.0 * u.mm, None)),
0108         seedingAlgorithm=SeedingAlgorithm.TruthSmeared
0109         if truthSmearedSeeded
0110         else SeedingAlgorithm.TruthEstimated
0111         if truthEstimatedSeeded
0112         else SeedingAlgorithm.Default,
0113         geoSelectionConfigFile=geometrySelection,
0114         outputDirRoot=outputDir,
0115         rnd=rnd,  # only used by SeedingAlgorithm.TruthSmeared
0116     )
0117
0118     addCKFTracks(
0119         s,
0120         trackingGeometry,
0121         field,
0122         outputDirRoot=outputDir,
0123         outputDirCsv=outputDir / "csv" if outputCsv else None,
0124     )
0125
0126     return s
0127
0128
0129 if "__main__" == __name__:
0130     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0131
0132     detector, trackingGeometry, decorators = GenericDetector.create()
0133
0134     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0135
0136     inputParticlePath = Path("particles.root")
0137     if not inputParticlePath.exists():
0138         inputParticlePath = None
0139
0140     runCKFTracks(
0141         trackingGeometry,
0142         decorators,
0143         field=field,
0144         geometrySelection=srcdir
0145         / "Examples/Algorithms/TrackFinding/share/geoSelection-genericDetector.json",
0146         digiConfigFile=srcdir
0147         / "Examples/Algorithms/Digitization/share/default-smearing-config-generic.json",
0148         truthSmearedSeeded=False,
0149         truthEstimatedSeeded=False,
0150         inputParticlePath=inputParticlePath,
0151         outputDir=Path.cwd(),
0152         outputCsv=True,
0153     ).run()