#!/usr/bin/env python3

import argparse

import pathlib

import acts

import acts.examples

from acts.examples.simulation import (

addPythia8,

ParticleSelectorConfig,

addGenParticleSelection,

addFatras,

addDigitization,

addDigiParticleSelection,

)

from acts.examples.reconstruction import (

addSeeding,

addCKFTracks,

TrackSelectorConfig,

SeedingAlgorithm,

)

from acts.examples.root import (

RootSpacePointWriter,

)

from pathlib import Path

from typing import Optional

from enum import Enum

DetectorName = Enum("DetectorName", "ODD generic")

HashingMetric = Enum("HashingMetric", "dphi dR")

class Config:

def __init__(

self,

mu: int = None,

bucketSize: int = 100,

maxSeedsPerSpM: int = 10,

detector: DetectorName = DetectorName.generic,

seedingAlgorithm: SeedingAlgorithm = SeedingAlgorithm.HashingPrototype,

metric: str = HashingMetric.dphi,

annoySeed: int = 123456789,

zBins: int = 0,

phiBins: int = 100,

):

self.mu = mu

self.bucketSize = bucketSize

self.maxSeedsPerSpM = maxSeedsPerSpM

self.detector = detector

self.seedingAlgorithm = (

seedingAlgorithm

if type(seedingAlgorithm) == SeedingAlgorithm

else SeedingAlgorithm[seedingAlgorithm]

)

self.metric = metric if type(metric) == HashingMetric else HashingMetric[metric]

self.annoySeed = annoySeed

self.zBins = zBins

self.phiBins = phiBins

if seedingAlgorithm == SeedingAlgorithm.GridTriplet:

self.bucketSize = 0

self.metric = HashingMetric.dphi

self.annoySeed = 123456789

self.zBins = 0

self.phiBins = 0

def __str__(self):

return (

f"mu={self.mu}, bucketSize={self.bucketSize}, maxSeedsPerSpM={self.maxSeedsPerSpM}, "

f"detector={self.detector}, seedingAlgorithm={self.seedingAlgorithm}, metric={self.metric}, "

f"annoySeed={self.annoySeed}, zBins={self.zBins}, phiBins={self.phiBins}"

)

def __repr__(self):

return self.__str__()

def save(self, path: pathlib.Path):

with open(path, "w") as f:

f.write(str(self))

@property

def doHashing(self):

return self.bucketSize > 0

@property

def directory(self):

outDir = f"detector_{extractEnumName(self.detector)}"

outDir += "_output"

doHashing = self.bucketSize > 0

if doHashing:

outDir += "_hashing"

outDir += f"_mu_{self.mu}"

if doHashing:

outDir += f"_bucket_{self.bucketSize}"

outDir += f"_maxSeedsPerSpM_{self.maxSeedsPerSpM}"

outDir += f"_seedFinderConfig_{'TrackML' if self.detector == DetectorName.generic else 'ODD'}"

outDir += f"_seedingAlgorithm_{extractEnumName(self.seedingAlgorithm)}Seeding"

if doHashing:

if self.metric != "angular":

outDir += f"_metric_{extractEnumName(self.metric)}"

outDir += f"_annoySeed_{self.annoySeed}"

if self.zBins != 0:

outDir += f"_zBins_{self.zBins}"

if self.phiBins != 0:

outDir += f"_phiBins_{self.phiBins}"

return outDir

def getDetectorInfo(self):

actsExamplesDir = Path(__file__).parent.parent.parent

if self.detector == DetectorName.ODD:

from acts.examples.odd import (

getOpenDataDetector,

getOpenDataDetectorDirectory,

)

geoDir = getOpenDataDetectorDirectory()

oddMaterialMap = geoDir / "data/odd-material-maps.root"

oddDigiConfig = actsExamplesDir / "Configs/odd-digi-smearing-config.json"

oddSeedingSel = actsExamplesDir / "Configs/odd-seeding-config.json"

oddMaterialDeco = acts.IMaterialDecorator.fromFile(oddMaterialMap)

detector = getOpenDataDetector(

odd_dir=geoDir, materialDecorator=oddMaterialDeco

)

trackingGeometry = detector.trackingGeometry()

digiConfig = oddDigiConfig

geoSelectionConfigFile = oddSeedingSel

elif self.detector == DetectorName.generic:

print("Create detector and tracking geometry")

detector = acts.examples.GenericDetector()

trackingGeometry = detector.trackingGeometry()

digiConfig = actsExamplesDir / "Configs/generic-digi-smearing-config.json"

geoSelectionConfigFile = (

actsExamplesDir / "Configs/generic-seeding-config.json"

)

else:

exit("Detector not supported")

return detector, trackingGeometry, digiConfig, geoSelectionConfigFile

def extractEnumName(enumvar):

return str(enumvar).split(".")[-1]

def runHashingSeeding(

nevents: int,

trackingGeometry: acts.TrackingGeometry,

field: acts.MagneticFieldProvider,

outputDir: pathlib.Path,

saveFiles: bool,

npileup: int,

seedingAlgorithm: SeedingAlgorithm,

maxSeedsPerSpM: int,

digiConfig: pathlib.Path,

geoSelectionConfigFile: pathlib.Path,

eta=4,

rnd: acts.examples.RandomNumbers = acts.examples.RandomNumbers(seed=42),

config: Config = None,

s=None,

):

from acts.examples.reconstruction import (

SeedFinderConfigArg,

SeedFinderOptionsArg,

HashingTrainingConfigArg,

HashingAlgorithmConfigArg,

)

u = acts.UnitConstants

outputDir = Path(outputDir)

s = s or acts.examples.Sequencer(

events=nevents,

numThreads=1,

outputDir=str(outputDir),

trackFpes=False,

)

addPythia8(

s,

hardProcess=["Top:qqbar2ttbar=on"],

npileup=npileup,

vtxGen=acts.examples.GaussianVertexGenerator(

stddev=acts.Vector4(0, 0, 50 * u.mm, 0),

mean=acts.Vector4(0, 0, 0, 0),

),

rnd=rnd,

)

addGenParticleSelection(

s,

ParticleSelectorConfig(

rho=(0.0, 24 * u.mm),

absZ=(0.0, 1.0 * u.m),

),

)

addFatras(

s,

trackingGeometry,

field,

enableInteractions=True,

# outputDirRoot=outputDir, # RootParticle ERROR when setting the outputDirRoot

outputDirCsv=outputDir if saveFiles else None,

rnd=rnd,

)

addDigitization(

s,

trackingGeometry,

field,

digiConfigFile=digiConfig,

outputDirRoot=outputDir,

outputDirCsv=outputDir if saveFiles else None,

rnd=rnd,

)

addDigiParticleSelection(

s,

ParticleSelectorConfig(

pt=(1.0 * u.GeV, None),

eta=(-eta, eta),

measurements=(9, None),

removeNeutral=True,

),

)

import numpy as np

cotThetaMax = 1 / (np.tan(2 * np.arctan(np.exp(-eta)))) # =1/tan(2×atan(e^(-eta)))

seedFinderConfigArg = SeedFinderConfigArg(

r=(None, 200 * u.mm), # rMin=default, 33mm

deltaR=(1 * u.mm, 300 * u.mm),

collisionRegion=(-250 * u.mm, 250 * u.mm),

z=(-2000 * u.mm, 2000 * u.mm),

maxSeedsPerSpM=maxSeedsPerSpM,

sigmaScattering=5,

radLengthPerSeed=0.1,

minPt=500 * u.MeV,

impactMax=3 * u.mm,

cotThetaMax=cotThetaMax,

)

seedFinderOptionsArg: SeedFinderOptionsArg = SeedFinderOptionsArg(bFieldInZ=2 * u.T)

hashingTrainingConfigArg: HashingTrainingConfigArg = HashingTrainingConfigArg(

annoySeed=config.annoySeed, f=2 if config.metric == HashingMetric.dR else 1

)

hashingAlgorithmConfigArg: HashingAlgorithmConfigArg = HashingAlgorithmConfigArg(

bucketSize=config.bucketSize,

zBins=config.zBins,

phiBins=config.phiBins,

)

initialSigmas: Optional[list] = [

1 * u.mm,

1 * u.mm,

1 * u.degree,

1 * u.degree,

0.1 / u.GeV,

1 * u.ns,

]

initialVarInflation: Optional[list] = [1.0] * 6

addSeeding(

s,

trackingGeometry,

field,

seedFinderConfigArg,

seedFinderOptionsArg,

hashingTrainingConfigArg,

hashingAlgorithmConfigArg,

seedingAlgorithm=seedingAlgorithm,

geoSelectionConfigFile=geoSelectionConfigFile,

initialSigmas=initialSigmas,

initialVarInflation=initialVarInflation,

outputDirRoot=outputDir,

outputDirCsv=outputDir if saveFiles else None,

)

return s

if __name__ == "__main__":

eta = 3

parser = argparse.ArgumentParser(

description="Example script to run seed finding with hashing"

)

parser.add_argument("--mu", type=int, default=50)

parser.add_argument("--bucketSize", type=int, default=100)

parser.add_argument(

"--nevents",

type=int,

default=20,

)

parser.add_argument("--maxSeedsPerSpM", type=int, default=10)

parser.add_argument("--seedingAlgorithm", type=str, default="HashingPrototype")

parser.add_argument("--saveFiles", type=bool, default=True)

parser.add_argument("--annoySeed", type=int, default=123456789)

parser.add_argument("--zBins", type=int, default=0)

parser.add_argument("--phiBins", type=int, default=100)

parser.add_argument("--metric", type=str, default="dphi")

args = parser.parse_args()

print(args)

mu = args.mu

bucketSize = args.bucketSize

nevents = args.nevents

saveFiles = args.saveFiles

annoySeed = args.annoySeed

zBins = args.zBins

phiBins = args.phiBins

maxSeedsPerSpM = args.maxSeedsPerSpM

metric = HashingMetric[args.metric]

seedingAlgorithm = SeedingAlgorithm[args.seedingAlgorithm]

detector = DetectorName.generic

u = acts.UnitConstants

config = Config(

mu=mu,

bucketSize=bucketSize,

maxSeedsPerSpM=maxSeedsPerSpM,

detector=detector,

seedingAlgorithm=seedingAlgorithm,

metric=metric,

annoySeed=annoySeed,

zBins=zBins,

phiBins=phiBins,

)

doHashing = config.doHashing

bucketSize = config.bucketSize

npileup = config.mu

maxSeedsPerSpM = config.maxSeedsPerSpM

outputDir = pathlib.Path.cwd() / config.directory

if not outputDir.exists():

outputDir.mkdir(parents=True)

config.save(outputDir / "config_file.txt")

field = acts.ConstantBField(acts.Vector3(0.0, 0.0, 2.0 * u.T))

rnd = acts.examples.RandomNumbers(seed=42)

_, trackingGeometry, digiConfig, geoSelectionConfigFile = config.getDetectorInfo()

s = runHashingSeeding(

nevents,

trackingGeometry,

field,

outputDir,

saveFiles,

npileup,

seedingAlgorithm,

maxSeedsPerSpM,

digiConfig,

geoSelectionConfigFile,

rnd=rnd,

eta=eta,

config=config,

)

logLevel = acts.logging.VERBOSE

rootSpacePointsWriter = RootSpacePointWriter(

level=logLevel,

inputSpacePoints="spacepoints",

filePath=str(outputDir / "spacepoints.root"),

)

s.addWriter(rootSpacePointsWriter)

rootSeedsWriter = acts.examples.root.RootSeedWriter(

level=logLevel,

inputSeeds="seeds",

filePath=str(outputDir / "seeds.root"),

)

s.addWriter(rootSeedsWriter)

addCKFTracks(

s,

trackingGeometry,

field,

TrackSelectorConfig(

pt=(1.0 * u.GeV, None),

absEta=(None, eta),

nMeasurementsMin=6,

),

twoWay=False,

outputDirRoot=outputDir,

)

s.run()