Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-08-06 08:28:58

0001 #!/usr/bin/env python3
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 # acts.examples.dump_args_calls()  # show python binding calls
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         # Pythia can sometime simulate particles outside the world volume, a cut on the Z of the track help mitigate this effect
0300         # Older version of G4 might not work, this as has been tested on version `geant4-11-00-patch-03`
0301         # For more detail see issue #1578
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,  # beam pipe
0404                 32,
0405                 4,  # beam pip gap
0406                 16,
0407                 17,
0408                 18,  # pixel
0409                 20,  # PST
0410                 23,
0411                 24,
0412                 25,  # short strip
0413                 26,
0414                 8,  # long strip gap
0415                 28,
0416                 29,
0417                 30,  # long strip
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     # ODD volume -> detector enum mapping used in parquet simhit export.
0493     # 0..8 are subdetector-specific enums; 255 marks unknown/unmatched.
0494     _odd_detector_resolver = makeVolumeIdDetectorResolver(
0495         {
0496             7: 0,  # pixel_neg_endcap
0497             8: 1,  # pixel_barrel
0498             9: 2,  # pixel_pos_endcap
0499             12: 3,  # short_neg_endcap
0500             13: 4,  # short_barrel
0501             14: 5,  # short_pos_endcap
0502             16: 6,  # long_neg_endcap
0503             17: 7,  # long_barrel
0504             18: 8,  # long_pos_endcap
0505         },
0506         255,
0507     )
0508 
0509     # Each converter parks an arrow::Table on the whiteboard under a fresh
0510     # key, and one ParquetWriter picks them all up.
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()