File indexing completed on 2025-01-18 09:12:07
0001
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 ParticleSelectorConfig,
0120 addFatras,
0121 addDigitization,
0122 )
0123
0124 from acts.examples.reconstruction import (
0125 addSeeding,
0126 TrackSmearingSigmas,
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 postSelectParticles=ParticleSelectorConfig(
0173 pt=(0.5 * u.GeV, None),
0174 hits=(9, None),
0175 removeNeutral=True,
0176 ),
0177 )
0178
0179 addDigitization(
0180 s,
0181 trackingGeometry,
0182 field,
0183 digiConfigFile=digiConfigFile,
0184 rnd=rnd,
0185 )
0186
0187 addSeeding(
0188 s,
0189 trackingGeometry,
0190 field,
0191 TrackSmearingSigmas(
0192
0193 loc0=0,
0194 loc0PtA=0,
0195 loc0PtB=0,
0196 loc1=0,
0197 loc1PtA=0,
0198 loc1PtB=0,
0199 time=0,
0200 phi=0,
0201 theta=0,
0202 ptRel=0,
0203 ),
0204 SeedFinderConfigArg(
0205 r=(None, 200 * u.mm),
0206 deltaR=(DeltaRMin * u.mm, DeltaRMax * u.mm),
0207 collisionRegion=(-250 * u.mm, 250 * u.mm),
0208 z=(-2000 * u.mm, 2000 * u.mm),
0209 maxSeedsPerSpM=MaxSeedsPerSpM,
0210 cotThetaMax=CotThetaMax,
0211 sigmaScattering=SigmaScattering,
0212 radLengthPerSeed=RadLengthPerSeed,
0213 maxPtScattering=MaxPtScattering * u.GeV,
0214 minPt=500 * u.MeV,
0215 impactMax=ImpactMax * u.mm,
0216 ),
0217 SeedFinderOptionsArg(bFieldInZ=2 * u.T, beamPos=(0.0, 0, 0)),
0218 TruthEstimatedSeedingAlgorithmConfigArg(deltaR=(10.0 * u.mm, None)),
0219 seedingAlgorithm=(
0220 SeedingAlgorithm.TruthSmeared
0221 if truthSmearedSeeded
0222 else (
0223 SeedingAlgorithm.TruthEstimated
0224 if truthEstimatedSeeded
0225 else SeedingAlgorithm.Default
0226 )
0227 ),
0228 initialSigmas=[
0229 1 * u.mm,
0230 1 * u.mm,
0231 1 * u.degree,
0232 1 * u.degree,
0233 0.1 * u.e / u.GeV,
0234 1 * u.ns,
0235 ],
0236 initialSigmaPtRel=0.01,
0237 initialVarInflation=[1.0] * 6,
0238 geoSelectionConfigFile=geometrySelection,
0239 outputDirRoot=outputDir,
0240 rnd=rnd,
0241 )
0242
0243 addCKFTracks(
0244 s,
0245 trackingGeometry,
0246 field,
0247 outputDirRoot=outputDir,
0248 outputDirCsv=outputDir / "csv" if outputCsv else None,
0249 )
0250
0251 return s
0252
0253
0254 if "__main__" == __name__:
0255 options = getArgumentParser().parse_args()
0256
0257 Inputdir = options.indir
0258 Outputdir = options.outdir
0259
0260 srcdir = Path(__file__).resolve().parent.parent.parent.parent
0261
0262 detector = GenericDetector()
0263 trackingGeometry = detector.trackingGeometry()
0264 decorators = detector.contextDecorators()
0265
0266 field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0267
0268 inputParticlePath = Path(Inputdir) / "particles.root"
0269 if not inputParticlePath.exists():
0270 inputParticlePath = None
0271
0272 runCKFTracks(
0273 trackingGeometry,
0274 decorators,
0275 field=field,
0276 geometrySelection=srcdir
0277 / "Examples/Algorithms/TrackFinding/share/geoSelection-genericDetector.json",
0278 digiConfigFile=srcdir
0279 / "Examples/Algorithms/Digitization/share/default-smearing-config-generic.json",
0280 outputCsv=True,
0281 truthSmearedSeeded=False,
0282 truthEstimatedSeeded=False,
0283 inputParticlePath=inputParticlePath,
0284 outputDir=Outputdir,
0285 NumEvents=options.nEvts,
0286 MaxSeedsPerSpM=options.sf_maxSeedsPerSpM,
0287 CotThetaMax=options.sf_cotThetaMax,
0288 SigmaScattering=options.sf_sigmaScattering,
0289 RadLengthPerSeed=options.sf_radLengthPerSeed,
0290 ImpactMax=options.sf_impactMax,
0291 MaxPtScattering=options.sf_maxPtScattering,
0292 DeltaRMin=options.sf_deltaRMin,
0293 DeltaRMax=options.sf_deltaRMax,
0294 ).run()
