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     addFatras,
0019     addGeant4,
0020     ParticleSelectorConfig,
0021     addDigitization,
0022     addParticleSelection,
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 # acts.examples.dump_args_calls(locals())  # show python binding calls
0155
0156 oddMaterialMap = (
0157     args.material_config
0158     if args.material_config
0159     else geoDir / "data/odd-material-maps.root"
0160 )
0161
0162 oddDigiConfig = (
0163     args.digi_config
0164     if args.digi_config
0165     else geoDir / "config/odd-digi-smearing-config.json"
0166 )
0167
0168 oddSeedingSel = geoDir / "config/odd-seeding-config.json"
0169 oddMaterialDeco = acts.IMaterialDecorator.fromFile(oddMaterialMap)
0170
0171 detector = getOpenDataDetector(odd_dir=geoDir, mdecorator=oddMaterialDeco)
0172 trackingGeometry = detector.trackingGeometry()
0173 decorators = detector.contextDecorators()
0174 field = acts.ConstantBField(acts.Vector3(0.0, 0.0, 2.0 * u.T))
0175 rnd = acts.examples.RandomNumbers(seed=42)
0176
0177 s = acts.examples.Sequencer(
0178     events=args.events,
0179     skip=args.skip,
0180     numThreads=1 if args.geant4 else -1,
0181     outputDir=str(outputDir),
0182 )
0183
0184 if args.edm4hep:
0185     import acts.examples.edm4hep
0186
0187     edm4hepReader = acts.examples.edm4hep.EDM4hepReader(
0188         inputPath=str(args.edm4hep),
0189         inputSimHits=[
0190             "PixelBarrelReadout",
0191             "PixelEndcapReadout",
0192             "ShortStripBarrelReadout",
0193             "ShortStripEndcapReadout",
0194             "LongStripBarrelReadout",
0195             "LongStripEndcapReadout",
0196         ],
0197         outputParticlesGenerator="particles_input",
0198         outputParticlesSimulation="particles_simulated",
0199         outputSimHits="simhits",
0200         graphvizOutput="graphviz",
0201         dd4hepDetector=detector,
0202         trackingGeometry=trackingGeometry,
0203         sortSimHitsInTime=True,
0204         level=acts.logging.INFO,
0205     )
0206     s.addReader(edm4hepReader)
0207     s.addWhiteboardAlias("particles", edm4hepReader.config.outputParticlesGenerator)
0208
0209     addParticleSelection(
0210         s,
0211         config=ParticleSelectorConfig(
0212             rho=(0.0, 24 * u.mm),
0213             absZ=(0.0, 1.0 * u.m),
0214             eta=(-3.0, 3.0),
0215             pt=(150 * u.MeV, None),
0216             removeNeutral=True,
0217         ),
0218         inputParticles="particles",
0219         outputParticles="particles_selected",
0220     )
0221 else:
0222     if not args.ttbar:
0223         addParticleGun(
0224             s,
0225             MomentumConfig(
0226                 args.gun_pt_range[0] * u.GeV,
0227                 args.gun_pt_range[1] * u.GeV,
0228                 transverse=True,
0229             ),
0230             EtaConfig(args.gun_eta_range[0], args.gun_eta_range[1]),
0231             PhiConfig(0.0, 360.0 * u.degree),
0232             ParticleConfig(
0233                 args.gun_particles, acts.PdgParticle.eMuon, randomizeCharge=True
0234             ),
0235             vtxGen=acts.examples.GaussianDisplacedVertexPositionGenerator(
0236                 rMean=2,
0237                 rStdDev=0.0125 * u.mm,
0238                 zMean=2,
0239                 zStdDev=55.5 * u.mm,
0240                 tMean=0,
0241                 tStdDev=1.0 * u.ns,
0242             ),
0243             multiplicity=args.gun_multiplicity,
0244             rnd=rnd,
0245         )
0246     else:
0247         addPythia8(
0248             s,
0249             hardProcess=["Top:qqbar2ttbar=on"],
0250             npileup=args.ttbar_pu,
0251             vtxGen=acts.examples.GaussianDisplacedVertexPositionGenerator(
0252                 mean=acts.Vector4(0, 0, 0, 0),
0253                 stddev=acts.Vector4(
0254                     0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 5.0 * u.ns
0255                 ),
0256             ),
0257             rnd=rnd,
0258             outputDirRoot=outputDir if args.output_root else None,
0259             outputDirCsv=outputDir if args.output_csv else None,
0260         )
0261
0262     if args.geant4:
0263         if s.config.numThreads != 1:
0264             raise ValueError("Geant 4 simulation does not support multi-threading")
0265
0266         # Pythia can sometime simulate particles outside the world volume, a cut on the Z of the track help mitigate this effect
0267         # Older version of G4 might not work, this as has been tested on version `geant4-11-00-patch-03`
0268         # For more detail see issue #1578
0269         addGeant4(
0270             s,
0271             detector,
0272             trackingGeometry,
0273             field,
0274             preSelectParticles=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             postSelectParticles=ParticleSelectorConfig(
0281                 pt=(1.0 * u.GeV, None),
0282                 eta=(-3.0, 3.0),
0283                 hits=(9, None),
0284                 removeNeutral=True,
0285             ),
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             preSelectParticles=(
0299                 ParticleSelectorConfig(
0300                     rho=(0.0, 24 * u.mm),
0301                     absZ=(0.0, 1.0 * u.m),
0302                     eta=(-3.0, 3.0),
0303                     pt=(150 * u.MeV, None),
0304                 )
0305                 if args.ttbar
0306                 else ParticleSelectorConfig()
0307             ),
0308             postSelectParticles=ParticleSelectorConfig(
0309                 pt=(1.0 * u.GeV, None),
0310                 eta=(-3.0, 3.0),
0311                 hits=(9, None),
0312                 removeNeutral=True,
0313             ),
0314             enableInteractions=True,
0315             outputDirRoot=outputDir if args.output_root else None,
0316             outputDirCsv=outputDir if args.output_csv else None,
0317             outputDirObj=outputDir if args.output_obj else None,
0318             rnd=rnd,
0319         )
0320
0321 addDigitization(
0322     s,
0323     trackingGeometry,
0324     field,
0325     digiConfigFile=oddDigiConfig,
0326     outputDirRoot=outputDir if args.output_root else None,
0327     outputDirCsv=outputDir if args.output_csv else None,
0328     rnd=rnd,
0329 )
0330
0331 if args.reco:
0332     addSeeding(
0333         s,
0334         trackingGeometry,
0335         field,
0336         initialSigmas=[
0337             1 * u.mm,
0338             1 * u.mm,
0339             1 * u.degree,
0340             1 * u.degree,
0341             0.1 * u.e / u.GeV,
0342             1 * u.ns,
0343         ],
0344         initialSigmaPtRel=0.1,
0345         initialVarInflation=[1.0] * 6,
0346         geoSelectionConfigFile=oddSeedingSel,
0347         outputDirRoot=outputDir if args.output_root else None,
0348         outputDirCsv=outputDir if args.output_csv else None,
0349     )
0350
0351     if seedFilter_ML:
0352         addSeedFilterML(
0353             s,
0354             SeedFilterMLDBScanConfig(
0355                 epsilonDBScan=0.03, minPointsDBScan=2, minSeedScore=0.1
0356             ),
0357             onnxModelFile=os.path.dirname(__file__)
0358             + "/MLAmbiguityResolution/seedDuplicateClassifier.onnx",
0359             outputDirRoot=outputDir if args.output_root else None,
0360             outputDirCsv=outputDir if args.output_csv else None,
0361         )
0362
0363     addCKFTracks(
0364         s,
0365         trackingGeometry,
0366         field,
0367         TrackSelectorConfig(
0368             pt=(1.0 * u.GeV if args.ttbar else 0.0, None),
0369             absEta=(None, 3.0),
0370             loc0=(-4.0 * u.mm, 4.0 * u.mm),
0371             nMeasurementsMin=7,
0372             maxHoles=2,
0373             maxOutliers=2,
0374         ),
0375         CkfConfig(
0376             chi2CutOffMeasurement=15.0,
0377             chi2CutOffOutlier=25.0,
0378             numMeasurementsCutOff=10,
0379             seedDeduplication=True,
0380             stayOnSeed=True,
0381             pixelVolumes=[16, 17, 18],
0382             stripVolumes=[23, 24, 25],
0383             maxPixelHoles=1,
0384             maxStripHoles=2,
0385             constrainToVolumes=[
0386                 2,  # beam pipe
0387                 32,
0388                 4,  # beam pip gap
0389                 16,
0390                 17,
0391                 18,  # pixel
0392                 20,  # PST
0393                 23,
0394                 24,
0395                 25,  # short strip
0396                 26,
0397                 8,  # long strip gap
0398                 28,
0399                 29,
0400                 30,  # long strip
0401             ],
0402         ),
0403         outputDirRoot=outputDir if args.output_root else None,
0404         outputDirCsv=outputDir if args.output_csv else None,
0405         writeCovMat=True,
0406     )
0407
0408     if ambi_ML:
0409         addAmbiguityResolutionML(
0410             s,
0411             AmbiguityResolutionMLConfig(
0412                 maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0413             ),
0414             outputDirRoot=outputDir if args.output_root else None,
0415             outputDirCsv=outputDir if args.output_csv else None,
0416             onnxModelFile=os.path.dirname(__file__)
0417             + "/MLAmbiguityResolution/duplicateClassifier.onnx",
0418         )
0419
0420     elif ambi_scoring:
0421         addScoreBasedAmbiguityResolution(
0422             s,
0423             ScoreBasedAmbiguityResolutionConfig(
0424                 minScore=0,
0425                 minScoreSharedTracks=1,
0426                 maxShared=2,
0427                 maxSharedTracksPerMeasurement=2,
0428                 pTMax=1400,
0429                 pTMin=0.5,
0430                 phiMax=math.pi,
0431                 phiMin=-math.pi,
0432                 etaMax=4,
0433                 etaMin=-4,
0434                 useAmbiguityFunction=False,
0435             ),
0436             outputDirRoot=outputDir if args.output_root else None,
0437             outputDirCsv=outputDir if args.output_csv else None,
0438             ambiVolumeFile=ambi_config,
0439             writeCovMat=True,
0440         )
0441     else:
0442         addAmbiguityResolution(
0443             s,
0444             AmbiguityResolutionConfig(
0445                 maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0446             ),
0447             outputDirRoot=outputDir if args.output_root else None,
0448             outputDirCsv=outputDir if args.output_csv else None,
0449             writeCovMat=True,
0450         )
0451
0452     addVertexFitting(
0453         s,
0454         field,
0455         vertexFinder=VertexFinder.AMVF,
0456         outputDirRoot=outputDir if args.output_root else None,
0457     )
0458
0459 s.run()