Scripts/Python/full_chain_odd.py

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     addParticleGun,
0011     MomentumConfig,
0012     EtaConfig,
0013     PhiConfig,
0014     ParticleConfig,
0015     addPythia8,
0016     addFatras,
0017     addGeant4,
0018     ParticleSelectorConfig,
0019     addDigitization,
0020     addParticleSelection,
0021 )
0022 from acts.examples.reconstruction import (
0023     addSeeding,
0024     TruthSeedRanges,
0025     CkfConfig,
0026     addCKFTracks,
0027     TrackSelectorConfig,
0028     addAmbiguityResolution,
0029     AmbiguityResolutionConfig,
0030     addAmbiguityResolutionML,
0031     AmbiguityResolutionMLConfig,
0032     addVertexFitting,
0033     VertexFinder,
0034     addSeedFilterML,
0035     SeedFilterMLDBScanConfig,
0036 )
0037 from acts.examples.odd import getOpenDataDetector, getOpenDataDetectorDirectory
0038
0039 u = acts.UnitConstants
0040
0041
0042 parser = argparse.ArgumentParser(description="Full chain with the OpenDataDetector")
0043 parser.add_argument(
0044     "--output",
0045     "-o",
0046     help="Output directory",
0047     type=pathlib.Path,
0048     default=pathlib.Path.cwd() / "odd_output",
0049 )
0050 parser.add_argument("--events", "-n", help="Number of events", type=int, default=100)
0051 parser.add_argument("--skip", "-s", help="Number of events", type=int, default=0)
0052 parser.add_argument("--edm4hep", help="Use edm4hep inputs", type=pathlib.Path)
0053 parser.add_argument(
0054     "--geant4", help="Use Geant4 instead of fatras", action="store_true"
0055 )
0056 parser.add_argument(
0057     "--ttbar",
0058     help="Use Pythia8 (ttbar, pile-up 200) instead of particle gun",
0059     action="store_true",
0060 )
0061 parser.add_argument(
0062     "--ttbar-pu",
0063     help="Number of pile-up events for ttbar",
0064     type=int,
0065     default=200,
0066 )
0067 parser.add_argument(
0068     "--gun-multiplicity",
0069     help="Multiplicity of the particle gun",
0070     type=int,
0071     default=200,
0072 )
0073 parser.add_argument(
0074     "--MLSolver",
0075     help="Use the Ml Ambiguity Solver instead of the classical one",
0076     action="store_true",
0077 )
0078 parser.add_argument(
0079     "--MLSeedFilter",
0080     help="Use the Ml seed filter to select seed after the seeding step",
0081     action="store_true",
0082 )
0083
0084 args = vars(parser.parse_args())
0085
0086 outputDir = args["output"]
0087 ttbar = args["ttbar"]
0088 g4_simulation = args["geant4"]
0089 ambiguity_MLSolver = args["MLSolver"]
0090 seedFilter_ML = args["MLSeedFilter"]
0091 geoDir = getOpenDataDetectorDirectory()
0092 # acts.examples.dump_args_calls(locals())  # show python binding calls
0093
0094 oddMaterialMap = geoDir / "data/odd-material-maps.root"
0095 oddDigiConfig = geoDir / "config/odd-digi-smearing-config.json"
0096 oddSeedingSel = geoDir / "config/odd-seeding-config.json"
0097 oddMaterialDeco = acts.IMaterialDecorator.fromFile(oddMaterialMap)
0098
0099 detector, trackingGeometry, decorators = getOpenDataDetector(
0100     odd_dir=geoDir, mdecorator=oddMaterialDeco
0101 )
0102 field = acts.ConstantBField(acts.Vector3(0.0, 0.0, 2.0 * u.T))
0103 rnd = acts.examples.RandomNumbers(seed=42)
0104
0105 s = acts.examples.Sequencer(
0106     events=args["events"],
0107     skip=args["skip"],
0108     numThreads=1 if g4_simulation else -1,
0109     outputDir=str(outputDir),
0110 )
0111
0112 if args["edm4hep"]:
0113     import acts.examples.edm4hep
0114
0115     edm4hepReader = acts.examples.edm4hep.EDM4hepReader(
0116         inputPath=str(args["edm4hep"]),
0117         inputSimHits=[
0118             "PixelBarrelReadout",
0119             "PixelEndcapReadout",
0120             "ShortStripBarrelReadout",
0121             "ShortStripEndcapReadout",
0122             "LongStripBarrelReadout",
0123             "LongStripEndcapReadout",
0124         ],
0125         outputParticlesGenerator="particles_input",
0126         outputParticlesInitial="particles_initial",
0127         outputParticlesFinal="particles_final",
0128         outputSimHits="simhits",
0129         graphvizOutput="graphviz",
0130         dd4hepDetector=detector,
0131         trackingGeometry=trackingGeometry,
0132         sortSimHitsInTime=True,
0133         level=acts.logging.INFO,
0134     )
0135     s.addReader(edm4hepReader)
0136     s.addWhiteboardAlias("particles", edm4hepReader.config.outputParticlesGenerator)
0137
0138     addParticleSelection(
0139         s,
0140         config=ParticleSelectorConfig(
0141             rho=(0.0, 24 * u.mm),
0142             absZ=(0.0, 1.0 * u.m),
0143             eta=(-3.0, 3.0),
0144             pt=(150 * u.MeV, None),
0145             removeNeutral=True,
0146         ),
0147         inputParticles="particles",
0148         outputParticles="particles_selected",
0149     )
0150 else:
0151     if not ttbar:
0152         addParticleGun(
0153             s,
0154             MomentumConfig(1.0 * u.GeV, 10.0 * u.GeV, transverse=True),
0155             EtaConfig(-3.0, 3.0),
0156             PhiConfig(0.0, 360.0 * u.degree),
0157             ParticleConfig(4, acts.PdgParticle.eMuon, randomizeCharge=True),
0158             vtxGen=acts.examples.GaussianVertexGenerator(
0159                 mean=acts.Vector4(0, 0, 0, 0),
0160                 stddev=acts.Vector4(
0161                     0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 1.0 * u.ns
0162                 ),
0163             ),
0164             multiplicity=args["gun_multiplicity"],
0165             rnd=rnd,
0166         )
0167     else:
0168         addPythia8(
0169             s,
0170             hardProcess=["Top:qqbar2ttbar=on"],
0171             npileup=args["ttbar_pu"],
0172             vtxGen=acts.examples.GaussianVertexGenerator(
0173                 mean=acts.Vector4(0, 0, 0, 0),
0174                 stddev=acts.Vector4(
0175                     0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 5.0 * u.ns
0176                 ),
0177             ),
0178             rnd=rnd,
0179             outputDirRoot=outputDir,
0180             # outputDirCsv=outputDir,
0181         )
0182
0183     if g4_simulation:
0184         if s.config.numThreads != 1:
0185             raise ValueError("Geant 4 simulation does not support multi-threading")
0186
0187         # Pythia can sometime simulate particles outside the world volume, a cut on the Z of the track help mitigate this effect
0188         # Older version of G4 might not work, this as has been tested on version `geant4-11-00-patch-03`
0189         # For more detail see issue #1578
0190         addGeant4(
0191             s,
0192             detector,
0193             trackingGeometry,
0194             field,
0195             preSelectParticles=ParticleSelectorConfig(
0196                 rho=(0.0, 24 * u.mm),
0197                 absZ=(0.0, 1.0 * u.m),
0198                 eta=(-3.0, 3.0),
0199                 pt=(150 * u.MeV, None),
0200                 removeNeutral=True,
0201             ),
0202             outputDirRoot=outputDir,
0203             # outputDirCsv=outputDir,
0204             rnd=rnd,
0205             killVolume=trackingGeometry.worldVolume,
0206             killAfterTime=25 * u.ns,
0207         )
0208     else:
0209         addFatras(
0210             s,
0211             trackingGeometry,
0212             field,
0213             preSelectParticles=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             if ttbar
0221             else ParticleSelectorConfig(),
0222             enableInteractions=True,
0223             outputDirRoot=outputDir,
0224             # outputDirCsv=outputDir,
0225             rnd=rnd,
0226         )
0227
0228 addDigitization(
0229     s,
0230     trackingGeometry,
0231     field,
0232     digiConfigFile=oddDigiConfig,
0233     outputDirRoot=outputDir,
0234     # outputDirCsv=outputDir,
0235     rnd=rnd,
0236 )
0237
0238 addSeeding(
0239     s,
0240     trackingGeometry,
0241     field,
0242     TruthSeedRanges(pt=(1.0 * u.GeV, None), eta=(-3.0, 3.0), nHits=(9, None))
0243     if ttbar
0244     else TruthSeedRanges(),
0245     geoSelectionConfigFile=oddSeedingSel,
0246     outputDirRoot=outputDir,
0247     # outputDirCsv=outputDir,
0248 )
0249
0250 if seedFilter_ML:
0251     addSeedFilterML(
0252         s,
0253         SeedFilterMLDBScanConfig(
0254             epsilonDBScan=0.03, minPointsDBScan=2, minSeedScore=0.1
0255         ),
0256         onnxModelFile=os.path.dirname(__file__)
0257         + "/MLAmbiguityResolution/seedDuplicateClassifier.onnx",
0258         outputDirRoot=outputDir,
0259         # outputDirCsv=outputDir,
0260     )
0261
0262 addCKFTracks(
0263     s,
0264     trackingGeometry,
0265     field,
0266     TrackSelectorConfig(
0267         pt=(1.0 * u.GeV if ttbar else 0.0, None),
0268         absEta=(None, 3.0),
0269         loc0=(-4.0 * u.mm, 4.0 * u.mm),
0270         nMeasurementsMin=7,
0271         maxHoles=2,
0272         maxOutliers=2,
0273     ),
0274     CkfConfig(
0275         seedDeduplication=True,
0276         stayOnSeed=True,
0277     ),
0278     outputDirRoot=outputDir,
0279     writeCovMat=True,
0280     # outputDirCsv=outputDir,
0281 )
0282
0283 if ambiguity_MLSolver:
0284     addAmbiguityResolutionML(
0285         s,
0286         AmbiguityResolutionMLConfig(
0287             maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0288         ),
0289         outputDirRoot=outputDir,
0290         # outputDirCsv=outputDir,
0291         onnxModelFile=os.path.dirname(__file__)
0292         + "/MLAmbiguityResolution/duplicateClassifier.onnx",
0293     )
0294 else:
0295     addAmbiguityResolution(
0296         s,
0297         AmbiguityResolutionConfig(
0298             maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0299         ),
0300         outputDirRoot=outputDir,
0301         writeCovMat=True,
0302         # outputDirCsv=outputDir,
0303     )
0304
0305 addVertexFitting(
0306     s,
0307     field,
0308     vertexFinder=VertexFinder.Iterative,
0309     outputDirRoot=outputDir,
0310 )
0311
0312 s.run()