Scripts/Python/full_chain_odd_LRT.py

0001 #!/usr/bin/env python3
0002
0003 import os
0004 import argparse
0005 import pathlib
0006 import math
0007
0008 import acts
0009 import acts.examples
0010 from acts.examples.simulation import (
0011     addParticleGun,
0012     MomentumConfig,
0013     EtaConfig,
0014     PhiConfig,
0015     ParticleConfig,
0016     ParticleSelectorConfig,
0017     addPythia8,
0018     ParticleSelectorConfig,
0019     addGenParticleSelection,
0020     addFatras,
0021     addGeant4,
0022     addSimParticleSelection,
0023     addDigitization,
0024     addDigiParticleSelection,
0025 )
0026 from acts.examples.reconstruction import (
0027     addSeeding,
0028     CkfConfig,
0029     addCKFTracks,
0030     TrackSelectorConfig,
0031     addAmbiguityResolution,
0032     AmbiguityResolutionConfig,
0033     addAmbiguityResolutionML,
0034     AmbiguityResolutionMLConfig,
0035     addScoreBasedAmbiguityResolution,
0036     ScoreBasedAmbiguityResolutionConfig,
0037     addVertexFitting,
0038     VertexFinder,
0039     addSeedFilterML,
0040     SeedFilterMLDBScanConfig,
0041 )
0042 from acts.examples.odd import getOpenDataDetector, getOpenDataDetectorDirectory
0043
0044 u = acts.UnitConstants
0045
0046
0047 parser = argparse.ArgumentParser(description="Full chain with the OpenDataDetector")
0048 parser.add_argument(
0049     "--output",
0050     "-o",
0051     help="Output directory",
0052     type=pathlib.Path,
0053     default=pathlib.Path.cwd() / "odd_output",
0054 )
0055 parser.add_argument("--events", "-n", help="Number of events", type=int, default=100)
0056 parser.add_argument("--skip", "-s", help="Number of events", type=int, default=0)
0057 parser.add_argument("--edm4hep", help="Use edm4hep inputs", type=pathlib.Path)
0058 parser.add_argument(
0059     "--geant4", help="Use Geant4 instead of fatras", action="store_true"
0060 )
0061 parser.add_argument(
0062     "--ttbar",
0063     help="Use Pythia8 (ttbar, pile-up 200) instead of particle gun",
0064     action="store_true",
0065 )
0066 parser.add_argument(
0067     "--ttbar-pu",
0068     help="Number of pile-up events for ttbar",
0069     type=int,
0070     default=200,
0071 )
0072 parser.add_argument(
0073     "--gun-particles",
0074     help="Multiplicity (no. of particles) of the particle gun",
0075     type=int,
0076     default=4,
0077 )
0078 parser.add_argument(
0079     "--gun-multiplicity",
0080     help="Multiplicity (no. of vertices) of the particle gun",
0081     type=int,
0082     default=200,
0083 )
0084 parser.add_argument(
0085     "--gun-eta-range",
0086     nargs=2,
0087     help="Eta range of the particle gun",
0088     type=float,
0089     default=[-3.0, 3.0],
0090 )
0091 parser.add_argument(
0092     "--gun-pt-range",
0093     nargs=2,
0094     help="Pt range of the particle gun (GeV)",
0095     type=float,
0096     default=[1.0 * u.GeV, 10.0 * u.GeV],
0097 )
0098 parser.add_argument(
0099     "--digi-config", help="Digitization configuration file", type=pathlib.Path
0100 )
0101 parser.add_argument(
0102     "--material-config", help="Material map configuration file", type=pathlib.Path
0103 )
0104 parser.add_argument(
0105     "--ambi-solver",
0106     help="Set which ambiguity solver to use, default is the classical one",
0107     type=str,
0108     choices=["greedy", "scoring", "ML"],
0109     default="greedy",
0110 )
0111 parser.add_argument(
0112     "--ambi-config",
0113     help="Set the configuration file for the Score Based ambiguity resolution",
0114     type=pathlib.Path,
0115     default=pathlib.Path.cwd() / "ambi_config.json",
0116 )
0117
0118 parser.add_argument(
0119     "--MLSeedFilter",
0120     help="Use the Ml seed filter to select seed after the seeding step",
0121     action="store_true",
0122 )
0123 parser.add_argument(
0124     "--reco",
0125     help="Switch reco on/off",
0126     default=True,
0127     action=argparse.BooleanOptionalAction,
0128 )
0129 parser.add_argument(
0130     "--output-root",
0131     help="Switch root output on/off",
0132     default=True,
0133     action=argparse.BooleanOptionalAction,
0134 )
0135 parser.add_argument(
0136     "--output-csv",
0137     help="Switch csv output on/off",
0138     default=True,
0139     action=argparse.BooleanOptionalAction,
0140 )
0141 parser.add_argument(
0142     "--output-obj",
0143     help="Switch obj output on/off",
0144     default=True,
0145     action=argparse.BooleanOptionalAction,
0146 )
0147
0148 args = parser.parse_args()
0149
0150 outputDir = args.output
0151 ambi_ML = args.ambi_solver == "ML"
0152 ambi_scoring = args.ambi_solver == "scoring"
0153 ambi_config = args.ambi_config
0154 seedFilter_ML = args.MLSeedFilter
0155 geoDir = getOpenDataDetectorDirectory()
0156 actsDir = pathlib.Path(__file__).parent.parent.parent.parent
0157 # acts.examples.dump_args_calls()  # show python binding calls
0158
0159 oddMaterialMap = (
0160     args.material_config
0161     if args.material_config
0162     else geoDir / "data/odd-material-maps.root"
0163 )
0164
0165 oddDigiConfig = (
0166     args.digi_config
0167     if args.digi_config
0168     else actsDir / "Examples/Configs/odd-digi-smearing-config.json"
0169 )
0170
0171 oddSeedingSel = actsDir / "Examples/Configs/odd-seeding-config.json"
0172 oddMaterialDeco = acts.IMaterialDecorator.fromFile(oddMaterialMap)
0173
0174 detector = getOpenDataDetector(odd_dir=geoDir, materialDecorator=oddMaterialDeco)
0175 trackingGeometry = detector.trackingGeometry()
0176 decorators = detector.contextDecorators()
0177 field = acts.ConstantBField(acts.Vector3(0.0, 0.0, 2.0 * u.T))
0178 rnd = acts.examples.RandomNumbers(seed=42)
0179
0180 s = acts.examples.Sequencer(
0181     events=args.events,
0182     skip=args.skip,
0183     numThreads=1 if args.geant4 else -1,
0184     outputDir=str(outputDir),
0185 )
0186
0187 if args.edm4hep:
0188     import acts.examples.edm4hep
0189     from acts.examples.edm4hep import PodioReader
0190
0191     s.addReader(
0192         PodioReader(
0193             level=acts.logging.DEBUG,
0194             inputPath=str(args.edm4hep),
0195             outputFrame="events",
0196             category="events",
0197         )
0198     )
0199
0200     edm4hepReader = acts.examples.edm4hep.EDM4hepSimInputConverter(
0201         inputFrame="events",
0202         inputSimHits=[
0203             "PixelBarrelReadout",
0204             "PixelEndcapReadout",
0205             "ShortStripBarrelReadout",
0206             "ShortStripEndcapReadout",
0207             "LongStripBarrelReadout",
0208             "LongStripEndcapReadout",
0209         ],
0210         outputParticlesGenerator="particles_generated",
0211         outputParticlesSimulation="particles_simulated",
0212         outputSimHits="simhits",
0213         outputSimVertices="vertices_truth",
0214         dd4hepDetector=detector,
0215         trackingGeometry=trackingGeometry,
0216         sortSimHitsInTime=False,
0217         particleRMax=1080 * u.mm,
0218         particleZ=(-3030 * u.mm, 3030 * u.mm),
0219         particlePtMin=150 * u.MeV,
0220         level=acts.logging.DEBUG,
0221     )
0222     s.addAlgorithm(edm4hepReader)
0223
0224     s.addWhiteboardAlias("particles", edm4hepReader.config.outputParticlesSimulation)
0225
0226     addSimParticleSelection(
0227         s,
0228         ParticleSelectorConfig(
0229             rho=(0.0, 24 * u.mm),
0230             absZ=(0.0, 1.0 * u.m),
0231             eta=(-3.0, 3.0),
0232             removeNeutral=True,
0233         ),
0234     )
0235 else:
0236     if not args.ttbar:
0237         addParticleGun(
0238             s,
0239             MomentumConfig(
0240                 args.gun_pt_range[0] * u.GeV,
0241                 args.gun_pt_range[1] * u.GeV,
0242                 transverse=True,
0243             ),
0244             EtaConfig(args.gun_eta_range[0], args.gun_eta_range[1]),
0245             PhiConfig(0.0, 360.0 * u.degree),
0246             ParticleConfig(
0247                 args.gun_particles, acts.PdgParticle.eMuon, randomizeCharge=True
0248             ),
0249             vtxGen=acts.examples.GaussianDisplacedVertexPositionGenerator(
0250                 rMean=50.0,
0251                 rStdDev=50.0 * u.mm,
0252                 zMean=2,
0253                 zStdDev=55.5 * u.mm,
0254                 tMean=0,
0255                 tStdDev=1.0 * u.ns,
0256             ),
0257             multiplicity=args.gun_multiplicity,
0258             rnd=rnd,
0259         )
0260     else:
0261         addPythia8(
0262             s,
0263             hardProcess=["Top:qqbar2ttbar=on"],
0264             npileup=args.ttbar_pu,
0265             vtxGen=acts.examples.GaussianDisplacedVertexPositionGenerator(
0266                 rMean=50.0,
0267                 rStdDev=50.0 * u.mm,
0268                 zMean=2,
0269                 zStdDev=55.5 * u.mm,
0270                 tMean=0,
0271                 tStdDev=5.0 * u.ns,
0272             ),
0273             rnd=rnd,
0274             outputDirRoot=outputDir if args.output_root else None,
0275             outputDirCsv=outputDir if args.output_csv else None,
0276         )
0277
0278         addGenParticleSelection(
0279             s,
0280             ParticleSelectorConfig(
0281                 rho=(0.0, 24 * u.mm),
0282                 absZ=(0.0, 1.0 * u.m),
0283                 eta=(-3.0, 3.0),
0284                 pt=(150 * u.MeV, None),
0285             ),
0286         )
0287
0288     if args.geant4:
0289         if s.config.numThreads != 1:
0290             raise ValueError("Geant 4 simulation does not support multi-threading")
0291
0292         # Pythia can sometime simulate particles outside the world volume, a cut on the Z of the track help mitigate this effect
0293         # Older version of G4 might not work, this as has been tested on version `geant4-11-00-patch-03`
0294         # For more detail see issue #1578
0295         addGeant4(
0296             s,
0297             detector,
0298             trackingGeometry,
0299             field,
0300             outputDirRoot=outputDir if args.output_root else None,
0301             outputDirCsv=outputDir if args.output_csv else None,
0302             outputDirObj=outputDir if args.output_obj else None,
0303             rnd=rnd,
0304             killVolume=trackingGeometry.highestTrackingVolume,
0305             killAfterTime=25 * u.ns,
0306         )
0307     else:
0308         addFatras(
0309             s,
0310             trackingGeometry,
0311             field,
0312             enableInteractions=True,
0313             outputDirRoot=outputDir if args.output_root else None,
0314             outputDirCsv=outputDir if args.output_csv else None,
0315             outputDirObj=outputDir if args.output_obj else None,
0316             rnd=rnd,
0317         )
0318
0319 addDigitization(
0320     s,
0321     trackingGeometry,
0322     field,
0323     digiConfigFile=oddDigiConfig,
0324     outputDirRoot=outputDir if args.output_root else None,
0325     outputDirCsv=outputDir if args.output_csv else None,
0326     rnd=rnd,
0327 )
0328
0329 addDigiParticleSelection(
0330     s,
0331     ParticleSelectorConfig(
0332         pt=(1.0 * u.GeV, None),
0333         eta=(-3.0, 3.0),
0334         measurements=(9, None),
0335         removeNeutral=True,
0336     ),
0337 )
0338
0339 if args.reco:
0340     addSeeding(
0341         s,
0342         trackingGeometry,
0343         field,
0344         initialSigmas=[
0345             1 * u.mm,
0346             1 * u.mm,
0347             1 * u.degree,
0348             1 * u.degree,
0349             0 * u.e / u.GeV,
0350             1 * u.ns,
0351         ],
0352         initialSigmaQoverPt=0.1 * u.e / u.GeV,
0353         initialSigmaPtRel=0.1,
0354         initialVarInflation=[1.0] * 6,
0355         particleHypothesis=acts.ParticleHypothesis.muon,
0356         geoSelectionConfigFile=oddSeedingSel,
0357         outputDirRoot=outputDir if args.output_root else None,
0358         outputDirCsv=outputDir if args.output_csv else None,
0359         logLevel=acts.logging.DEBUG,
0360     )
0361
0362     if seedFilter_ML:
0363         addSeedFilterML(
0364             s,
0365             SeedFilterMLDBScanConfig(
0366                 epsilonDBScan=0.03, minPointsDBScan=2, minSeedScore=0.1
0367             ),
0368             onnxModelFile=os.path.dirname(__file__)
0369             + "/MLAmbiguityResolution/seedDuplicateClassifier.onnx",
0370             outputDirRoot=outputDir if args.output_root else None,
0371             outputDirCsv=outputDir if args.output_csv else None,
0372         )
0373
0374     addCKFTracks(
0375         s,
0376         trackingGeometry,
0377         field,
0378         TrackSelectorConfig(
0379             pt=(1.0 * u.GeV if args.ttbar else 0.0, None),
0380             absEta=(None, 3.0),
0381             loc0=(-4.0 * u.mm, 4.0 * u.mm),
0382             nMeasurementsMin=7,
0383             maxHoles=2,
0384             maxOutliers=2,
0385         ),
0386         CkfConfig(
0387             chi2CutOffMeasurement=15.0,
0388             chi2CutOffOutlier=25.0,
0389             numMeasurementsCutOff=2,
0390             seedDeduplication=True,
0391             stayOnSeed=True,
0392             pixelVolumes=[16, 17, 18],
0393             stripVolumes=[23, 24, 25],
0394             maxPixelHoles=1,
0395             maxStripHoles=2,
0396             constrainToVolumes=[
0397                 2,  # beam pipe
0398                 32,
0399                 4,  # beam pip gap
0400                 16,
0401                 17,
0402                 18,  # pixel
0403                 20,  # PST
0404                 23,
0405                 24,
0406                 25,  # short strip
0407                 26,
0408                 8,  # long strip gap
0409                 28,
0410                 29,
0411                 30,  # long strip
0412             ],
0413         ),
0414         outputDirRoot=outputDir if args.output_root else None,
0415         outputDirCsv=outputDir if args.output_csv else None,
0416         writeCovMat=True,
0417         logLevel=acts.logging.DEBUG,
0418     )
0419
0420     if ambi_ML:
0421         addAmbiguityResolutionML(
0422             s,
0423             AmbiguityResolutionMLConfig(
0424                 maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0425             ),
0426             outputDirRoot=outputDir if args.output_root else None,
0427             outputDirCsv=outputDir if args.output_csv else None,
0428             onnxModelFile=os.path.dirname(__file__)
0429             + "/MLAmbiguityResolution/duplicateClassifier.onnx",
0430         )
0431
0432     elif ambi_scoring:
0433         addScoreBasedAmbiguityResolution(
0434             s,
0435             ScoreBasedAmbiguityResolutionConfig(
0436                 minScore=0,
0437                 minScoreSharedTracks=1,
0438                 maxShared=2,
0439                 minUnshared=3,
0440                 maxSharedTracksPerMeasurement=2,
0441                 useAmbiguityScoring=False,
0442             ),
0443             outputDirRoot=outputDir if args.output_root else None,
0444             outputDirCsv=outputDir if args.output_csv else None,
0445             ambiVolumeFile=ambi_config,
0446             writeCovMat=True,
0447         )
0448     else:
0449         addAmbiguityResolution(
0450             s,
0451             AmbiguityResolutionConfig(
0452                 maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0453             ),
0454             outputDirRoot=outputDir if args.output_root else None,
0455             outputDirCsv=outputDir if args.output_csv else None,
0456             writeCovMat=True,
0457         )
0458
0459     addVertexFitting(
0460         s,
0461         field,
0462         vertexFinder=VertexFinder.AMVF,
0463         outputDirRoot=outputDir if args.output_root else None,
0464         outputDirCsv=outputDir if args.output_csv else None,
0465     )
0466
0467 s.run()