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