Scripts/Optimization/ckf.py

0001 #!/usr/bin/env python3
0002
0003 from pathlib import Path
0004 from typing import Optional
0005 import argparse
0006
0007 import acts
0008 from acts import UnitConstants as u
0009 from acts.examples import GenericDetector, RootParticleReader
0010
0011
0012 def getArgumentParser():
0013     """Get arguments from command line"""
0014     parser = argparse.ArgumentParser(description="Command line arguments for CKF")
0015     parser.add_argument(
0016         "-i",
0017         "--indir",
0018         dest="indir",
0019         help="Directory with input root files",
0020         default="./",
0021     )
0022     parser.add_argument(
0023         "-o",
0024         "--output",
0025         dest="outdir",
0026         help="Output directory for new ntuples",
0027         default="./",
0028     )
0029     parser.add_argument(
0030         "-n", "--nEvents", dest="nEvts", help="Number of events to run over", default=1
0031     )
0032     parser.add_argument(
0033         "--sf_maxSeedsPerSpM",
0034         dest="sf_maxSeedsPerSpM",
0035         help="Number of compatible seeds considered for middle seed",
0036         type=int,
0037         default=1,
0038     )
0039     parser.add_argument(
0040         "--sf_cotThetaMax",
0041         dest="sf_cotThetaMax",
0042         help="cot of maximum theta angle",
0043         type=float,
0044         default=7.40627,
0045     )
0046     parser.add_argument(
0047         "--sf_sigmaScattering",
0048         dest="sf_sigmaScattering",
0049         help="How many sigmas of scattering to include in seeds",
0050         type=float,
0051         default=5,
0052     )
0053     parser.add_argument(
0054         "--sf_radLengthPerSeed",
0055         dest="sf_radLengthPerSeed",
0056         help="Average Radiation Length",
0057         type=float,
0058         default=0.1,
0059     )
0060     parser.add_argument(
0061         "--sf_impactMax",
0062         dest="sf_impactMax",
0063         help="max impact parameter in mm",
0064         type=float,
0065         default=3.0,
0066     )
0067     parser.add_argument(
0068         "--sf_maxPtScattering",
0069         dest="sf_maxPtScattering",
0070         help="maximum Pt for scattering cut in GeV",
0071         type=float,
0072         default=10.0,
0073     )
0074     parser.add_argument(
0075         "--sf_deltaRMin",
0076         dest="sf_deltaRMin",
0077         help="minimum value for deltaR separation in mm",
0078         type=float,
0079         default=1.0,
0080     )
0081     parser.add_argument(
0082         "--sf_deltaRMax",
0083         dest="sf_deltaRMax",
0084         help="maximum value for deltaR separation in mm",
0085         type=float,
0086         default=60.0,
0087     )
0088
0089     return parser
0090
0091
0092 def runCKFTracks(
0093     trackingGeometry,
0094     decorators,
0095     geometrySelection: Path,
0096     digiConfigFile: Path,
0097     field,
0098     outputDir: Path,
0099     NumEvents=1,
0100     truthSmearedSeeded=False,
0101     truthEstimatedSeeded=False,
0102     outputCsv=True,
0103     inputParticlePath: Optional[Path] = None,
0104     s=None,
0105     MaxSeedsPerSpM=1,
0106     CotThetaMax=7.40627,
0107     SigmaScattering=5,
0108     RadLengthPerSeed=0.1,
0109     ImpactMax=3.0,
0110     MaxPtScattering=10.0,
0111     DeltaRMin=1.0,
0112     DeltaRMax=60.0,
0113 ):
0114     from acts.examples.simulation import (
0115         addParticleGun,
0116         EtaConfig,
0117         PhiConfig,
0118         ParticleConfig,
0119         addFatras,
0120         addDigitization,
0121     )
0122
0123     from acts.examples.reconstruction import (
0124         addSeeding,
0125         TruthSeedRanges,
0126         ParticleSmearingSigmas,
0127         SeedFinderConfigArg,
0128         SeedFinderOptionsArg,
0129         SeedingAlgorithm,
0130         TruthEstimatedSeedingAlgorithmConfigArg,
0131         addCKFTracks,
0132     )
0133
0134     s = s or acts.examples.Sequencer(
0135         events=int(NumEvents),
0136         numThreads=-1,
0137         logLevel=acts.logging.INFO,
0138         outputDir=outputDir,
0139     )
0140     for d in decorators:
0141         s.addContextDecorator(d)
0142     rnd = acts.examples.RandomNumbers(seed=42)
0143     outputDir = Path(outputDir)
0144
0145     if inputParticlePath is None:
0146         addParticleGun(
0147             s,
0148             EtaConfig(-2.0, 2.0),
0149             ParticleConfig(4, acts.PdgParticle.eMuon, True),
0150             PhiConfig(0.0, 360.0 * u.degree),
0151             multiplicity=2,
0152             rnd=rnd,
0153         )
0154     else:
0155         acts.logging.getLogger("CKFExample").info(
0156             "Reading particles from %s", inputParticlePath.resolve()
0157         )
0158         assert inputParticlePath.exists()
0159         s.addReader(
0160             RootParticleReader(
0161                 level=acts.logging.INFO,
0162                 filePath=str(inputParticlePath.resolve()),
0163                 outputParticles="particles_input",
0164             )
0165         )
0166
0167     addFatras(
0168         s,
0169         trackingGeometry,
0170         field,
0171         rnd=rnd,
0172     )
0173
0174     addDigitization(
0175         s,
0176         trackingGeometry,
0177         field,
0178         digiConfigFile=digiConfigFile,
0179         rnd=rnd,
0180     )
0181
0182     addSeeding(
0183         s,
0184         trackingGeometry,
0185         field,
0186         TruthSeedRanges(pt=(500.0 * u.MeV, None), nHits=(9, None)),
0187         ParticleSmearingSigmas(pRel=0.01),  # only used by SeedingAlgorithm.TruthSmeared
0188         SeedFinderConfigArg(
0189             r=(None, 200 * u.mm),  # rMin=default, 33mm
0190             deltaR=(DeltaRMin * u.mm, DeltaRMax * u.mm),
0191             collisionRegion=(-250 * u.mm, 250 * u.mm),
0192             z=(-2000 * u.mm, 2000 * u.mm),
0193             maxSeedsPerSpM=MaxSeedsPerSpM,
0194             cotThetaMax=CotThetaMax,
0195             sigmaScattering=SigmaScattering,
0196             radLengthPerSeed=RadLengthPerSeed,
0197             maxPtScattering=MaxPtScattering * u.GeV,
0198             minPt=500 * u.MeV,
0199             impactMax=ImpactMax * u.mm,
0200         ),
0201         SeedFinderOptionsArg(bFieldInZ=2 * u.T, beamPos=(0.0, 0, 0)),
0202         TruthEstimatedSeedingAlgorithmConfigArg(deltaR=(10.0 * u.mm, None)),
0203         seedingAlgorithm=SeedingAlgorithm.TruthSmeared
0204         if truthSmearedSeeded
0205         else SeedingAlgorithm.TruthEstimated
0206         if truthEstimatedSeeded
0207         else SeedingAlgorithm.Default,
0208         geoSelectionConfigFile=geometrySelection,
0209         outputDirRoot=outputDir,
0210         rnd=rnd,  # only used by SeedingAlgorithm.TruthSmeared
0211     )
0212
0213     addCKFTracks(
0214         s,
0215         trackingGeometry,
0216         field,
0217         outputDirRoot=outputDir,
0218         outputDirCsv=outputDir / "csv" if outputCsv else None,
0219     )
0220
0221     return s
0222
0223
0224 if "__main__" == __name__:
0225     options = getArgumentParser().parse_args()
0226
0227     Inputdir = options.indir
0228     Outputdir = options.outdir
0229
0230     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0231
0232     detector, trackingGeometry, decorators = GenericDetector.create()
0233
0234     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0235
0236     inputParticlePath = Path(Inputdir) / "pythia8_particles.root"
0237     if not inputParticlePath.exists():
0238         inputParticlePath = None
0239
0240     runCKFTracks(
0241         trackingGeometry,
0242         decorators,
0243         field=field,
0244         geometrySelection=srcdir
0245         / "Examples/Algorithms/TrackFinding/share/geoSelection-genericDetector.json",
0246         digiConfigFile=srcdir
0247         / "Examples/Algorithms/Digitization/share/default-smearing-config-generic.json",
0248         outputCsv=True,
0249         truthSmearedSeeded=False,
0250         truthEstimatedSeeded=False,
0251         inputParticlePath=inputParticlePath,
0252         outputDir=Outputdir,
0253         NumEvents=options.nEvts,
0254         MaxSeedsPerSpM=options.sf_maxSeedsPerSpM,
0255         CotThetaMax=options.sf_cotThetaMax,
0256         SigmaScattering=options.sf_sigmaScattering,
0257         RadLengthPerSeed=options.sf_radLengthPerSeed,
0258         ImpactMax=options.sf_impactMax,
0259         MaxPtScattering=options.sf_maxPtScattering,
0260         DeltaRMin=options.sf_deltaRMin,
0261         DeltaRMax=options.sf_deltaRMax,
0262     ).run()