File indexing completed on 2025-06-30 07:52:48
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("--edm4hep", help="Use edm4hep inputs", type=pathlib.Path)
0056 parser.add_argument(
0057 "--geant4", help="Use Geant4 instead of fatras", action="store_true"
0058 )
0059 parser.add_argument(
0060 "--ttbar",
0061 help="Use Pythia8 (ttbar, pile-up 200) instead of particle gun",
0062 action="store_true",
0063 )
0064 parser.add_argument(
0065 "--ttbar-pu",
0066 help="Number of pile-up events for ttbar",
0067 type=int,
0068 default=200,
0069 )
0070 parser.add_argument(
0071 "--gun-particles",
0072 help="Multiplicity (no. of particles) of the particle gun",
0073 type=int,
0074 default=4,
0075 )
0076 parser.add_argument(
0077 "--gun-multiplicity",
0078 help="Multiplicity (no. of vertices) of the particle gun",
0079 type=int,
0080 default=200,
0081 )
0082 parser.add_argument(
0083 "--gun-eta-range",
0084 nargs=2,
0085 help="Eta range of the particle gun",
0086 type=float,
0087 default=[-3.0, 3.0],
0088 )
0089 parser.add_argument(
0090 "--gun-pt-range",
0091 nargs=2,
0092 help="Pt range of the particle gun (GeV)",
0093 type=float,
0094 default=[1.0 * u.GeV, 10.0 * u.GeV],
0095 )
0096 parser.add_argument(
0097 "--digi-config", help="Digitization configuration file", type=pathlib.Path
0098 )
0099 parser.add_argument(
0100 "--material-config", help="Material map configuration file", type=pathlib.Path
0101 )
0102 parser.add_argument(
0103 "--ambi-solver",
0104 help="Set which ambiguity solver to use, default is the classical one",
0105 type=str,
0106 choices=["greedy", "scoring", "ML"],
0107 default="greedy",
0108 )
0109 parser.add_argument(
0110 "--ambi-config",
0111 help="Set the configuration file for the Score Based ambiguity resolution",
0112 type=pathlib.Path,
0113 default=pathlib.Path.cwd() / "ambi_config.json",
0114 )
0115
0116 parser.add_argument(
0117 "--MLSeedFilter",
0118 help="Use the Ml seed filter to select seed after the seeding step",
0119 action="store_true",
0120 )
0121 parser.add_argument(
0122 "--reco",
0123 help="Switch reco on/off",
0124 default=True,
0125 action=argparse.BooleanOptionalAction,
0126 )
0127 parser.add_argument(
0128 "--output-root",
0129 help="Switch root output on/off",
0130 default=True,
0131 action=argparse.BooleanOptionalAction,
0132 )
0133 parser.add_argument(
0134 "--output-csv",
0135 help="Switch csv output on/off",
0136 default=True,
0137 action=argparse.BooleanOptionalAction,
0138 )
0139 parser.add_argument(
0140 "--output-obj",
0141 help="Switch obj output on/off",
0142 default=True,
0143 action=argparse.BooleanOptionalAction,
0144 )
0145
0146 args = parser.parse_args()
0147
0148 outputDir = args.output
0149 ambi_ML = args.ambi_solver == "ML"
0150 ambi_scoring = args.ambi_solver == "scoring"
0151 ambi_config = args.ambi_config
0152 seedFilter_ML = args.MLSeedFilter
0153 geoDir = getOpenDataDetectorDirectory()
0154 actsDir = pathlib.Path(__file__).parent.parent.parent.parent
0155
0156
0157 oddMaterialMap = (
0158 args.material_config
0159 if args.material_config
0160 else geoDir / "data/odd-material-maps.root"
0161 )
0162
0163 oddDigiConfig = (
0164 args.digi_config
0165 if args.digi_config
0166 else actsDir / "Examples/Configs/odd-digi-smearing-config.json"
0167 )
0168
0169 oddSeedingSel = actsDir / "Examples/Configs/odd-seeding-config.json"
0170 oddMaterialDeco = acts.IMaterialDecorator.fromFile(oddMaterialMap)
0171
0172 detector = getOpenDataDetector(odd_dir=geoDir, materialDecorator=oddMaterialDeco)
0173 trackingGeometry = detector.trackingGeometry()
0174 decorators = detector.contextDecorators()
0175 field = acts.ConstantBField(acts.Vector3(0.0, 0.0, 2.0 * u.T))
0176 rnd = acts.examples.RandomNumbers(seed=42)
0177
0178 s = acts.examples.Sequencer(
0179 events=args.events,
0180 skip=args.skip,
0181 numThreads=1 if args.geant4 else -1,
0182 outputDir=str(outputDir),
0183 )
0184
0185 if args.edm4hep:
0186 import acts.examples.edm4hep
0187
0188 edm4hepReader = acts.examples.edm4hep.EDM4hepSimReader(
0189 inputPath=str(args.edm4hep),
0190 inputSimHits=[
0191 "PixelBarrelReadout",
0192 "PixelEndcapReadout",
0193 "ShortStripBarrelReadout",
0194 "ShortStripEndcapReadout",
0195 "LongStripBarrelReadout",
0196 "LongStripEndcapReadout",
0197 ],
0198 outputParticlesGenerator="particles_generated",
0199 outputParticlesSimulation="particles_simulated",
0200 outputSimHits="simhits",
0201 graphvizOutput="graphviz",
0202 dd4hepDetector=detector,
0203 trackingGeometry=trackingGeometry,
0204 sortSimHitsInTime=True,
0205 level=acts.logging.INFO,
0206 )
0207 s.addReader(edm4hepReader)
0208
0209 s.addWhiteboardAlias("particles", edm4hepReader.config.outputParticlesGenerator)
0210
0211 addSimParticleSelection(
0212 s,
0213 ParticleSelectorConfig(
0214 rho=(0.0, 24 * u.mm),
0215 absZ=(0.0, 1.0 * u.m),
0216 eta=(-3.0, 3.0),
0217 pt=(150 * u.MeV, None),
0218 removeNeutral=True,
0219 ),
0220 )
0221 else:
0222 if not args.ttbar:
0223 addParticleGun(
0224 s,
0225 MomentumConfig(
0226 args.gun_pt_range[0] * u.GeV,
0227 args.gun_pt_range[1] * u.GeV,
0228 transverse=True,
0229 ),
0230 EtaConfig(args.gun_eta_range[0], args.gun_eta_range[1]),
0231 PhiConfig(0.0, 360.0 * u.degree),
0232 ParticleConfig(
0233 args.gun_particles, acts.PdgParticle.eMuon, randomizeCharge=True
0234 ),
0235 vtxGen=acts.examples.GaussianVertexGenerator(
0236 mean=acts.Vector4(0, 0, 0, 0),
0237 stddev=acts.Vector4(
0238 0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 1.0 * u.ns
0239 ),
0240 ),
0241 multiplicity=args.gun_multiplicity,
0242 rnd=rnd,
0243 )
0244 else:
0245 addPythia8(
0246 s,
0247 hardProcess=["Top:qqbar2ttbar=on"],
0248 npileup=args.ttbar_pu,
0249 vtxGen=acts.examples.GaussianVertexGenerator(
0250 mean=acts.Vector4(0, 0, 0, 0),
0251 stddev=acts.Vector4(
0252 0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 5.0 * u.ns
0253 ),
0254 ),
0255 rnd=rnd,
0256 outputDirRoot=outputDir if args.output_root else None,
0257 outputDirCsv=outputDir if args.output_csv else None,
0258 )
0259
0260 addGenParticleSelection(
0261 s,
0262 ParticleSelectorConfig(
0263 rho=(0.0, 24 * u.mm),
0264 absZ=(0.0, 1.0 * u.m),
0265 eta=(-3.0, 3.0),
0266 pt=(150 * u.MeV, None),
0267 ),
0268 )
0269
0270 if args.geant4:
0271 if s.config.numThreads != 1:
0272 raise ValueError("Geant 4 simulation does not support multi-threading")
0273
0274
0275
0276
0277 addGeant4(
0278 s,
0279 detector,
0280 trackingGeometry,
0281 field,
0282 outputDirRoot=outputDir if args.output_root else None,
0283 outputDirCsv=outputDir if args.output_csv else None,
0284 outputDirObj=outputDir if args.output_obj else None,
0285 rnd=rnd,
0286 killVolume=trackingGeometry.highestTrackingVolume,
0287 killAfterTime=25 * u.ns,
0288 )
0289 else:
0290 addFatras(
0291 s,
0292 trackingGeometry,
0293 field,
0294 enableInteractions=True,
0295 outputDirRoot=outputDir if args.output_root else None,
0296 outputDirCsv=outputDir if args.output_csv else None,
0297 outputDirObj=outputDir if args.output_obj else None,
0298 rnd=rnd,
0299 )
0300
0301 addDigitization(
0302 s,
0303 trackingGeometry,
0304 field,
0305 digiConfigFile=oddDigiConfig,
0306 outputDirRoot=outputDir if args.output_root else None,
0307 outputDirCsv=outputDir if args.output_csv else None,
0308 rnd=rnd,
0309 )
0310
0311 addDigiParticleSelection(
0312 s,
0313 ParticleSelectorConfig(
0314 pt=(1.0 * u.GeV, None),
0315 eta=(-3.0, 3.0),
0316 measurements=(9, None),
0317 removeNeutral=True,
0318 ),
0319 )
0320
0321 if args.reco:
0322 addSeeding(
0323 s,
0324 trackingGeometry,
0325 field,
0326 initialSigmas=[
0327 1 * u.mm,
0328 1 * u.mm,
0329 1 * u.degree,
0330 1 * u.degree,
0331 0 * u.e / u.GeV,
0332 1 * u.ns,
0333 ],
0334 initialSigmaQoverPt=0.1 * u.e / u.GeV,
0335 initialSigmaPtRel=0.1,
0336 initialVarInflation=[1.0] * 6,
0337 geoSelectionConfigFile=oddSeedingSel,
0338 outputDirRoot=outputDir if args.output_root else None,
0339 outputDirCsv=outputDir if args.output_csv else None,
0340 )
0341
0342 if seedFilter_ML:
0343 addSeedFilterML(
0344 s,
0345 SeedFilterMLDBScanConfig(
0346 epsilonDBScan=0.03, minPointsDBScan=2, minSeedScore=0.1
0347 ),
0348 onnxModelFile=os.path.dirname(__file__)
0349 + "/MLAmbiguityResolution/seedDuplicateClassifier.onnx",
0350 outputDirRoot=outputDir if args.output_root else None,
0351 outputDirCsv=outputDir if args.output_csv else None,
0352 )
0353
0354 addCKFTracks(
0355 s,
0356 trackingGeometry,
0357 field,
0358 TrackSelectorConfig(
0359 pt=(1.0 * u.GeV if args.ttbar else 0.0, None),
0360 absEta=(None, 3.0),
0361 loc0=(-4.0 * u.mm, 4.0 * u.mm),
0362 nMeasurementsMin=7,
0363 maxHoles=2,
0364 maxOutliers=2,
0365 ),
0366 CkfConfig(
0367 chi2CutOffMeasurement=15.0,
0368 chi2CutOffOutlier=25.0,
0369 numMeasurementsCutOff=10,
0370 seedDeduplication=True,
0371 stayOnSeed=True,
0372 pixelVolumes=[16, 17, 18],
0373 stripVolumes=[23, 24, 25],
0374 maxPixelHoles=1,
0375 maxStripHoles=2,
0376 constrainToVolumes=[
0377 2,
0378 32,
0379 4,
0380 16,
0381 17,
0382 18,
0383 20,
0384 23,
0385 24,
0386 25,
0387 26,
0388 8,
0389 28,
0390 29,
0391 30,
0392 ],
0393 ),
0394 outputDirRoot=outputDir if args.output_root else None,
0395 outputDirCsv=outputDir if args.output_csv else None,
0396 writeCovMat=True,
0397 )
0398
0399 if ambi_ML:
0400 addAmbiguityResolutionML(
0401 s,
0402 AmbiguityResolutionMLConfig(
0403 maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0404 ),
0405 outputDirRoot=outputDir if args.output_root else None,
0406 outputDirCsv=outputDir if args.output_csv else None,
0407 onnxModelFile=os.path.dirname(__file__)
0408 + "/MLAmbiguityResolution/duplicateClassifier.onnx",
0409 )
0410
0411 elif ambi_scoring:
0412 addScoreBasedAmbiguityResolution(
0413 s,
0414 ScoreBasedAmbiguityResolutionConfig(
0415 minScore=0,
0416 minScoreSharedTracks=1,
0417 maxShared=2,
0418 minUnshared=3,
0419 maxSharedTracksPerMeasurement=2,
0420 useAmbiguityScoring=False,
0421 ),
0422 outputDirRoot=outputDir if args.output_root else None,
0423 outputDirCsv=outputDir if args.output_csv else None,
0424 ambiVolumeFile=ambi_config,
0425 writeCovMat=True,
0426 )
0427 else:
0428 addAmbiguityResolution(
0429 s,
0430 AmbiguityResolutionConfig(
0431 maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0432 ),
0433 outputDirRoot=outputDir if args.output_root else None,
0434 outputDirCsv=outputDir if args.output_csv else None,
0435 writeCovMat=True,
0436 )
0437
0438 addVertexFitting(
0439 s,
0440 field,
0441 vertexFinder=VertexFinder.AMVF,
0442 outputDirRoot=outputDir if args.output_root else None,
0443 )
0444
0445 s.run()
