File indexing completed on 2026-08-06 08:28:58
0001
0002
0003 import os
0004 import argparse
0005 import pathlib
0006
0007 import acts
0008 import acts.examples
0009 from acts.examples.simulation import (
0010 MomentumConfig,
0011 EtaConfig,
0012 PhiConfig,
0013 ParticleConfig,
0014 ParticleSelectorConfig,
0015 addParticleGun,
0016 addPythia8,
0017 addGenParticleSelection,
0018 addFatras,
0019 addGeant4,
0020 addSimParticleSelection,
0021 addDigitization,
0022 addDigiParticleSelection,
0023 )
0024 from acts.examples.reconstruction import (
0025 addSeeding,
0026 CkfConfig,
0027 addCKFTracks,
0028 TrackSelectorConfig,
0029 addAmbiguityResolution,
0030 AmbiguityResolutionConfig,
0031 addAmbiguityResolutionML,
0032 AmbiguityResolutionMLConfig,
0033 addScoreBasedAmbiguityResolution,
0034 ScoreBasedAmbiguityResolutionConfig,
0035 addVertexFitting,
0036 VertexFinder,
0037 addSeedFilterML,
0038 SeedFilterMLDBScanConfig,
0039 )
0040 from acts.examples.odd import getOpenDataDetector, getOpenDataDetectorDirectory
0041
0042 u = acts.UnitConstants
0043
0044
0045 parser = argparse.ArgumentParser(description="Full chain with the OpenDataDetector")
0046 parser.add_argument(
0047 "--output",
0048 "-o",
0049 help="Output directory",
0050 type=pathlib.Path,
0051 default=pathlib.Path.cwd() / "odd_output",
0052 )
0053 parser.add_argument("--events", "-n", help="Number of events", type=int, default=100)
0054 parser.add_argument("--skip", "-s", help="Number of events", type=int, default=0)
0055 parser.add_argument(
0056 "--jobs",
0057 "-j",
0058 help="Number of worker threads (-1 uses all cores)",
0059 type=int,
0060 default=None,
0061 )
0062 parser.add_argument("--edm4hep", help="Use edm4hep inputs", type=pathlib.Path)
0063 parser.add_argument(
0064 "--geant4", help="Use Geant4 instead of fatras", action="store_true"
0065 )
0066 parser.add_argument(
0067 "--ttbar",
0068 help="Use Pythia8 (ttbar, pile-up 200) instead of particle gun",
0069 action="store_true",
0070 )
0071 parser.add_argument(
0072 "--ttbar-pu",
0073 help="Number of pile-up events for ttbar",
0074 type=int,
0075 default=200,
0076 )
0077 parser.add_argument(
0078 "--gun-particles",
0079 help="Multiplicity (no. of particles) of the particle gun",
0080 type=int,
0081 default=4,
0082 )
0083 parser.add_argument(
0084 "--gun-multiplicity",
0085 help="Multiplicity (no. of vertices) of the particle gun",
0086 type=int,
0087 default=200,
0088 )
0089 parser.add_argument(
0090 "--gun-eta-range",
0091 nargs=2,
0092 help="Eta range of the particle gun",
0093 type=float,
0094 default=[-3.0, 3.0],
0095 )
0096 parser.add_argument(
0097 "--gun-pt-range",
0098 nargs=2,
0099 help="Pt range of the particle gun (GeV)",
0100 type=float,
0101 default=[1.0 * u.GeV, 10.0 * u.GeV],
0102 )
0103 parser.add_argument(
0104 "--digi-config", help="Digitization configuration file", type=pathlib.Path
0105 )
0106 parser.add_argument(
0107 "--material-config", help="Material map configuration file", type=pathlib.Path
0108 )
0109 parser.add_argument(
0110 "--ambi-solver",
0111 help="Set which ambiguity solver to use, default is the classical one",
0112 type=str,
0113 choices=["greedy", "scoring", "ML"],
0114 default="greedy",
0115 )
0116 parser.add_argument(
0117 "--ambi-config",
0118 help="Set the configuration file for the Score Based ambiguity resolution",
0119 type=pathlib.Path,
0120 default=pathlib.Path.cwd() / "ambi_config.json",
0121 )
0122
0123 parser.add_argument(
0124 "--MLSeedFilter",
0125 help="Use the Ml seed filter to select seed after the seeding step",
0126 action="store_true",
0127 )
0128 parser.add_argument(
0129 "--reco",
0130 help="Switch reco on/off",
0131 default=True,
0132 action=argparse.BooleanOptionalAction,
0133 )
0134 parser.add_argument(
0135 "--output-root",
0136 help="Switch root output on/off",
0137 default=False,
0138 action=argparse.BooleanOptionalAction,
0139 )
0140 parser.add_argument(
0141 "--output-csv",
0142 help="Switch csv output on/off",
0143 default=False,
0144 action=argparse.BooleanOptionalAction,
0145 )
0146 parser.add_argument(
0147 "--output-obj",
0148 help="Switch obj output on/off",
0149 default=False,
0150 action=argparse.BooleanOptionalAction,
0151 )
0152 parser.add_argument(
0153 "--output-parquet",
0154 help="Switch parquet output on/off (requires ACTS_BUILD_EXAMPLES_PARQUET=ON)",
0155 default=False,
0156 action=argparse.BooleanOptionalAction,
0157 )
0158
0159 args = parser.parse_args()
0160
0161 outputDir = args.output
0162 ambi_ML = args.ambi_solver == "ML"
0163 ambi_scoring = args.ambi_solver == "scoring"
0164 ambi_config = args.ambi_config
0165 seedFilter_ML = args.MLSeedFilter
0166 geoDir = getOpenDataDetectorDirectory()
0167 actsDir = pathlib.Path(__file__).parent.parent.parent.parent
0168
0169
0170 oddMaterialMap = (
0171 args.material_config
0172 if args.material_config
0173 else geoDir / "data/odd-material-maps.root"
0174 )
0175
0176 oddDigiConfig = (
0177 args.digi_config
0178 if args.digi_config
0179 else actsDir / "Examples/Configs/odd-digi-smearing-config.json"
0180 )
0181
0182 oddSeedingSel = actsDir / "Examples/Configs/odd-seeding-config.json"
0183 oddMaterialDeco = acts.IMaterialDecorator.fromFile(oddMaterialMap)
0184
0185 detector = getOpenDataDetector(odd_dir=geoDir, materialDecorator=oddMaterialDeco)
0186 trackingGeometry = detector.trackingGeometry()
0187 decorators = detector.contextDecorators()
0188 field = acts.ConstantBField(acts.Vector3(0.0, 0.0, 2.0 * u.T))
0189 rnd = acts.examples.RandomNumbers(seed=42)
0190
0191 s = acts.examples.Sequencer(
0192 events=args.events,
0193 skip=args.skip,
0194 numThreads=args.jobs if args.jobs is not None else (1 if args.geant4 else -1),
0195 outputDir=str(outputDir),
0196 )
0197
0198 if args.edm4hep:
0199 import acts.examples.edm4hep
0200 from acts.examples.edm4hep import PodioReader
0201
0202 s.addReader(
0203 PodioReader(
0204 level=acts.logging.DEBUG,
0205 inputPath=str(args.edm4hep),
0206 outputFrame="events",
0207 category="events",
0208 )
0209 )
0210
0211 edm4hepReader = acts.examples.edm4hep.EDM4hepSimInputConverter(
0212 inputFrame="events",
0213 inputSimHits=[
0214 "PixelBarrelReadout",
0215 "PixelEndcapReadout",
0216 "ShortStripBarrelReadout",
0217 "ShortStripEndcapReadout",
0218 "LongStripBarrelReadout",
0219 "LongStripEndcapReadout",
0220 ],
0221 outputParticlesGenerator="particles_generated",
0222 outputParticlesSimulation="particles_simulated",
0223 outputSimHits="simhits",
0224 outputSimVertices="vertices_truth",
0225 dd4hepDetector=detector,
0226 trackingGeometry=trackingGeometry,
0227 sortSimHitsInTime=False,
0228 particleRMax=1080 * u.mm,
0229 particleZ=(-3030 * u.mm, 3030 * u.mm),
0230 particlePtMin=150 * u.MeV,
0231 level=acts.logging.DEBUG,
0232 )
0233 s.addAlgorithm(edm4hepReader)
0234
0235 s.addWhiteboardAlias("particles", edm4hepReader.config.outputParticlesSimulation)
0236
0237 addSimParticleSelection(
0238 s,
0239 ParticleSelectorConfig(
0240 rho=(0.0, 24 * u.mm),
0241 absZ=(0.0, 1.0 * u.m),
0242 eta=(-3.0, 3.0),
0243 removeNeutral=True,
0244 ),
0245 )
0246 else:
0247 if not args.ttbar:
0248 addParticleGun(
0249 s,
0250 MomentumConfig(
0251 args.gun_pt_range[0] * u.GeV,
0252 args.gun_pt_range[1] * u.GeV,
0253 transverse=True,
0254 ),
0255 EtaConfig(args.gun_eta_range[0], args.gun_eta_range[1]),
0256 PhiConfig(0.0, 360.0 * u.degree),
0257 ParticleConfig(
0258 args.gun_particles, acts.PdgParticle.eMuon, randomizeCharge=True
0259 ),
0260 vtxGen=acts.examples.GaussianVertexGenerator(
0261 mean=acts.Vector4(0, 0, 0, 0),
0262 stddev=acts.Vector4(
0263 0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 1.0 * u.ns
0264 ),
0265 ),
0266 multiplicity=args.gun_multiplicity,
0267 rnd=rnd,
0268 )
0269 else:
0270 addPythia8(
0271 s,
0272 hardProcess=["Top:qqbar2ttbar=on"],
0273 npileup=args.ttbar_pu,
0274 vtxGen=acts.examples.GaussianVertexGenerator(
0275 mean=acts.Vector4(0, 0, 0, 0),
0276 stddev=acts.Vector4(
0277 0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 5.0 * u.ns
0278 ),
0279 ),
0280 rnd=rnd,
0281 outputDirRoot=outputDir if args.output_root else None,
0282 outputDirCsv=outputDir if args.output_csv else None,
0283 )
0284
0285 addGenParticleSelection(
0286 s,
0287 ParticleSelectorConfig(
0288 rho=(0.0, 24 * u.mm),
0289 absZ=(0.0, 1.0 * u.m),
0290 eta=(-3.0, 3.0),
0291 pt=(150 * u.MeV, None),
0292 ),
0293 )
0294
0295 if args.geant4:
0296 if s.config.numThreads != 1:
0297 raise ValueError("Geant 4 simulation does not support multi-threading")
0298
0299
0300
0301
0302 addGeant4(
0303 s,
0304 detector,
0305 trackingGeometry,
0306 field,
0307 outputDirRoot=outputDir if args.output_root else None,
0308 outputDirCsv=outputDir if args.output_csv else None,
0309 outputDirObj=outputDir if args.output_obj else None,
0310 rnd=rnd,
0311 killVolume=trackingGeometry.highestTrackingVolume,
0312 killAfterTime=25 * u.ns,
0313 )
0314 else:
0315 addFatras(
0316 s,
0317 trackingGeometry,
0318 field,
0319 enableInteractions=True,
0320 outputDirRoot=outputDir if args.output_root else None,
0321 outputDirCsv=outputDir if args.output_csv else None,
0322 outputDirObj=outputDir if args.output_obj else None,
0323 rnd=rnd,
0324 )
0325
0326 addDigitization(
0327 s,
0328 trackingGeometry,
0329 field,
0330 digiConfigFile=oddDigiConfig,
0331 outputDirRoot=outputDir if args.output_root else None,
0332 outputDirCsv=outputDir if args.output_csv else None,
0333 rnd=rnd,
0334 )
0335
0336 addDigiParticleSelection(
0337 s,
0338 ParticleSelectorConfig(
0339 pt=(1.0 * u.GeV, None),
0340 eta=(-3.0, 3.0),
0341 measurements=(9, None),
0342 removeNeutral=True,
0343 ),
0344 )
0345
0346 if args.reco:
0347 addSeeding(
0348 s,
0349 trackingGeometry,
0350 field,
0351 initialSigmas=[
0352 1 * u.mm,
0353 1 * u.mm,
0354 1 * u.degree,
0355 1 * u.degree,
0356 0 * u.e / u.GeV,
0357 1 * u.ns,
0358 ],
0359 initialSigmaQoverPt=0.1 * u.e / u.GeV,
0360 initialSigmaPtRel=0.1,
0361 initialVarInflation=[1.0] * 6,
0362 particleHypothesis=acts.ParticleHypothesis.muon,
0363 geoSelectionConfigFile=oddSeedingSel,
0364 outputDirRoot=outputDir if args.output_root else None,
0365 outputDirCsv=outputDir if args.output_csv else None,
0366 )
0367
0368 if seedFilter_ML:
0369 addSeedFilterML(
0370 s,
0371 SeedFilterMLDBScanConfig(
0372 epsilonDBScan=0.03, minPointsDBScan=2, minSeedScore=0.1
0373 ),
0374 onnxModelFile=os.path.dirname(__file__)
0375 + "/MLAmbiguityResolution/seedDuplicateClassifier.onnx",
0376 outputDirRoot=outputDir if args.output_root else None,
0377 outputDirCsv=outputDir if args.output_csv else None,
0378 )
0379
0380 addCKFTracks(
0381 s,
0382 trackingGeometry,
0383 field,
0384 TrackSelectorConfig(
0385 pt=(1.0 * u.GeV if args.ttbar else 0.0, None),
0386 absEta=(None, 3.0),
0387 loc0=(-4.0 * u.mm, 4.0 * u.mm),
0388 nMeasurementsMin=7,
0389 maxHoles=2,
0390 maxOutliers=2,
0391 ),
0392 CkfConfig(
0393 chi2CutOffMeasurement=15.0,
0394 chi2CutOffOutlier=25.0,
0395 numMeasurementsCutOff=2,
0396 seedDeduplication=True,
0397 stayOnSeed=True,
0398 pixelVolumes=[16, 17, 18],
0399 stripVolumes=[23, 24, 25],
0400 maxPixelHoles=1,
0401 maxStripHoles=2,
0402 constrainToVolumes=[
0403 2,
0404 32,
0405 4,
0406 16,
0407 17,
0408 18,
0409 20,
0410 23,
0411 24,
0412 25,
0413 26,
0414 8,
0415 28,
0416 29,
0417 30,
0418 ],
0419 ),
0420 outputDirRoot=outputDir if args.output_root else None,
0421 outputDirCsv=outputDir if args.output_csv else None,
0422 writeCovMat=True,
0423 )
0424
0425 if ambi_ML:
0426 addAmbiguityResolutionML(
0427 s,
0428 AmbiguityResolutionMLConfig(
0429 maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0430 ),
0431 outputDirRoot=outputDir if args.output_root else None,
0432 outputDirCsv=outputDir if args.output_csv else None,
0433 onnxModelFile=os.path.dirname(__file__)
0434 + "/MLAmbiguityResolution/duplicateClassifier.onnx",
0435 )
0436
0437 elif ambi_scoring:
0438 addScoreBasedAmbiguityResolution(
0439 s,
0440 ScoreBasedAmbiguityResolutionConfig(
0441 minScore=0,
0442 minScoreSharedTracks=1,
0443 maxShared=2,
0444 minUnshared=3,
0445 maxSharedTracksPerMeasurement=2,
0446 useAmbiguityScoring=False,
0447 ),
0448 outputDirRoot=outputDir if args.output_root else None,
0449 outputDirCsv=outputDir if args.output_csv else None,
0450 ambiVolumeFile=ambi_config,
0451 writeCovMat=True,
0452 )
0453 else:
0454 addAmbiguityResolution(
0455 s,
0456 AmbiguityResolutionConfig(
0457 maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0458 ),
0459 outputDirRoot=outputDir if args.output_root else None,
0460 outputDirCsv=outputDir if args.output_csv else None,
0461 writeCovMat=True,
0462 )
0463
0464 addVertexFitting(
0465 s,
0466 field,
0467 vertexFinder=VertexFinder.AMVF,
0468 outputDirRoot=outputDir if args.output_root else None,
0469 outputDirCsv=outputDir if args.output_csv else None,
0470 )
0471
0472 if args.output_parquet:
0473 try:
0474 from acts.arrow import (
0475 particleSchema,
0476 simHitSchema,
0477 trackSchema,
0478 )
0479 from acts.examples.arrow import (
0480 ArrowParticleOutputConverter,
0481 ArrowSimHitOutputConverter,
0482 ArrowTrackOutputConverter,
0483 makeVolumeIdDetectorResolver,
0484 ParquetWriter,
0485 )
0486 except ImportError as e:
0487 raise RuntimeError(
0488 "Parquet output requested but acts.examples.arrow is not available; "
0489 "rebuild with ACTS_BUILD_EXAMPLES_PARQUET=ON."
0490 ) from e
0491
0492
0493
0494 _odd_detector_resolver = makeVolumeIdDetectorResolver(
0495 {
0496 7: 0,
0497 8: 1,
0498 9: 2,
0499 12: 3,
0500 13: 4,
0501 14: 5,
0502 16: 6,
0503 17: 7,
0504 18: 8,
0505 },
0506 255,
0507 )
0508
0509
0510
0511 arrParticleConv = ArrowParticleOutputConverter(
0512 level=acts.logging.INFO,
0513 inputParticles="particles_simulated",
0514 outputTable="particles_arrow",
0515 )
0516 s.addAlgorithm(arrParticleConv)
0517
0518 arrSimHitConv = ArrowSimHitOutputConverter(
0519 level=acts.logging.INFO,
0520 inputSimHits="simhits",
0521 inputParticles="particles_simulated",
0522 inputClusters="clusters",
0523 inputSimHitMeasurementsMap="simhit_measurements_map",
0524 outputTable="simhits_arrow",
0525 detectorResolver=_odd_detector_resolver,
0526 )
0527 s.addAlgorithm(arrSimHitConv)
0528
0529 if args.reco:
0530 arrTrackConv = ArrowTrackOutputConverter(
0531 level=acts.logging.INFO,
0532 inputTracks="tracks",
0533 inputTrackParticleMatching="track_particle_matching",
0534 inputParticles="particles_simulated",
0535 inputMeasurementSimHitsMap="measurement_simhits_map",
0536 outputTable="tracks_arrow",
0537 )
0538 s.addAlgorithm(arrTrackConv)
0539
0540 s.addWriter(
0541 ParquetWriter(
0542 level=acts.logging.INFO,
0543 outputDir=str(outputDir),
0544 collections={
0545 arrSimHitConv.config.outputTable: "simhits",
0546 arrTrackConv.config.outputTable: "tracks",
0547 arrParticleConv.config.outputTable: "particles",
0548 },
0549 expectedSchemas={
0550 arrSimHitConv.config.outputTable: simHitSchema(),
0551 arrTrackConv.config.outputTable: trackSchema(),
0552 arrParticleConv.config.outputTable: particleSchema(),
0553 },
0554 )
0555 )
0556
0557 s.run()