Scripts/Python/full_chain_odd.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.GaussianVertexGenerator(
0236                 mean=acts.Vector4(0, 0, 0, 0),
0237                 stddev=acts.Vector4(
0238                     0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 1.0 * u.ns
0239                 ),
0240             ),
0241             multiplicity=args.gun_multiplicity,
0242             rnd=rnd,
0243         )
0244     else:
0245         addPythia8(
0246             s,
0247             hardProcess=["Top:qqbar2ttbar=on"],
0248             npileup=args.ttbar_pu,
0249             vtxGen=acts.examples.GaussianVertexGenerator(
0250                 mean=acts.Vector4(0, 0, 0, 0),
0251                 stddev=acts.Vector4(
0252                     0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 5.0 * u.ns
0253                 ),
0254             ),
0255             rnd=rnd,
0256             outputDirRoot=outputDir if args.output_root else None,
0257             outputDirCsv=outputDir if args.output_csv else None,
0258         )
0259
0260     if args.geant4:
0261         if s.config.numThreads != 1:
0262             raise ValueError("Geant 4 simulation does not support multi-threading")
0263
0264         # Pythia can sometime simulate particles outside the world volume, a cut on the Z of the track help mitigate this effect
0265         # Older version of G4 might not work, this as has been tested on version `geant4-11-00-patch-03`
0266         # For more detail see issue #1578
0267         addGeant4(
0268             s,
0269             detector,
0270             trackingGeometry,
0271             field,
0272             preSelectParticles=ParticleSelectorConfig(
0273                 rho=(0.0, 24 * u.mm),
0274                 absZ=(0.0, 1.0 * u.m),
0275                 eta=(-3.0, 3.0),
0276                 pt=(150 * u.MeV, None),
0277             ),
0278             postSelectParticles=ParticleSelectorConfig(
0279                 pt=(1.0 * u.GeV, None),
0280                 eta=(-3.0, 3.0),
0281                 hits=(9, None),
0282                 removeNeutral=True,
0283             ),
0284             outputDirRoot=outputDir if args.output_root else None,
0285             outputDirCsv=outputDir if args.output_csv else None,
0286             outputDirObj=outputDir if args.output_obj else None,
0287             rnd=rnd,
0288             killVolume=trackingGeometry.highestTrackingVolume,
0289             killAfterTime=25 * u.ns,
0290         )
0291     else:
0292         addFatras(
0293             s,
0294             trackingGeometry,
0295             field,
0296             preSelectParticles=(
0297                 ParticleSelectorConfig(
0298                     rho=(0.0, 24 * u.mm),
0299                     absZ=(0.0, 1.0 * u.m),
0300                     eta=(-3.0, 3.0),
0301                     pt=(150 * u.MeV, None),
0302                 )
0303                 if args.ttbar
0304                 else ParticleSelectorConfig()
0305             ),
0306             postSelectParticles=ParticleSelectorConfig(
0307                 pt=(1.0 * u.GeV, None),
0308                 eta=(-3.0, 3.0),
0309                 hits=(9, None),
0310                 removeNeutral=True,
0311             ),
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 if args.reco:
0330     addSeeding(
0331         s,
0332         trackingGeometry,
0333         field,
0334         initialSigmas=[
0335             1 * u.mm,
0336             1 * u.mm,
0337             1 * u.degree,
0338             1 * u.degree,
0339             0.1 * u.e / u.GeV,
0340             1 * u.ns,
0341         ],
0342         initialSigmaPtRel=0.1,
0343         initialVarInflation=[1.0] * 6,
0344         geoSelectionConfigFile=oddSeedingSel,
0345         outputDirRoot=outputDir if args.output_root else None,
0346         outputDirCsv=outputDir if args.output_csv else None,
0347     )
0348
0349     if seedFilter_ML:
0350         addSeedFilterML(
0351             s,
0352             SeedFilterMLDBScanConfig(
0353                 epsilonDBScan=0.03, minPointsDBScan=2, minSeedScore=0.1
0354             ),
0355             onnxModelFile=os.path.dirname(__file__)
0356             + "/MLAmbiguityResolution/seedDuplicateClassifier.onnx",
0357             outputDirRoot=outputDir if args.output_root else None,
0358             outputDirCsv=outputDir if args.output_csv else None,
0359         )
0360
0361     addCKFTracks(
0362         s,
0363         trackingGeometry,
0364         field,
0365         TrackSelectorConfig(
0366             pt=(1.0 * u.GeV if args.ttbar else 0.0, None),
0367             absEta=(None, 3.0),
0368             loc0=(-4.0 * u.mm, 4.0 * u.mm),
0369             nMeasurementsMin=7,
0370             maxHoles=2,
0371             maxOutliers=2,
0372         ),
0373         CkfConfig(
0374             chi2CutOffMeasurement=15.0,
0375             chi2CutOffOutlier=25.0,
0376             numMeasurementsCutOff=10,
0377             seedDeduplication=True,
0378             stayOnSeed=True,
0379             pixelVolumes=[16, 17, 18],
0380             stripVolumes=[23, 24, 25],
0381             maxPixelHoles=1,
0382             maxStripHoles=2,
0383             constrainToVolumes=[
0384                 2,  # beam pipe
0385                 32,
0386                 4,  # beam pip gap
0387                 16,
0388                 17,
0389                 18,  # pixel
0390                 20,  # PST
0391                 23,
0392                 24,
0393                 25,  # short strip
0394                 26,
0395                 8,  # long strip gap
0396                 28,
0397                 29,
0398                 30,  # long strip
0399             ],
0400         ),
0401         outputDirRoot=outputDir if args.output_root else None,
0402         outputDirCsv=outputDir if args.output_csv else None,
0403         writeCovMat=True,
0404     )
0405
0406     if ambi_ML:
0407         addAmbiguityResolutionML(
0408             s,
0409             AmbiguityResolutionMLConfig(
0410                 maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0411             ),
0412             outputDirRoot=outputDir if args.output_root else None,
0413             outputDirCsv=outputDir if args.output_csv else None,
0414             onnxModelFile=os.path.dirname(__file__)
0415             + "/MLAmbiguityResolution/duplicateClassifier.onnx",
0416         )
0417
0418     elif ambi_scoring:
0419         addScoreBasedAmbiguityResolution(
0420             s,
0421             ScoreBasedAmbiguityResolutionConfig(
0422                 minScore=0,
0423                 minScoreSharedTracks=1,
0424                 maxShared=2,
0425                 maxSharedTracksPerMeasurement=2,
0426                 pTMax=1400,
0427                 pTMin=0.5,
0428                 phiMax=math.pi,
0429                 phiMin=-math.pi,
0430                 etaMax=4,
0431                 etaMin=-4,
0432                 useAmbiguityFunction=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()