Scripts/Python/onnx_seeding.py

0001 #!/usr/bin/env python3
0002 """
0003 Full reconstruction chain on the OpenDataDetector using an ONNX-based
0004 PythonCallable seeder. Demonstrates how to plug an ONNX seeding model
0005 (here: GUNTAM) into ACTS via SeedingAlgorithm.PythonCallable.
0006
0007 Usage:
0008     python onnx_seeding.py --model-path /path/to/model.onnx
0009     python onnx_seeding.py --model-path /path/to/model.onnx --ttbar --ttbar-pu 200
0010 """
0011
0012 import argparse
0013 import pathlib
0014
0015 import acts
0016 import acts.examples
0017 from acts.examples.simulation import (
0018     MomentumConfig,
0019     EtaConfig,
0020     PhiConfig,
0021     ParticleConfig,
0022     ParticleSelectorConfig,
0023     addParticleGun,
0024     addPythia8,
0025     addGenParticleSelection,
0026     addFatras,
0027     addDigitization,
0028     addDigiParticleSelection,
0029 )
0030 from acts.examples.reconstruction import (
0031     SeedingAlgorithm,
0032     addSeeding,
0033     CkfConfig,
0034     addCKFTracks,
0035     TrackSelectorConfig,
0036     addAmbiguityResolution,
0037     AmbiguityResolutionConfig,
0038     addVertexFitting,
0039     VertexFinder,
0040 )
0041 from acts.examples.odd import getOpenDataDetector, getOpenDataDetectorDirectory
0042 from guntam_transformer_seeder import guntam_transformer_seeder
0043
0044 u = acts.UnitConstants
0045
0046
0047 parser = argparse.ArgumentParser(
0048     description="ODD full chain with an ONNX PythonCallable seeder (GUNTAM)"
0049 )
0050 parser.add_argument(
0051     "--output",
0052     "-o",
0053     help="Output directory",
0054     type=pathlib.Path,
0055     default=pathlib.Path.cwd() / "onnx_seeding_output",
0056 )
0057 parser.add_argument("--events", "-n", help="Number of events", type=int, default=100)
0058 parser.add_argument(
0059     "--threads",
0060     help="Number of sequencer threads, i.e. events processed in parallel (-1 = all cores). "
0061     "For CPU inference ensure --threads * --onnx-threads does not exceed available cores.",
0062     type=int,
0063     default=-1,
0064 )
0065 parser.add_argument(
0066     "--onnx-threads",
0067     help="ONNX intra- and inter-op thread count for CPU inference (default 1).",
0068     type=int,
0069     default=1,
0070 )
0071 parser.add_argument(
0072     "--gpu",
0073     help="Use CUDA GPU for ONNX inference (CUDAExecutionProvider with CPU fallback).",
0074     action="store_true",
0075 )
0076 parser.add_argument(
0077     "--model-path",
0078     help="Path to the ONNX seeding model",
0079     type=pathlib.Path,
0080     required=True,
0081 )
0082 parser.add_argument(
0083     "--ttbar",
0084     help="Use Pythia8 ttbar instead of particle gun",
0085     action="store_true",
0086 )
0087 parser.add_argument(
0088     "--ttbar-pu",
0089     help="Number of pile-up events (only used with --ttbar)",
0090     type=int,
0091     default=200,
0092 )
0093 parser.add_argument(
0094     "--output-root",
0095     help="Write ROOT output files",
0096     default=True,
0097     action=argparse.BooleanOptionalAction,
0098 )
0099 parser.add_argument(
0100     "--output-csv",
0101     help="Write CSV output files",
0102     default=True,
0103     action=argparse.BooleanOptionalAction,
0104 )
0105
0106 args = parser.parse_args()
0107
0108 outputDir = args.output
0109 outputDir.mkdir(parents=True, exist_ok=True)
0110
0111 geoDir = getOpenDataDetectorDirectory()
0112 actsDir = pathlib.Path(__file__).resolve().parents[3]
0113
0114 oddMaterialMap = geoDir / "data/odd-material-maps.root"
0115 oddDigiConfig = actsDir / "Examples/Configs/odd-digi-smearing-config.json"
0116 oddSeedingSel = actsDir / "Examples/Configs/odd-seeding-config.json"
0117
0118 oddMaterialDeco = acts.IMaterialDecorator.fromFile(oddMaterialMap)
0119 detector = getOpenDataDetector(odd_dir=geoDir, materialDecorator=oddMaterialDeco)
0120 trackingGeometry = detector.trackingGeometry()
0121
0122 field = acts.ConstantBField(acts.Vector3(0.0, 0.0, 2.0 * u.T))
0123 rnd = acts.examples.RandomNumbers(seed=42)
0124
0125 s = acts.examples.Sequencer(
0126     events=args.events,
0127     numThreads=args.threads,
0128     outputDir=str(outputDir),
0129     # ONNX SIMD kernels raise harmless FP underflow signals; disabling avoids false-positive failures unrelated to this script.
0130     failOnUnmaskedFpe=False,
0131 )
0132
0133 if not args.ttbar:
0134     addParticleGun(
0135         s,
0136         MomentumConfig(1.0 * u.GeV, 10.0 * u.GeV, transverse=True),
0137         EtaConfig(-3.0, 3.0),
0138         PhiConfig(0.0, 360.0 * u.degree),
0139         ParticleConfig(4, acts.PdgParticle.eMuon, randomizeCharge=True),
0140         vtxGen=acts.examples.GaussianVertexGenerator(
0141             mean=acts.Vector4(0, 0, 0, 0),
0142             stddev=acts.Vector4(0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 1.0 * u.ns),
0143         ),
0144         multiplicity=200,
0145         rnd=rnd,
0146     )
0147 else:
0148     addPythia8(
0149         s,
0150         hardProcess=["Top:qqbar2ttbar=on"],
0151         npileup=args.ttbar_pu,
0152         vtxGen=acts.examples.GaussianVertexGenerator(
0153             mean=acts.Vector4(0, 0, 0, 0),
0154             stddev=acts.Vector4(0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 5.0 * u.ns),
0155         ),
0156         rnd=rnd,
0157         outputDirRoot=outputDir if args.output_root else None,
0158         outputDirCsv=outputDir if args.output_csv else None,
0159     )
0160
0161     addGenParticleSelection(
0162         s,
0163         ParticleSelectorConfig(
0164             rho=(0.0, 24 * u.mm),
0165             absZ=(0.0, 1.0 * u.m),
0166             eta=(-3.0, 3.0),
0167             pt=(150 * u.MeV, None),
0168         ),
0169     )
0170
0171 addFatras(
0172     s,
0173     trackingGeometry,
0174     field,
0175     enableInteractions=True,
0176     outputDirRoot=outputDir if args.output_root else None,
0177     outputDirCsv=outputDir if args.output_csv else None,
0178     rnd=rnd,
0179 )
0180
0181 addDigitization(
0182     s,
0183     trackingGeometry,
0184     field,
0185     digiConfigFile=oddDigiConfig,
0186     outputDirRoot=outputDir if args.output_root else None,
0187     outputDirCsv=outputDir if args.output_csv else None,
0188     rnd=rnd,
0189 )
0190
0191 addDigiParticleSelection(
0192     s,
0193     ParticleSelectorConfig(
0194         pt=(1.0 * u.GeV, None),
0195         eta=(-3.0, 3.0),
0196         measurements=(9, None),
0197         removeNeutral=True,
0198     ),
0199 )
0200
0201 addSeeding(
0202     s,
0203     trackingGeometry,
0204     field,
0205     seedingAlgorithm=SeedingAlgorithm.PythonCallable,
0206     customSeeder=guntam_transformer_seeder,
0207     customSeederConfig={
0208         "model_path": str(args.model_path),
0209         "num_threads": args.onnx_threads,
0210         "providers": (
0211             ["CUDAExecutionProvider", "CPUExecutionProvider"] if args.gpu else None
0212         ),
0213     },
0214     geoSelectionConfigFile=oddSeedingSel,
0215     initialSigmas=[
0216         1 * u.mm,
0217         1 * u.mm,
0218         1 * u.degree,
0219         1 * u.degree,
0220         0 * u.e / u.GeV,
0221         1 * u.ns,
0222     ],
0223     initialSigmaQoverPt=0.1 * u.e / u.GeV,
0224     initialSigmaPtRel=0.1,
0225     initialVarInflation=[1.0] * 6,
0226     particleHypothesis=acts.ParticleHypothesis.muon,
0227     outputDirRoot=outputDir if args.output_root else None,
0228     outputDirCsv=outputDir if args.output_csv else None,
0229 )
0230
0231 addCKFTracks(
0232     s,
0233     trackingGeometry,
0234     field,
0235     TrackSelectorConfig(
0236         pt=(1.0 * u.GeV if args.ttbar else 0.0, None),
0237         absEta=(None, 3.0),
0238         loc0=(-4.0 * u.mm, 4.0 * u.mm),
0239         nMeasurementsMin=7,
0240         maxHoles=2,
0241         maxOutliers=2,
0242     ),
0243     CkfConfig(
0244         chi2CutOffMeasurement=15.0,
0245         chi2CutOffOutlier=25.0,
0246         numMeasurementsCutOff=2,
0247         seedDeduplication=True,
0248         stayOnSeed=False,
0249         pixelVolumes=[16, 17, 18],
0250         stripVolumes=[23, 24, 25],
0251         maxPixelHoles=1,
0252         maxStripHoles=2,
0253         constrainToVolumes=[
0254             2,  # beam pipe
0255             32,
0256             4,  # beam pipe gap
0257             16,
0258             17,
0259             18,  # pixel
0260             20,  # PST
0261             23,
0262             24,
0263             25,  # short strip
0264             26,
0265             8,  # long strip gap
0266             28,
0267             29,
0268             30,  # long strip
0269         ],
0270     ),
0271     outputDirRoot=outputDir if args.output_root else None,
0272     outputDirCsv=outputDir if args.output_csv else None,
0273     writeCovMat=True,
0274 )
0275
0276 addAmbiguityResolution(
0277     s,
0278     AmbiguityResolutionConfig(
0279         maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0280     ),
0281     outputDirRoot=outputDir if args.output_root else None,
0282     outputDirCsv=outputDir if args.output_csv else None,
0283     writeCovMat=True,
0284 )
0285
0286 addVertexFitting(
0287     s,
0288     field,
0289     vertexFinder=VertexFinder.AMVF,
0290     outputDirRoot=outputDir if args.output_root else None,
0291     outputDirCsv=outputDir if args.output_csv else None,
0292 )
0293
0294 s.run()