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(locals())  # 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
0190     edm4hepReader = acts.examples.edm4hep.EDM4hepReader(
0191         inputPath=str(args.edm4hep),
0192         inputSimHits=[
0193             "PixelBarrelReadout",
0194             "PixelEndcapReadout",
0195             "ShortStripBarrelReadout",
0196             "ShortStripEndcapReadout",
0197             "LongStripBarrelReadout",
0198             "LongStripEndcapReadout",
0199         ],
0200         outputParticlesGenerator="particles_generated",
0201         outputParticlesSimulation="particles_simulated",
0202         outputSimHits="simhits",
0203         graphvizOutput="graphviz",
0204         dd4hepDetector=detector,
0205         trackingGeometry=trackingGeometry,
0206         sortSimHitsInTime=True,
0207         level=acts.logging.INFO,
0208     )
0209     s.addReader(edm4hepReader)
0210
0211     s.addWhiteboardAlias("particles", edm4hepReader.config.outputParticlesGenerator)
0212
0213     addSimParticleSelection(
0214         s,
0215         ParticleSelectorConfig(
0216             rho=(0.0, 24 * u.mm),
0217             absZ=(0.0, 1.0 * u.m),
0218             eta=(-3.0, 3.0),
0219             pt=(150 * u.MeV, None),
0220             removeNeutral=True,
0221         ),
0222     )
0223 else:
0224     if not args.ttbar:
0225         addParticleGun(
0226             s,
0227             MomentumConfig(
0228                 args.gun_pt_range[0] * u.GeV,
0229                 args.gun_pt_range[1] * u.GeV,
0230                 transverse=True,
0231             ),
0232             EtaConfig(args.gun_eta_range[0], args.gun_eta_range[1]),
0233             PhiConfig(0.0, 360.0 * u.degree),
0234             ParticleConfig(
0235                 args.gun_particles, acts.PdgParticle.eMuon, randomizeCharge=True
0236             ),
0237             vtxGen=acts.examples.GaussianDisplacedVertexPositionGenerator(
0238                 rMean=50.0,
0239                 rStdDev=50.0 * u.mm,
0240                 zMean=2,
0241                 zStdDev=55.5 * u.mm,
0242                 tMean=0,
0243                 tStdDev=1.0 * u.ns,
0244             ),
0245             multiplicity=args.gun_multiplicity,
0246             rnd=rnd,
0247         )
0248     else:
0249         addPythia8(
0250             s,
0251             hardProcess=["Top:qqbar2ttbar=on"],
0252             npileup=args.ttbar_pu,
0253             vtxGen=acts.examples.GaussianDisplacedVertexPositionGenerator(
0254                 mean=acts.Vector4(0, 0, 0, 0),
0255                 stddev=acts.Vector4(
0256                     0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 5.0 * u.ns
0257                 ),
0258             ),
0259             rnd=rnd,
0260             outputDirRoot=outputDir if args.output_root else None,
0261             outputDirCsv=outputDir if args.output_csv else None,
0262         )
0263
0264         addGenParticleSelection(
0265             s,
0266             ParticleSelectorConfig(
0267                 rho=(0.0, 24 * u.mm),
0268                 absZ=(0.0, 1.0 * u.m),
0269                 eta=(-3.0, 3.0),
0270                 pt=(150 * u.MeV, None),
0271             ),
0272         )
0273
0274     if args.geant4:
0275         if s.config.numThreads != 1:
0276             raise ValueError("Geant 4 simulation does not support multi-threading")
0277
0278         # Pythia can sometime simulate particles outside the world volume, a cut on the Z of the track help mitigate this effect
0279         # Older version of G4 might not work, this as has been tested on version `geant4-11-00-patch-03`
0280         # For more detail see issue #1578
0281         addGeant4(
0282             s,
0283             detector,
0284             trackingGeometry,
0285             field,
0286             outputDirRoot=outputDir if args.output_root else None,
0287             outputDirCsv=outputDir if args.output_csv else None,
0288             outputDirObj=outputDir if args.output_obj else None,
0289             rnd=rnd,
0290             killVolume=trackingGeometry.highestTrackingVolume,
0291             killAfterTime=25 * u.ns,
0292         )
0293     else:
0294         addFatras(
0295             s,
0296             trackingGeometry,
0297             field,
0298             enableInteractions=True,
0299             outputDirRoot=outputDir if args.output_root else None,
0300             outputDirCsv=outputDir if args.output_csv else None,
0301             outputDirObj=outputDir if args.output_obj else None,
0302             rnd=rnd,
0303         )
0304
0305 addDigitization(
0306     s,
0307     trackingGeometry,
0308     field,
0309     digiConfigFile=oddDigiConfig,
0310     outputDirRoot=outputDir if args.output_root else None,
0311     outputDirCsv=outputDir if args.output_csv else None,
0312     rnd=rnd,
0313 )
0314
0315 addDigiParticleSelection(
0316     s,
0317     ParticleSelectorConfig(
0318         pt=(1.0 * u.GeV, None),
0319         eta=(-3.0, 3.0),
0320         measurements=(9, None),
0321         removeNeutral=True,
0322     ),
0323 )
0324
0325 if args.reco:
0326     addSeeding(
0327         s,
0328         trackingGeometry,
0329         field,
0330         initialSigmas=[
0331             1 * u.mm,
0332             1 * u.mm,
0333             1 * u.degree,
0334             1 * u.degree,
0335             0 * u.e / u.GeV,
0336             1 * u.ns,
0337         ],
0338         initialSigmaQoverPt=0.1 * u.e / u.GeV,
0339         initialSigmaPtRel=0.1,
0340         initialVarInflation=[1.0] * 6,
0341         geoSelectionConfigFile=oddSeedingSel,
0342         outputDirRoot=outputDir if args.output_root else None,
0343         outputDirCsv=outputDir if args.output_csv else None,
0344     )
0345
0346     if seedFilter_ML:
0347         addSeedFilterML(
0348             s,
0349             SeedFilterMLDBScanConfig(
0350                 epsilonDBScan=0.03, minPointsDBScan=2, minSeedScore=0.1
0351             ),
0352             onnxModelFile=os.path.dirname(__file__)
0353             + "/MLAmbiguityResolution/seedDuplicateClassifier.onnx",
0354             outputDirRoot=outputDir if args.output_root else None,
0355             outputDirCsv=outputDir if args.output_csv else None,
0356         )
0357
0358     addCKFTracks(
0359         s,
0360         trackingGeometry,
0361         field,
0362         TrackSelectorConfig(
0363             pt=(1.0 * u.GeV if args.ttbar else 0.0, None),
0364             absEta=(None, 3.0),
0365             loc0=(-4.0 * u.mm, 4.0 * u.mm),
0366             nMeasurementsMin=7,
0367             maxHoles=2,
0368             maxOutliers=2,
0369         ),
0370         CkfConfig(
0371             chi2CutOffMeasurement=15.0,
0372             chi2CutOffOutlier=25.0,
0373             numMeasurementsCutOff=10,
0374             seedDeduplication=True,
0375             stayOnSeed=True,
0376             pixelVolumes=[16, 17, 18],
0377             stripVolumes=[23, 24, 25],
0378             maxPixelHoles=1,
0379             maxStripHoles=2,
0380             constrainToVolumes=[
0381                 2,  # beam pipe
0382                 32,
0383                 4,  # beam pip gap
0384                 16,
0385                 17,
0386                 18,  # pixel
0387                 20,  # PST
0388                 23,
0389                 24,
0390                 25,  # short strip
0391                 26,
0392                 8,  # long strip gap
0393                 28,
0394                 29,
0395                 30,  # long strip
0396             ],
0397         ),
0398         outputDirRoot=outputDir if args.output_root else None,
0399         outputDirCsv=outputDir if args.output_csv else None,
0400         writeCovMat=True,
0401     )
0402
0403     if ambi_ML:
0404         addAmbiguityResolutionML(
0405             s,
0406             AmbiguityResolutionMLConfig(
0407                 maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0408             ),
0409             outputDirRoot=outputDir if args.output_root else None,
0410             outputDirCsv=outputDir if args.output_csv else None,
0411             onnxModelFile=os.path.dirname(__file__)
0412             + "/MLAmbiguityResolution/duplicateClassifier.onnx",
0413         )
0414
0415     elif ambi_scoring:
0416         addScoreBasedAmbiguityResolution(
0417             s,
0418             ScoreBasedAmbiguityResolutionConfig(
0419                 minScore=0,
0420                 minScoreSharedTracks=1,
0421                 maxShared=2,
0422                 minUnshared=3,
0423                 maxSharedTracksPerMeasurement=2,
0424                 useAmbiguityScoring=False,
0425             ),
0426             outputDirRoot=outputDir if args.output_root else None,
0427             outputDirCsv=outputDir if args.output_csv else None,
0428             ambiVolumeFile=ambi_config,
0429             writeCovMat=True,
0430         )
0431     else:
0432         addAmbiguityResolution(
0433             s,
0434             AmbiguityResolutionConfig(
0435                 maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0436             ),
0437             outputDirRoot=outputDir if args.output_root else None,
0438             outputDirCsv=outputDir if args.output_csv else None,
0439             writeCovMat=True,
0440         )
0441
0442     addVertexFitting(
0443         s,
0444         field,
0445         vertexFinder=VertexFinder.AMVF,
0446         outputDirRoot=outputDir if args.output_root else None,
0447     )
0448
0449 s.run()