Scripts/Python/hashing_seeding.py

0001 #!/usr/bin/env python3
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 = 1000,
0043         detector: DetectorName = DetectorName.generic,
0044         seedingAlgorithm: SeedingAlgorithm = SeedingAlgorithm.Hashing,
0045         metric: str = HashingMetric.dphi,
0046         annoySeed: int = 123456789,
0047         zBins: int = 100_000,
0048         phiBins: int = 0,
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         # outputDirRoot=outputDir,  # RootParticle ERROR when setting the outputDirRoot
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))))  # =1/tan(2×atan(e^(-eta)))
0247     seedFinderConfigArg = SeedFinderConfigArg(
0248         r=(None, 200 * u.mm),  # rMin=default, 33mm
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 = 4
0301     # eta = 2.5
0302
0303     parser = argparse.ArgumentParser(
0304         description="Example script to run seed finding with hashing"
0305     )
0306     parser.add_argument("--mu", type=int, default=50)
0307     parser.add_argument("--bucketSize", type=int, default=100)
0308     parser.add_argument(
0309         "--nevents",
0310         type=int,
0311         default=20,
0312     )
0313     parser.add_argument("--maxSeedsPerSpM", type=int, default=1000)
0314     parser.add_argument("--seedingAlgorithm", type=str, default="Hashing")
0315     parser.add_argument("--saveFiles", type=bool, default=True)
0316     parser.add_argument("--annoySeed", type=int, default=123456789)
0317     parser.add_argument("--zBins", type=int, default=100000)
0318     parser.add_argument("--phiBins", type=int, default=0)
0319     parser.add_argument("--metric", type=str, default="dphi")
0320     args = parser.parse_args()
0321
0322     print(args)
0323
0324     mu = args.mu
0325     bucketSize = args.bucketSize
0326     nevents = args.nevents
0327     saveFiles = args.saveFiles
0328     annoySeed = args.annoySeed
0329     zBins = args.zBins
0330     phiBins = args.phiBins
0331     maxSeedsPerSpM = args.maxSeedsPerSpM
0332
0333     metric = HashingMetric[args.metric]
0334
0335     seedingAlgorithm = SeedingAlgorithm[args.seedingAlgorithm]
0336
0337     detector = DetectorName.generic
0338
0339     u = acts.UnitConstants
0340
0341     config = Config(
0342         mu=mu,
0343         bucketSize=bucketSize,
0344         maxSeedsPerSpM=maxSeedsPerSpM,
0345         detector=detector,
0346         seedingAlgorithm=seedingAlgorithm,
0347         metric=metric,
0348         annoySeed=annoySeed,
0349         zBins=zBins,
0350         phiBins=phiBins,
0351     )
0352
0353     doHashing = config.doHashing
0354     bucketSize = config.bucketSize
0355     npileup = config.mu
0356     maxSeedsPerSpM = config.maxSeedsPerSpM
0357
0358     outputDir = pathlib.Path.cwd() / config.directory
0359
0360     if not outputDir.exists():
0361         outputDir.mkdir(parents=True)
0362
0363     config.save(outputDir / "config_file.txt")
0364
0365     field = acts.ConstantBField(acts.Vector3(0.0, 0.0, 2.0 * u.T))
0366     rnd = acts.examples.RandomNumbers(seed=42)
0367
0368     _, trackingGeometry, digiConfig, geoSelectionConfigFile = config.getDetectorInfo()
0369
0370     s = runHashingSeeding(
0371         nevents,
0372         trackingGeometry,
0373         field,
0374         outputDir,
0375         saveFiles,
0376         npileup,
0377         seedingAlgorithm,
0378         maxSeedsPerSpM,
0379         digiConfig,
0380         geoSelectionConfigFile,
0381         rnd=rnd,
0382         eta=eta,
0383         config=config,
0384     )
0385
0386     logLevel = acts.logging.VERBOSE
0387     rootSpacepointsWriter = RootSpacepointWriter(
0388         level=logLevel,
0389         inputSpacepoints="spacepoints",
0390         filePath=str(outputDir / "spacepoints.root"),
0391     )
0392     s.addWriter(rootSpacepointsWriter)
0393
0394     rootSeedsWriter = acts.examples.root.RootSeedWriter(
0395         level=logLevel,
0396         inputSeeds="seeds",
0397         filePath=str(outputDir / "seeds.root"),
0398     )
0399     s.addWriter(rootSeedsWriter)
0400
0401     # TrackFinding ERROR no intersection found; TrackExtrapolationError:2
0402     addCKFTracks(
0403         s,
0404         trackingGeometry,
0405         field,
0406         TrackSelectorConfig(
0407             pt=(1.0 * u.GeV, None),
0408             absEta=(None, eta),
0409             nMeasurementsMin=6,
0410         ),
0411         outputDirRoot=outputDir,
0412         writeTrajectories=False,
0413         twoWay=False,
0414     )
0415
0416     s.run()