File indexing completed on 2026-05-19 07:35:08
0001
0002
0003 import argparse
0004 import pathlib
0005
0006 import acts
0007 import acts.examples
0008 from acts.examples.simulation import (
0009 addPythia8,
0010 ParticleSelectorConfig,
0011 addGenParticleSelection,
0012 addFatras,
0013 addDigitization,
0014 addDigiParticleSelection,
0015 )
0016
0017 from acts.examples.reconstruction import (
0018 addSeeding,
0019 addCKFTracks,
0020 TrackSelectorConfig,
0021 SeedingAlgorithm,
0022 )
0023
0024 from acts.examples.root import (
0025 RootSpacePointWriter,
0026 )
0027 from pathlib import Path
0028
0029 from typing import Optional
0030 from enum import Enum
0031
0032 DetectorName = Enum("DetectorName", "ODD generic")
0033
0034 HashingMetric = Enum("HashingMetric", "dphi dR")
0035
0036
0037 class Config:
0038 def __init__(
0039 self,
0040 mu: int = None,
0041 bucketSize: int = 100,
0042 maxSeedsPerSpM: int = 10,
0043 detector: DetectorName = DetectorName.generic,
0044 seedingAlgorithm: SeedingAlgorithm = SeedingAlgorithm.HashingPrototype,
0045 metric: str = HashingMetric.dphi,
0046 annoySeed: int = 123456789,
0047 zBins: int = 0,
0048 phiBins: int = 100,
0049 ):
0050 self.mu = mu
0051 self.bucketSize = bucketSize
0052 self.maxSeedsPerSpM = maxSeedsPerSpM
0053 self.detector = detector
0054 self.seedingAlgorithm = (
0055 seedingAlgorithm
0056 if type(seedingAlgorithm) == SeedingAlgorithm
0057 else SeedingAlgorithm[seedingAlgorithm]
0058 )
0059 self.metric = metric if type(metric) == HashingMetric else HashingMetric[metric]
0060 self.annoySeed = annoySeed
0061 self.zBins = zBins
0062 self.phiBins = phiBins
0063
0064 if seedingAlgorithm == SeedingAlgorithm.GridTriplet:
0065 self.bucketSize = 0
0066 self.metric = HashingMetric.dphi
0067 self.annoySeed = 123456789
0068 self.zBins = 0
0069 self.phiBins = 0
0070
0071 def __str__(self):
0072 return (
0073 f"mu={self.mu}, bucketSize={self.bucketSize}, maxSeedsPerSpM={self.maxSeedsPerSpM}, "
0074 f"detector={self.detector}, seedingAlgorithm={self.seedingAlgorithm}, metric={self.metric}, "
0075 f"annoySeed={self.annoySeed}, zBins={self.zBins}, phiBins={self.phiBins}"
0076 )
0077
0078 def __repr__(self):
0079 return self.__str__()
0080
0081 def save(self, path: pathlib.Path):
0082 with open(path, "w") as f:
0083 f.write(str(self))
0084
0085 @property
0086 def doHashing(self):
0087 return self.bucketSize > 0
0088
0089 @property
0090 def directory(self):
0091 outDir = f"detector_{extractEnumName(self.detector)}"
0092 outDir += "_output"
0093 doHashing = self.bucketSize > 0
0094 if doHashing:
0095 outDir += "_hashing"
0096
0097 outDir += f"_mu_{self.mu}"
0098 if doHashing:
0099 outDir += f"_bucket_{self.bucketSize}"
0100
0101 outDir += f"_maxSeedsPerSpM_{self.maxSeedsPerSpM}"
0102
0103 outDir += f"_seedFinderConfig_{'TrackML' if self.detector == DetectorName.generic else 'ODD'}"
0104
0105 outDir += f"_seedingAlgorithm_{extractEnumName(self.seedingAlgorithm)}Seeding"
0106 if doHashing:
0107 if self.metric != "angular":
0108 outDir += f"_metric_{extractEnumName(self.metric)}"
0109 outDir += f"_annoySeed_{self.annoySeed}"
0110 if self.zBins != 0:
0111 outDir += f"_zBins_{self.zBins}"
0112 if self.phiBins != 0:
0113 outDir += f"_phiBins_{self.phiBins}"
0114
0115 return outDir
0116
0117 def getDetectorInfo(self):
0118 actsExamplesDir = Path(__file__).parent.parent.parent
0119
0120 if self.detector == DetectorName.ODD:
0121 from acts.examples.odd import (
0122 getOpenDataDetector,
0123 getOpenDataDetectorDirectory,
0124 )
0125
0126 geoDir = getOpenDataDetectorDirectory()
0127
0128 oddMaterialMap = geoDir / "data/odd-material-maps.root"
0129 oddDigiConfig = actsExamplesDir / "Configs/odd-digi-smearing-config.json"
0130 oddSeedingSel = actsExamplesDir / "Configs/odd-seeding-config.json"
0131 oddMaterialDeco = acts.IMaterialDecorator.fromFile(oddMaterialMap)
0132
0133 detector = getOpenDataDetector(
0134 odd_dir=geoDir, materialDecorator=oddMaterialDeco
0135 )
0136 trackingGeometry = detector.trackingGeometry()
0137
0138 digiConfig = oddDigiConfig
0139
0140 geoSelectionConfigFile = oddSeedingSel
0141
0142 elif self.detector == DetectorName.generic:
0143 print("Create detector and tracking geometry")
0144
0145 detector = acts.examples.GenericDetector()
0146 trackingGeometry = detector.trackingGeometry()
0147 digiConfig = actsExamplesDir / "Configs/generic-digi-smearing-config.json"
0148 geoSelectionConfigFile = (
0149 actsExamplesDir / "Configs/generic-seeding-config.json"
0150 )
0151 else:
0152 exit("Detector not supported")
0153
0154 return detector, trackingGeometry, digiConfig, geoSelectionConfigFile
0155
0156
0157 def extractEnumName(enumvar):
0158 return str(enumvar).split(".")[-1]
0159
0160
0161 def runHashingSeeding(
0162 nevents: int,
0163 trackingGeometry: acts.TrackingGeometry,
0164 field: acts.MagneticFieldProvider,
0165 outputDir: pathlib.Path,
0166 saveFiles: bool,
0167 npileup: int,
0168 seedingAlgorithm: SeedingAlgorithm,
0169 maxSeedsPerSpM: int,
0170 digiConfig: pathlib.Path,
0171 geoSelectionConfigFile: pathlib.Path,
0172 eta=4,
0173 rnd: acts.examples.RandomNumbers = acts.examples.RandomNumbers(seed=42),
0174 config: Config = None,
0175 s=None,
0176 ):
0177 from acts.examples.reconstruction import (
0178 SeedFinderConfigArg,
0179 SeedFinderOptionsArg,
0180 HashingTrainingConfigArg,
0181 HashingAlgorithmConfigArg,
0182 )
0183
0184 u = acts.UnitConstants
0185
0186 outputDir = Path(outputDir)
0187
0188 s = s or acts.examples.Sequencer(
0189 events=nevents,
0190 numThreads=1,
0191 outputDir=str(outputDir),
0192 trackFpes=False,
0193 )
0194
0195 addPythia8(
0196 s,
0197 hardProcess=["Top:qqbar2ttbar=on"],
0198 npileup=npileup,
0199 vtxGen=acts.examples.GaussianVertexGenerator(
0200 stddev=acts.Vector4(0, 0, 50 * u.mm, 0),
0201 mean=acts.Vector4(0, 0, 0, 0),
0202 ),
0203 rnd=rnd,
0204 )
0205
0206 addGenParticleSelection(
0207 s,
0208 ParticleSelectorConfig(
0209 rho=(0.0, 24 * u.mm),
0210 absZ=(0.0, 1.0 * u.m),
0211 ),
0212 )
0213
0214 addFatras(
0215 s,
0216 trackingGeometry,
0217 field,
0218 enableInteractions=True,
0219
0220 outputDirCsv=outputDir if saveFiles else None,
0221 rnd=rnd,
0222 )
0223
0224 addDigitization(
0225 s,
0226 trackingGeometry,
0227 field,
0228 digiConfigFile=digiConfig,
0229 outputDirRoot=outputDir,
0230 outputDirCsv=outputDir if saveFiles else None,
0231 rnd=rnd,
0232 )
0233
0234 addDigiParticleSelection(
0235 s,
0236 ParticleSelectorConfig(
0237 pt=(1.0 * u.GeV, None),
0238 eta=(-eta, eta),
0239 measurements=(9, None),
0240 removeNeutral=True,
0241 ),
0242 )
0243
0244 import numpy as np
0245
0246 cotThetaMax = 1 / (np.tan(2 * np.arctan(np.exp(-eta))))
0247 seedFinderConfigArg = SeedFinderConfigArg(
0248 r=(None, 200 * u.mm),
0249 deltaR=(1 * u.mm, 300 * u.mm),
0250 collisionRegion=(-250 * u.mm, 250 * u.mm),
0251 z=(-2000 * u.mm, 2000 * u.mm),
0252 maxSeedsPerSpM=maxSeedsPerSpM,
0253 sigmaScattering=5,
0254 radLengthPerSeed=0.1,
0255 minPt=500 * u.MeV,
0256 impactMax=3 * u.mm,
0257 cotThetaMax=cotThetaMax,
0258 )
0259
0260 seedFinderOptionsArg: SeedFinderOptionsArg = SeedFinderOptionsArg(bFieldInZ=2 * u.T)
0261 hashingTrainingConfigArg: HashingTrainingConfigArg = HashingTrainingConfigArg(
0262 annoySeed=config.annoySeed, f=2 if config.metric == HashingMetric.dR else 1
0263 )
0264 hashingAlgorithmConfigArg: HashingAlgorithmConfigArg = HashingAlgorithmConfigArg(
0265 bucketSize=config.bucketSize,
0266 zBins=config.zBins,
0267 phiBins=config.phiBins,
0268 )
0269 initialSigmas: Optional[list] = [
0270 1 * u.mm,
0271 1 * u.mm,
0272 1 * u.degree,
0273 1 * u.degree,
0274 0.1 / u.GeV,
0275 1 * u.ns,
0276 ]
0277
0278 initialVarInflation: Optional[list] = [1.0] * 6
0279
0280 addSeeding(
0281 s,
0282 trackingGeometry,
0283 field,
0284 seedFinderConfigArg,
0285 seedFinderOptionsArg,
0286 hashingTrainingConfigArg,
0287 hashingAlgorithmConfigArg,
0288 seedingAlgorithm=seedingAlgorithm,
0289 geoSelectionConfigFile=geoSelectionConfigFile,
0290 initialSigmas=initialSigmas,
0291 initialVarInflation=initialVarInflation,
0292 outputDirRoot=outputDir,
0293 outputDirCsv=outputDir if saveFiles else None,
0294 )
0295
0296 return s
0297
0298
0299 if __name__ == "__main__":
0300 eta = 3
0301
0302 parser = argparse.ArgumentParser(
0303 description="Example script to run seed finding with hashing"
0304 )
0305 parser.add_argument("--mu", type=int, default=50)
0306 parser.add_argument("--bucketSize", type=int, default=100)
0307 parser.add_argument(
0308 "--nevents",
0309 type=int,
0310 default=20,
0311 )
0312 parser.add_argument("--maxSeedsPerSpM", type=int, default=10)
0313 parser.add_argument("--seedingAlgorithm", type=str, default="HashingPrototype")
0314 parser.add_argument("--saveFiles", type=bool, default=True)
0315 parser.add_argument("--annoySeed", type=int, default=123456789)
0316 parser.add_argument("--zBins", type=int, default=0)
0317 parser.add_argument("--phiBins", type=int, default=100)
0318 parser.add_argument("--metric", type=str, default="dphi")
0319 args = parser.parse_args()
0320
0321 print(args)
0322
0323 mu = args.mu
0324 bucketSize = args.bucketSize
0325 nevents = args.nevents
0326 saveFiles = args.saveFiles
0327 annoySeed = args.annoySeed
0328 zBins = args.zBins
0329 phiBins = args.phiBins
0330 maxSeedsPerSpM = args.maxSeedsPerSpM
0331
0332 metric = HashingMetric[args.metric]
0333
0334 seedingAlgorithm = SeedingAlgorithm[args.seedingAlgorithm]
0335
0336 detector = DetectorName.generic
0337
0338 u = acts.UnitConstants
0339
0340 config = Config(
0341 mu=mu,
0342 bucketSize=bucketSize,
0343 maxSeedsPerSpM=maxSeedsPerSpM,
0344 detector=detector,
0345 seedingAlgorithm=seedingAlgorithm,
0346 metric=metric,
0347 annoySeed=annoySeed,
0348 zBins=zBins,
0349 phiBins=phiBins,
0350 )
0351
0352 doHashing = config.doHashing
0353 bucketSize = config.bucketSize
0354 npileup = config.mu
0355 maxSeedsPerSpM = config.maxSeedsPerSpM
0356
0357 outputDir = pathlib.Path.cwd() / config.directory
0358
0359 if not outputDir.exists():
0360 outputDir.mkdir(parents=True)
0361
0362 config.save(outputDir / "config_file.txt")
0363
0364 field = acts.ConstantBField(acts.Vector3(0.0, 0.0, 2.0 * u.T))
0365 rnd = acts.examples.RandomNumbers(seed=42)
0366
0367 _, trackingGeometry, digiConfig, geoSelectionConfigFile = config.getDetectorInfo()
0368
0369 s = runHashingSeeding(
0370 nevents,
0371 trackingGeometry,
0372 field,
0373 outputDir,
0374 saveFiles,
0375 npileup,
0376 seedingAlgorithm,
0377 maxSeedsPerSpM,
0378 digiConfig,
0379 geoSelectionConfigFile,
0380 rnd=rnd,
0381 eta=eta,
0382 config=config,
0383 )
0384
0385 logLevel = acts.logging.VERBOSE
0386 rootSpacePointsWriter = RootSpacePointWriter(
0387 level=logLevel,
0388 inputSpacePoints="spacepoints",
0389 filePath=str(outputDir / "spacepoints.root"),
0390 )
0391 s.addWriter(rootSpacePointsWriter)
0392
0393 rootSeedsWriter = acts.examples.root.RootSeedWriter(
0394 level=logLevel,
0395 inputSeeds="seeds",
0396 filePath=str(outputDir / "seeds.root"),
0397 )
0398 s.addWriter(rootSeedsWriter)
0399
0400 addCKFTracks(
0401 s,
0402 trackingGeometry,
0403 field,
0404 TrackSelectorConfig(
0405 pt=(1.0 * u.GeV, None),
0406 absEta=(None, eta),
0407 nMeasurementsMin=6,
0408 ),
0409 twoWay=False,
0410 outputDirRoot=outputDir,
0411 )
0412
0413 s.run()
