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