Scripts/Python/full_chain_odd_LRT_second_pass.py

0001 #!/usr/bin/env python3
0002 #
0003 # Two-pass tracking demonstration.
0004 #
0005 # Pass 1: standard prompt tracking (same as full_chain_odd_LRT.py).
0006 # Pass 2: LRT (large-radius tracking) on the measurements NOT used by pass 1.
0007 #         MeasurementFilterAlgorithm removes pass-1-consumed measurements from
0008 #         the initial full subset produced by digitization, and the LRT seeding
0009 #         + CKF runs on the remainder.
0010
0011 import os
0012 import argparse
0013 import pathlib
0014
0015 import acts
0016 import acts.examples
0017 from acts.examples.simulation import (
0018     addParticleGun,
0019     MomentumConfig,
0020     EtaConfig,
0021     PhiConfig,
0022     ParticleConfig,
0023     ParticleSelectorConfig,
0024     addPythia8,
0025     addGenParticleSelection,
0026     addFatras,
0027     addGeant4,
0028     addSimParticleSelection,
0029     addDigitization,
0030     addDigiParticleSelection,
0031 )
0032 from acts.examples.reconstruction import (
0033     addSeeding,
0034     SeedFinderConfigArg,
0035     SeedFinderOptionsArg,
0036     SeedFilterConfigArg,
0037     SpacePointGridConfigArg,
0038     SeedingAlgorithmConfigArg,
0039     CkfConfig,
0040     addCKFTracks,
0041     TrackSelectorConfig,
0042     addAmbiguityResolution,
0043     AmbiguityResolutionConfig,
0044 )
0045 from acts.examples.odd import getOpenDataDetector, getOpenDataDetectorDirectory
0046
0047 u = acts.UnitConstants
0048
0049
0050 parser = argparse.ArgumentParser(
0051     description="Two-pass tracking with UsedMeasurementMap (prompt + LRT)"
0052 )
0053 parser.add_argument(
0054     "--output",
0055     "-o",
0056     help="Output directory",
0057     type=pathlib.Path,
0058     default=pathlib.Path.cwd() / "odd_output_lrt",
0059 )
0060 parser.add_argument("--events", "-n", help="Number of events", type=int, default=1000)
0061 parser.add_argument(
0062     "--skip", "-s", help="Number of events to skip", type=int, default=0
0063 )
0064 parser.add_argument(
0065     "--geant4", help="Use Geant4 instead of fatras", action="store_true"
0066 )
0067 parser.add_argument(
0068     "--ttbar",
0069     help="Use Pythia8 (ttbar, pile-up 200) instead of particle gun",
0070     action="store_true",
0071 )
0072 parser.add_argument(
0073     "--ttbar-pu",
0074     help="Number of pile-up events for ttbar",
0075     type=int,
0076     default=200,
0077 )
0078 parser.add_argument(
0079     "--gun-particles",
0080     help="Multiplicity (no. of particles) of the particle gun",
0081     type=int,
0082     default=5,
0083 )
0084 parser.add_argument(
0085     "--gun-multiplicity",
0086     help="Multiplicity (no. of vertices) of the particle gun",
0087     type=int,
0088     default=200,
0089 )
0090 parser.add_argument(
0091     "--gun-eta-range",
0092     nargs=2,
0093     help="Eta range of the particle gun",
0094     type=float,
0095     default=[-3.0, 3.0],
0096 )
0097 parser.add_argument(
0098     "--gun-pt-range",
0099     nargs=2,
0100     help="Pt range of the particle gun (GeV)",
0101     type=float,
0102     default=[1.0, 10.0],
0103 )
0104 parser.add_argument(
0105     "--digi-config", help="Digitization configuration file", type=pathlib.Path
0106 )
0107 parser.add_argument(
0108     "--material-config", help="Material map configuration file", type=pathlib.Path
0109 )
0110 parser.add_argument(
0111     "--reco",
0112     help="Switch reco on/off",
0113     default=True,
0114     action=argparse.BooleanOptionalAction,
0115 )
0116 parser.add_argument(
0117     "--output-root",
0118     help="Switch root output on/off",
0119     default=True,
0120     action=argparse.BooleanOptionalAction,
0121 )
0122 parser.add_argument(
0123     "--output-csv",
0124     help="Switch csv output on/off",
0125     default=True,
0126     action=argparse.BooleanOptionalAction,
0127 )
0128
0129 args = parser.parse_args()
0130
0131 outputDir = args.output
0132 geoDir = getOpenDataDetectorDirectory()
0133 actsDir = pathlib.Path(__file__).parent.parent.parent.parent
0134
0135 oddMaterialMap = (
0136     args.material_config
0137     if args.material_config
0138     else geoDir / "data/odd-material-maps.root"
0139 )
0140 oddDigiConfig = (
0141     args.digi_config
0142     if args.digi_config
0143     else actsDir / "Examples/Configs/odd-digi-smearing-config.json"
0144 )
0145 oddSeedingSel = actsDir / "Examples/Configs/odd-seeding-config.json"
0146 oddStripSeedingSel = actsDir / "Examples/Configs/odd-strip-spacepoint-selection.json"
0147 oddShortStripSeedingSel = (
0148     actsDir / "Examples/Configs/odd-short-strip-spacepoint-selection.json"
0149 )
0150 oddMaterialDeco = acts.IMaterialDecorator.fromFile(oddMaterialMap)
0151
0152 detector = getOpenDataDetector(odd_dir=geoDir, materialDecorator=oddMaterialDeco)
0153 trackingGeometry = detector.trackingGeometry()
0154 decorators = detector.contextDecorators()
0155 field = acts.ConstantBField(acts.Vector3(0.0, 0.0, 2.0 * u.T))
0156 rnd = acts.examples.RandomNumbers(seed=42)
0157
0158 s = acts.examples.Sequencer(
0159     events=args.events,
0160     skip=args.skip,
0161     numThreads=1 if args.geant4 else -1,
0162     outputDir=str(outputDir),
0163 )
0164
0165 # -------------------------------------------------------------------------
0166 # Simulation
0167 # -------------------------------------------------------------------------
0168
0169 if not args.ttbar:
0170     addParticleGun(
0171         s,
0172         MomentumConfig(
0173             args.gun_pt_range[0] * u.GeV,
0174             args.gun_pt_range[1] * u.GeV,
0175             transverse=True,
0176         ),
0177         EtaConfig(args.gun_eta_range[0], args.gun_eta_range[1]),
0178         PhiConfig(0.0, 360.0 * u.degree),
0179         ParticleConfig(
0180             args.gun_particles, acts.PdgParticle.eMuon, randomizeCharge=True
0181         ),
0182         vtxGen=acts.examples.GaussianDisplacedVertexPositionGenerator(
0183             rMean=100.0,
0184             rStdDev=100.0 * u.mm,
0185             zMean=2,
0186             zStdDev=55.5 * u.mm,
0187             tMean=0,
0188             tStdDev=1.0 * u.ns,
0189         ),
0190         multiplicity=args.gun_multiplicity,
0191         rnd=rnd,
0192     )
0193 else:
0194     addPythia8(
0195         s,
0196         hardProcess=["Top:qqbar2ttbar=on"],
0197         npileup=args.ttbar_pu,
0198         vtxGen=acts.examples.GaussianDisplacedVertexPositionGenerator(
0199             rMean=50.0,
0200             rStdDev=50.0 * u.mm,
0201             zMean=2,
0202             zStdDev=55.5 * u.mm,
0203             tMean=0,
0204             tStdDev=5.0 * u.ns,
0205         ),
0206         rnd=rnd,
0207         outputDirRoot=outputDir if args.output_root else None,
0208         outputDirCsv=outputDir if args.output_csv else None,
0209     )
0210     addGenParticleSelection(
0211         s,
0212         ParticleSelectorConfig(
0213             rho=(0.0, 24 * u.mm),
0214             absZ=(0.0, 1.0 * u.m),
0215             eta=(-3.0, 3.0),
0216             pt=(150 * u.MeV, None),
0217         ),
0218     )
0219
0220 if args.geant4:
0221     if s.config.numThreads != 1:
0222         raise ValueError("Geant4 simulation does not support multi-threading")
0223     addGeant4(
0224         s,
0225         detector,
0226         trackingGeometry,
0227         field,
0228         outputDirRoot=outputDir if args.output_root else None,
0229         outputDirCsv=outputDir if args.output_csv else None,
0230         rnd=rnd,
0231         killVolume=trackingGeometry.highestTrackingVolume,
0232         killAfterTime=25 * u.ns,
0233     )
0234 else:
0235     addFatras(
0236         s,
0237         trackingGeometry,
0238         field,
0239         enableInteractions=True,
0240         outputDirRoot=outputDir if args.output_root else None,
0241         outputDirCsv=outputDir if args.output_csv else None,
0242         rnd=rnd,
0243     )
0244
0245 addDigitization(
0246     s,
0247     trackingGeometry,
0248     field,
0249     digiConfigFile=oddDigiConfig,
0250     outputDirRoot=outputDir if args.output_root else None,
0251     outputDirCsv=outputDir if args.output_csv else None,
0252     rnd=rnd,
0253 )
0254
0255 addDigiParticleSelection(
0256     s,
0257     ParticleSelectorConfig(
0258         pt=(150 * u.MeV, None),
0259         eta=(-3.0, 3.0),
0260         measurements=(9, None),
0261         removeNeutral=True,
0262     ),
0263 )
0264
0265 if not args.reco:
0266     s.run()
0267     raise SystemExit(0)
0268
0269 # =========================================================================
0270 # Pass 1: standard prompt tracking
0271 # Writes: "measurements" (original, from digitization)
0272 #         "spacepoints", "seeds", "estimatedparameters", "estimatedseeds"
0273 #         "ckf_tracks", "tracks" alias -> "ckf_tracks"
0274 #         "ambi_tracks", "tracks" alias -> "ambi_tracks"
0275 # =========================================================================
0276
0277 addSeeding(
0278     s,
0279     trackingGeometry,
0280     field,
0281     initialSigmas=[
0282         1 * u.mm,
0283         1 * u.mm,
0284         1 * u.degree,
0285         1 * u.degree,
0286         0 * u.e / u.GeV,
0287         1 * u.ns,
0288     ],
0289     initialSigmaQoverPt=0.1 * u.e / u.GeV,
0290     initialSigmaPtRel=0.1,
0291     initialVarInflation=[1.0] * 6,
0292     particleHypothesis=acts.ParticleHypothesis.muon,
0293     geoSelectionConfigFile=oddSeedingSel,
0294     outputDirRoot=outputDir if args.output_root else None,
0295     outputDirCsv=outputDir if args.output_csv else None,
0296 )
0297
0298 addCKFTracks(
0299     s,
0300     trackingGeometry,
0301     field,
0302     TrackSelectorConfig(
0303         absEta=(None, 3.0),
0304         loc0=(-4.0 * u.mm, 4.0 * u.mm),
0305         nMeasurementsMin=7,
0306         maxHoles=2,
0307         maxOutliers=2,
0308     ),
0309     CkfConfig(
0310         chi2CutOffMeasurement=15.0,
0311         chi2CutOffOutlier=25.0,
0312         numMeasurementsCutOff=2,
0313         seedDeduplication=True,
0314         stayOnSeed=True,
0315         pixelVolumes=[16, 17, 18],
0316         stripVolumes=[23, 24, 25],
0317         maxPixelHoles=1,
0318         maxStripHoles=2,
0319         constrainToVolumes=[
0320             2,
0321             32,
0322             4,  # beam pipe
0323             16,
0324             17,
0325             18,  # pixel
0326             20,  # PST
0327             23,
0328             24,
0329             25,  # short strip
0330             26,
0331             8,
0332             28,
0333             29,
0334             30,  # long strip
0335         ],
0336     ),
0337     outputDirRoot=outputDir if args.output_root else None,
0338     outputDirCsv=outputDir if args.output_csv else None,
0339     writeCovMat=True,
0340 )
0341
0342 # After addCKFTracks, "tracks" alias -> "ckf_tracks"
0343 addAmbiguityResolution(
0344     s,
0345     AmbiguityResolutionConfig(
0346         maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0347     ),
0348     outputDirRoot=outputDir if args.output_root else None,
0349     outputDirCsv=outputDir if args.output_csv else None,
0350     writeCovMat=True,
0351 )
0352 # After addAmbiguityResolution, "tracks" alias -> "ambi_tracks"
0353
0354 # =========================================================================
0355 # Measurement filtering: remove pass-1-consumed measurements for the LRT pass.
0356 # Reads:  "tracks" (-> "ambi_tracks"), "measurement_subset" (all measurements,
0357 #         written by digitization)
0358 # Writes: "lrt_measurement_subset" (remaining measurements, original indices)
0359 # =========================================================================
0360
0361 s.addAlgorithm(
0362     acts.examples.MeasurementFilterAlgorithm(
0363         level=acts.logging.INFO,
0364         inputTracks="tracks",
0365         inputMeasurementSubset="measurement_subset",
0366         outputMeasurementSubset="lrt_measurement_subset",
0367     )
0368 )
0369
0370 # =========================================================================
0371 # Pass 2: LRT (large-radius tracking) on filtered measurements
0372 #
0373 # LRT uses wider impact-parameter and collision-region cuts to recover
0374 # displaced tracks missed by the prompt pass. Seeding uses short strips
0375 # (2D, geometrySelection) and long-strip stereo pairs (stripGeometrySelection).
0376 # =========================================================================
0377
0378 addSeeding(
0379     s,
0380     trackingGeometry,
0381     field,
0382     geoSelectionConfigFile=oddShortStripSeedingSel,
0383     stripGeoSelectionConfigFile=oddStripSeedingSel,
0384     initialSigmas=[
0385         1 * u.mm,
0386         1 * u.mm,
0387         1 * u.degree,
0388         1 * u.degree,
0389         0 * u.e / u.GeV,
0390         1 * u.ns,
0391     ],
0392     initialSigmaQoverPt=0.1 * u.e / u.GeV,
0393     initialSigmaPtRel=0.1,
0394     initialVarInflation=[1.0] * 6,
0395     particleHypothesis=acts.ParticleHypothesis.muon,
0396     seedFinderConfigArg=SeedFinderConfigArg(
0397         impactMax=300.0 * u.mm,
0398         collisionRegion=(-350.0 * u.mm, 350.0 * u.mm),
0399         r=(200.0 * u.mm, 1000.0 * u.mm),
0400         rMinMiddle=200.0 * u.mm,
0401         rMaxMiddle=700.0 * u.mm,
0402     ),
0403     seedFinderOptionsArg=SeedFinderOptionsArg(bFieldInZ=2.0 * u.T),
0404     seedFilterConfigArg=SeedFilterConfigArg(),
0405     spacePointGridConfigArg=SpacePointGridConfigArg(),
0406     seedingAlgorithmConfigArg=SeedingAlgorithmConfigArg(),
0407     prefix="lrt_",
0408     logLevel=acts.logging.INFO,
0409     outputDirRoot=outputDir if args.output_root else None,
0410     outputDirCsv=outputDir if args.output_csv else None,
0411 )
0412
0413 addCKFTracks(
0414     s,
0415     trackingGeometry,
0416     field,
0417     TrackSelectorConfig(
0418         absEta=(None, 3.0),
0419         loc0=(-300.0 * u.mm, 300.0 * u.mm),
0420         nMeasurementsMin=8,
0421         maxHoles=2,
0422         maxOutliers=2,
0423     ),
0424     CkfConfig(
0425         chi2CutOffMeasurement=15.0,
0426         chi2CutOffOutlier=25.0,
0427         numMeasurementsCutOff=2,
0428         seedDeduplication=True,
0429         stayOnSeed=True,
0430         pixelVolumes=[16, 17, 18],
0431         stripVolumes=[23, 24, 25],
0432         maxPixelHoles=1,
0433         maxStripHoles=1,
0434         constrainToVolumes=[
0435             2,
0436             32,
0437             4,
0438             16,
0439             17,
0440             18,
0441             20,
0442             23,
0443             24,
0444             25,
0445             26,
0446             8,
0447             28,
0448             29,
0449             30,
0450         ],
0451     ),
0452     prefix="lrt_",
0453     logLevel=acts.logging.INFO,
0454     outputDirRoot=outputDir if args.output_root else None,
0455     outputDirCsv=outputDir if args.output_csv else None,
0456     writeCovMat=True,
0457 )
0458
0459 addAmbiguityResolution(
0460     s,
0461     AmbiguityResolutionConfig(
0462         maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0463     ),
0464     prefix="lrt_",
0465     logLevel=acts.logging.INFO,
0466     outputDirRoot=outputDir if args.output_root else None,
0467     outputDirCsv=outputDir if args.output_csv else None,
0468     writeCovMat=True,
0469 )
0470
0471 s.run()