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