Scripts/Python/full_chain_test.py

0001 #!/usr/bin/env python3
0002
0003 import sys, os, argparse, pathlib
0004
0005
0006 def parse_args():
0007     from acts.examples.reconstruction import SeedingAlgorithm
0008
0009     parser = argparse.ArgumentParser(description="""
0010 Script to test the full chain ACTS simulation and reconstruction.
0011
0012 This script is provided for interactive developer testing only.
0013 It is not intended (and not supported) for end user use, automated testing,
0014 and certainly should never be called in production. The Python API is the
0015 proper way to access the ActsExamples from scripts. The other Examples/Scripts
0016 are much better examples of how to do that. physmon in the CI is the proper
0017 way to do automated integration tests. This script is only for the case of
0018 interactive testing with one-off configuration specified by command-line options.
0019 """)
0020     parser.add_argument(
0021         "-G",
0022         "--generic-detector",
0023         action="store_true",
0024         help="Use generic detector geometry and config",
0025     )
0026     parser.add_argument(
0027         "--odd",
0028         default=True,
0029         action=argparse.BooleanOptionalAction,
0030         help="Use Open Data Detector geometry and config (default unless overridden by -G or -A). Requires ACTS_BUILD_ODD.",
0031     )
0032     parser.add_argument(
0033         "-A",
0034         "--itk",
0035         action="store_true",
0036         help="Use ATLAS ITk geometry and config. Requires acts-itk/ in current directory.",
0037     )
0038     parser.add_argument(
0039         "-g",
0040         "--geant4",
0041         action="store_true",
0042         help="Use Geant4 instead of Fatras for detector simulation",
0043     )
0044     parser.add_argument(
0045         "--edm4hep",
0046         type=pathlib.Path,
0047         help="Use edm4hep inputs",
0048     )
0049     parser.add_argument(
0050         "-b",
0051         "--bf-constant",
0052         action="store_true",
0053         help="Use constant 2T B-field also for ITk; and don't include material map",
0054     )
0055     parser.add_argument(
0056         "-j",
0057         "--threads",
0058         type=int,
0059         default=-1,
0060         help="Number of parallel threads, negative for automatic (default).",
0061     )
0062     parser.add_argument(
0063         "-o",
0064         "--output-dir",
0065         "--output",
0066         default=None,
0067         type=pathlib.Path,
0068         help="Directory to write outputs to",
0069     )
0070     parser.add_argument(
0071         "-O",
0072         "--output-detail",
0073         action="count",
0074         default=0,
0075         help="fewer output files. Use -OO for more output files. Use -OOO to disable all output.",
0076     )
0077     parser.add_argument(
0078         "-c",
0079         "--output-csv",
0080         action="count",
0081         default=0,
0082         help="Use CSV output instead of ROOT. Specify -cc to output all formats (ROOT, CSV, and obj).",
0083     )
0084     parser.add_argument(
0085         "--output-obj",
0086         action="store_true",
0087         help="Enable obj output",
0088     )
0089     parser.add_argument(
0090         "-n",
0091         "--events",
0092         type=int,
0093         default=100,
0094         help="The number of events to process (default=%(default)d).",
0095     )
0096     parser.add_argument(
0097         "-s",
0098         "--skip",
0099         type=int,
0100         default=0,
0101         help="Number of events to skip (default=%(default)d)",
0102     )
0103     # Many of the following option names were inherited from the old examples binaries and full_chain_odd.py.
0104     # To maintain compatibility, both option names are supported.
0105     parser.add_argument(
0106         "-N",
0107         "--gen-nparticles",
0108         "--gun-particles",
0109         type=int,
0110         default=4,
0111         help="Number of generated particles per vertex from the particle gun (default=%(default)d).",
0112     )
0113     parser.add_argument(
0114         "-M",
0115         "--gen-nvertices",
0116         "--gun-multiplicity",
0117         "--ttbar-pu",
0118         type=int,
0119         default=200,
0120         help="Number of vertices per event (multiplicity) from the particle gun; or number of pileup events (default=%(default)d)",
0121     )
0122     parser.add_argument(
0123         "-t",
0124         "--ttbar-pu200",
0125         "--ttbar",
0126         action="store_true",
0127         help="Generate ttbar + mu=200 pile-up using Pythia8",
0128     )
0129     parser.add_argument(
0130         "-p",
0131         "--gen-pt-range",
0132         "--gun-pt-range",
0133         default="1:10",
0134         help="transverse momentum (pT) range (min:max) of the particle gun in GeV (default=%(default)s)",
0135     )
0136     parser.add_argument(
0137         "--gen-eta-range",
0138         "--gun-eta-range",
0139         help="Eta range (min:max) of the particle gun (default -2:2 (Generic), -3:3 (ODD), -4:4 (ITk))",
0140     )
0141     parser.add_argument(
0142         "--gen-cos-theta",
0143         action="store_true",
0144         help="Sample eta as cos(theta) and not uniform",
0145     )
0146     parser.add_argument(
0147         "-r",
0148         "--random-seed",
0149         type=int,
0150         default=42,
0151         help="Random number seed (default=%(default)d)",
0152     )
0153     parser.add_argument(
0154         "-F",
0155         "--disable-fpemon",
0156         action="store_true",
0157         help="sets ACTS_SEQUENCER_DISABLE_FPEMON=1",
0158     )
0159     parser.add_argument(
0160         "-l",
0161         "--loglevel",
0162         type=int,
0163         default=2,
0164         help="The output log level. Please set the wished number (0 = VERBOSE, 1 = DEBUG, 2 = INFO (default), 3 = WARNING, 4 = ERROR, 5 = FATAL).",
0165     )
0166     parser.add_argument(
0167         "-d",
0168         "--dump-args-calls",
0169         action="store_true",
0170         help="Show pybind function call details",
0171     )
0172     parser.add_argument(
0173         "--digi-config",
0174         type=pathlib.Path,
0175         help="Digitization configuration file",
0176     )
0177     parser.add_argument(
0178         "--material-config",
0179         type=pathlib.Path,
0180         help="Material map configuration file",
0181     )
0182     parser.add_argument(
0183         "-S",
0184         "--seeding-algorithm",
0185         action=EnumAction,
0186         enum=SeedingAlgorithm,
0187         default=SeedingAlgorithm.GridTriplet,
0188         help="Select the seeding algorithm to use",
0189     )
0190     parser.add_argument(
0191         "--gbts-layer-map",
0192         type=pathlib.Path,
0193         help="GBTS layer map file (default: the detector's shipped one)",
0194     )
0195     parser.add_argument(
0196         "--gbts-connection-table",
0197         type=pathlib.Path,
0198         help="GBTS layer connection table file (default: the detector's shipped one)",
0199     )
0200     parser.add_argument(
0201         "--ckf",
0202         default=True,
0203         action=argparse.BooleanOptionalAction,
0204         help="Switch CKF on/off",
0205     )
0206     parser.add_argument(
0207         "--reco",
0208         default=True,
0209         action=argparse.BooleanOptionalAction,
0210         help="Switch reco on/off",
0211     )
0212     parser.add_argument(
0213         "--vertexing",
0214         default=True,
0215         action=argparse.BooleanOptionalAction,
0216         help="Switch vertexing on/off",
0217     )
0218     parser.add_argument(
0219         "--MLSeedFilter",
0220         action="store_true",
0221         help="Use the ML seed filter to select seed after the seeding step",
0222     )
0223     parser.add_argument(
0224         "--ambi-solver",
0225         type=str,
0226         choices=["greedy", "scoring", "ML", "none"],
0227         default="greedy",
0228         help="Set which ambiguity solver to use (default=%(default)s)",
0229     )
0230     parser.add_argument(
0231         "--ambi-config",
0232         type=pathlib.Path,
0233         default=pathlib.Path.cwd() / "ambi_config.json",
0234         help="Set the configuration file for the Score Based ambiguity resolution (default=%(default)s)",
0235     )
0236     return parser.parse_args()
0237
0238
0239 def full_chain(args):
0240     import acts, acts.root
0241
0242     # keep these in memory after we return the sequence
0243     global detector, trackingGeometry, decorators, field, rnd
0244     global logger
0245
0246     if args.disable_fpemon:
0247         os.environ["ACTS_SEQUENCER_DISABLE_FPEMON"] = "1"
0248
0249     if args.dump_args_calls:
0250         acts.examples.dump_args_calls()
0251
0252     logger = acts.getDefaultLogger("full_chain_test", acts.logging.Level(args.loglevel))
0253
0254     nDetArgs = [args.generic_detector, args.odd, args.itk].count(True)
0255     if nDetArgs == 0:
0256         args.generic_detector = True
0257     elif nDetArgs == 2:
0258         args.odd = False
0259     nDetArgs = [args.generic_detector, args.odd, args.itk].count(True)
0260     if nDetArgs != 1:
0261         logger.fatal("require exactly one of: --generic-detector --odd --itk")
0262         sys.exit(2)
0263     if args.generic_detector:
0264         detname = "gen"
0265     elif args.itk:
0266         detname = "itk"
0267     elif args.odd:
0268         detname = "odd"
0269
0270     u = acts.UnitConstants
0271
0272     if args.output_detail == 3:
0273         outputDirLess = None
0274     elif args.output_dir is None:
0275         outputDirLess = pathlib.Path.cwd() / f"{detname}_output"
0276     else:
0277         outputDirLess = args.output_dir
0278
0279     outputDir = None if args.output_detail == 1 else outputDirLess
0280     outputDirMore = None if args.output_detail in (0, 1) else outputDirLess
0281
0282     outputDirRoot = outputDir if args.output_csv != 1 else None
0283     outputDirLessRoot = outputDirLess if args.output_csv != 1 else None
0284     outputDirMoreRoot = outputDirMore if args.output_csv != 1 else None
0285     outputDirCsv = outputDir if args.output_csv != 0 else None
0286     outputDirLessCsv = outputDirLess if args.output_csv != 0 else None
0287     outputDirMoreCsv = outputDirMore if args.output_csv != 0 else None
0288     outputDirObj = (
0289         outputDirLess
0290         if args.output_obj
0291         else outputDir if args.output_csv == 2 else None
0292     )
0293
0294     actsDir = pathlib.Path(__file__).resolve().parent.parent.parent.parent
0295
0296     # fmt: off
0297     if args.generic_detector:
0298         etaRange = (-2.0, 2.0)
0299         ptMin = 0.5 * u.GeV
0300         rhoMax = 24.0 * u.mm
0301         if args.loglevel <= 2:
0302             logger.info(f"Load Generic Detector from {actsDir}")
0303         if args.digi_config is None:
0304             args.digi_config = actsDir / "Examples/Configs/generic-digi-smearing-config.json"
0305         seedingConfigFile = actsDir / "Examples/Configs/generic-seeding-config.json"
0306         if args.gbts_layer_map is None:
0307             args.gbts_layer_map = actsDir / "Examples/Configs/generic-gbts-layer-map.json"
0308         if args.gbts_connection_table is None:
0309             args.gbts_connection_table = actsDir / "Examples/Configs/generic-gbts-connection-table.json"
0310         args.bf_constant = True
0311         detector = acts.examples.GenericDetector()
0312         trackingGeometry = detector.trackingGeometry()
0313         decorators = detector.contextDecorators()
0314     elif args.odd:
0315         import acts.examples.odd
0316         etaRange = (-3.0, 3.0)
0317         ptMin = 1.0 * u.GeV
0318         rhoMax = 24.0 * u.mm
0319         geoDir = acts.examples.odd.getOpenDataDetectorDirectory()
0320         if args.loglevel <= 2:
0321             logger.info(f"Load Open Data Detector from {geoDir.resolve()}")
0322         if args.digi_config is None:
0323             args.digi_config = actsDir / "Examples/Configs/odd-digi-smearing-config.json"
0324         seedingConfigFile = actsDir / "Examples/Configs/odd-seeding-config.json"
0325         if args.gbts_layer_map is None:
0326             args.gbts_layer_map = actsDir / "Examples/Configs/odd-gbts-layer-map.json"
0327         if args.gbts_connection_table is None:
0328             args.gbts_connection_table = actsDir / "Examples/Configs/odd-gbts-connection-table.json"
0329         if args.material_config is None:
0330             args.material_config = geoDir / "data/odd-material-maps.root"
0331         args.bf_constant = True
0332         detector = acts.examples.odd.getOpenDataDetector(
0333             odd_dir=geoDir,
0334             materialDecorator=acts.IMaterialDecorator.fromFile(args.material_config),
0335         )
0336         trackingGeometry = detector.trackingGeometry()
0337         decorators = detector.contextDecorators()
0338     elif args.itk:
0339         import acts.examples.itk as itk
0340         etaRange = (-4.0, 4.0)
0341         ptMin = 1.0 * u.GeV
0342         rhoMax = 28.0 * u.mm
0343         geoDir = pathlib.Path("acts-itk")
0344         if args.loglevel <= 2:
0345             logger.info(f"Load ATLAS ITk from {geoDir.resolve()}")
0346         if args.digi_config is None:
0347             args.digi_config = geoDir / "itk-hgtd/itk-smearing-config.json"
0348         seedingConfigFile = geoDir / "itk-hgtd/geoSelection-ITk.json"
0349         if args.gbts_layer_map is None:
0350             args.gbts_layer_map = geoDir / "itk-hgtd/gbts_layers.json"
0351         if args.gbts_connection_table is None:
0352             args.gbts_connection_table = geoDir / "itk-hgtd/gbts_connections.json"
0353         # args.material_config defaulted in itk.buildITkGeometry: geoDir / "itk-hgtd/material-maps-ITk-HGTD.json"
0354         bFieldFile = geoDir / "bfield/ATLAS-BField-xyz.root"
0355         detector = itk.buildITkGeometry(
0356             geoDir,
0357             customMaterialFile=args.material_config,
0358             material=not args.bf_constant,
0359             logLevel=acts.logging.Level(args.loglevel),
0360         )
0361         trackingGeometry = detector.trackingGeometry()
0362         decorators = detector.contextDecorators()
0363     # fmt: on
0364
0365     if args.bf_constant:
0366         field = acts.ConstantBField(acts.Vector3(0.0, 0.0, 2.0 * u.T))
0367     else:
0368         logger.info("Create magnetic field map from {}", bFieldFile)
0369         field = acts.root.MagneticFieldMapXyz(str(bFieldFile))
0370     rnd = acts.examples.RandomNumbers(seed=args.random_seed)
0371
0372     from acts.examples.simulation import (
0373         MomentumConfig,
0374         EtaConfig,
0375         PhiConfig,
0376         ParticleConfig,
0377         ParticleSelectorConfig,
0378         addDigitization,
0379         addGenParticleSelection,
0380         addSimParticleSelection,
0381         addDigiParticleSelection,
0382     )
0383
0384     s = acts.examples.Sequencer(
0385         events=args.events,
0386         skip=args.skip,
0387         numThreads=args.threads if not (args.geant4 and args.threads == -1) else 1,
0388         logLevel=acts.logging.Level(args.loglevel),
0389         outputDir="" if outputDirLess is None else str(outputDirLess),
0390     )
0391
0392     # is this needed?
0393     for d in decorators:
0394         s.addContextDecorator(d)
0395
0396     if args.edm4hep:
0397         import acts.examples.edm4hep
0398
0399         s.addReader(
0400             acts.examples.edm4hep.PodioReader(
0401                 level=acts.logging.DEBUG,
0402                 inputPath=str(args.edm4hep),
0403                 outputFrame="events",
0404                 category="events",
0405             )
0406         )
0407
0408         edm4hepReader = acts.examples.edm4hep.EDM4hepSimInputConverter(
0409             inputFrame="events",
0410             inputSimHits=[
0411                 "PixelBarrelReadout",
0412                 "PixelEndcapReadout",
0413                 "ShortStripBarrelReadout",
0414                 "ShortStripEndcapReadout",
0415                 "LongStripBarrelReadout",
0416                 "LongStripEndcapReadout",
0417             ],
0418             outputParticlesGenerator="particles_generated",
0419             outputParticlesSimulation="particles_simulated",
0420             outputSimHits="simhits",
0421             outputSimVertices="vertices_truth",
0422             dd4hepDetector=detector,
0423             trackingGeometry=trackingGeometry,
0424             sortSimHitsInTime=False,
0425             particleRMax=1080 * u.mm,
0426             particleZ=(-3030 * u.mm, 3030 * u.mm),
0427             particlePtMin=150 * u.MeV,
0428             level=acts.logging.DEBUG,
0429         )
0430         s.addAlgorithm(edm4hepReader)
0431
0432         s.addWhiteboardAlias(
0433             "particles", edm4hepReader.config.outputParticlesSimulation
0434         )
0435
0436         addSimParticleSelection(
0437             s,
0438             ParticleSelectorConfig(
0439                 rho=(0.0 * u.mm, rhoMax),
0440                 absZ=(0.0 * u.mm, 1.0 * u.m),
0441                 eta=etaRange,
0442                 removeNeutral=True,
0443             ),
0444         )
0445
0446     else:
0447
0448         if not args.ttbar_pu200:
0449             from acts.examples.simulation import addParticleGun
0450
0451             addParticleGun(
0452                 s,
0453                 MomentumConfig(
0454                     *strToRange(args.gen_pt_range, "--gen-pt-range", u.GeV),
0455                     transverse=True,
0456                 ),
0457                 EtaConfig(
0458                     *(
0459                         strToRange(args.gen_eta_range, "--gen-eta-range")
0460                         if args.gen_eta_range
0461                         else etaRange
0462                     ),
0463                     uniform=(
0464                         not args.gen_cos_theta
0465                         if args.gen_cos_theta or not args.odd
0466                         else None
0467                     ),
0468                 ),
0469                 PhiConfig(0.0, 360.0 * u.degree) if not args.itk else PhiConfig(),
0470                 ParticleConfig(
0471                     args.gen_nparticles, acts.PdgParticle.eMuon, randomizeCharge=True
0472                 ),
0473                 vtxGen=(
0474                     acts.examples.GaussianVertexGenerator(
0475                         mean=acts.Vector4(0, 0, 0, 0),
0476                         stddev=acts.Vector4(
0477                             0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 1.0 * u.ns
0478                         ),
0479                     )
0480                     if args.odd
0481                     else None
0482                 ),
0483                 multiplicity=args.gen_nvertices,
0484                 rnd=rnd,
0485                 outputDirRoot=outputDirMoreRoot,
0486                 outputDirCsv=outputDirMoreCsv,
0487             )
0488         else:
0489             from acts.examples.simulation import addPythia8
0490
0491             addPythia8(
0492                 s,
0493                 hardProcess=["Top:qqbar2ttbar=on"],
0494                 npileup=args.gen_nvertices,
0495                 vtxGen=acts.examples.GaussianVertexGenerator(
0496                     stddev=acts.Vector4(
0497                         0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 5.0 * u.ns
0498                     ),
0499                     mean=acts.Vector4(0, 0, 0, 0),
0500                 ),
0501                 rnd=rnd,
0502                 outputDirRoot=outputDirRoot,
0503                 outputDirCsv=outputDirCsv,
0504             )
0505             addGenParticleSelection(
0506                 s,
0507                 ParticleSelectorConfig(
0508                     rho=(0.0 * u.mm, rhoMax),
0509                     absZ=(0.0 * u.mm, 1.0 * u.m),
0510                     eta=etaRange,
0511                     pt=(150 * u.MeV, None),
0512                 ),
0513             )
0514
0515         if not args.geant4:
0516             from acts.examples.simulation import addFatras
0517
0518             addFatras(
0519                 s,
0520                 trackingGeometry,
0521                 field,
0522                 rnd=rnd,
0523                 outputDirRoot=outputDirRoot,
0524                 outputDirCsv=outputDirCsv,
0525                 outputDirObj=outputDirObj,
0526             )
0527         else:
0528             if s.config.numThreads != 1:
0529                 logger.fatal(
0530                     f"Geant 4 simulation does not support multi-threading (threads={s.config.numThreads})"
0531                 )
0532                 sys.exit(2)
0533
0534             from acts.examples.simulation import addGeant4
0535
0536             # Pythia can sometime simulate particles outside the world volume, a cut on the Z of the track help mitigate this effect
0537             # Older version of G4 might not work, this as has been tested on version `geant4-11-00-patch-03`
0538             # For more detail see issue #1578
0539             addGeant4(
0540                 s,
0541                 detector,
0542                 trackingGeometry,
0543                 field,
0544                 rnd=rnd,
0545                 killVolume=trackingGeometry.highestTrackingVolume,
0546                 killAfterTime=25 * u.ns,
0547                 outputDirRoot=outputDirRoot,
0548                 outputDirCsv=outputDirCsv,
0549                 outputDirObj=outputDirObj,
0550             )
0551
0552     addDigitization(
0553         s,
0554         trackingGeometry,
0555         field,
0556         digiConfigFile=args.digi_config,
0557         rnd=rnd,
0558         outputDirRoot=outputDirRoot,
0559         outputDirCsv=outputDirCsv,
0560     )
0561
0562     addDigiParticleSelection(
0563         s,
0564         ParticleSelectorConfig(
0565             pt=(ptMin, None),
0566             eta=etaRange if not args.generic_detector else (None, None),
0567             measurements=(9, None),
0568             removeNeutral=True,
0569         ),
0570     )
0571
0572     if not args.reco:
0573         return s
0574
0575     from acts.examples.reconstruction import (
0576         addSeeding,
0577         TrackSmearingSigmas,
0578         addCKFTracks,
0579         CkfConfig,
0580         SeedingAlgorithm,
0581         TrackSelectorConfig,
0582         addAmbiguityResolution,
0583         AmbiguityResolutionConfig,
0584         addVertexFitting,
0585         VertexFinder,
0586     )
0587
0588     if args.seeding_algorithm == SeedingAlgorithm.Gbts and (
0589         pathlib.Path(args.gbts_layer_map).is_file() == False
0590         or pathlib.Path(args.gbts_connection_table).is_file() == False
0591     ):
0592         logger.fatal(
0593             "GBTS seeding needs --gbts-layer-map and --gbts-connection-table "
0594             "for this detector"
0595         )
0596         sys.exit(2)
0597
0598     if args.itk and args.seeding_algorithm in (
0599         SeedingAlgorithm.GridTriplet,
0600         SeedingAlgorithm.Gbts,
0601     ):
0602         seedingAlgConfig = itk.itkSeedingAlgConfig(
0603             itk.InputSpacePointsType.PixelSpacePoints
0604         )
0605     else:
0606         seedingAlgConfig = []
0607
0608     addSeeding(
0609         s,
0610         trackingGeometry,
0611         field,
0612         *seedingAlgConfig,
0613         seedingAlgorithm=args.seeding_algorithm,
0614         **(
0615             dict(
0616                 trackSmearingSigmas=TrackSmearingSigmas(ptRel=0.01),
0617                 rnd=rnd,
0618             )
0619             if args.seeding_algorithm == SeedingAlgorithm.TruthSmeared
0620             else (
0621                 dict(
0622                     layerMappingConfigFile=args.gbts_layer_map,
0623                     connectorInputConfigFile=args.gbts_connection_table,
0624                 )
0625                 if args.seeding_algorithm == SeedingAlgorithm.Gbts
0626                 else {}
0627             )
0628         ),
0629         initialSigmas=[
0630             1 * u.mm,
0631             1 * u.mm,
0632             1 * u.degree,
0633             1 * u.degree,
0634             0 * u.e / u.GeV,
0635             1 * u.ns,
0636         ],
0637         initialSigmaQoverPt=0.1 * u.e / u.GeV,
0638         initialSigmaPtRel=0.1,
0639         initialVarInflation=[1.0] * 6,
0640         geoSelectionConfigFile=seedingConfigFile,
0641         outputDirRoot=outputDirLessRoot,
0642         outputDirCsv=outputDirLessCsv,
0643     )
0644
0645     if args.MLSeedFilter:
0646         from acts.examples.reconstruction import (
0647             addSeedFilterML,
0648             SeedFilterMLDBScanConfig,
0649         )
0650
0651         addSeedFilterML(
0652             s,
0653             SeedFilterMLDBScanConfig(
0654                 epsilonDBScan=0.03, minPointsDBScan=2, minSeedScore=0.1
0655             ),
0656             onnxModelFile=str(
0657                 actsDir
0658                 / "Examples/Scripts/Python/MLAmbiguityResolution/seedDuplicateClassifier.onnx"
0659             ),
0660             outputDirRoot=outputDirLessRoot,
0661             outputDirCsv=outputDirLessCsv,
0662         )
0663
0664     if not args.ckf:
0665         return s
0666
0667     if args.seeding_algorithm != SeedingAlgorithm.TruthSmeared:
0668         ckfConfig = CkfConfig(
0669             seedDeduplication=True,
0670             stayOnSeed=True,
0671         )
0672     else:
0673         ckfConfig = CkfConfig()
0674
0675     if not args.itk:
0676         trackSelectorConfig = TrackSelectorConfig(
0677             pt=(ptMin if args.ttbar_pu200 else 0.0, None),
0678             absEta=(None, 3.0),
0679             loc0=(-4.0 * u.mm, 4.0 * u.mm),
0680             nMeasurementsMin=7,
0681             maxHoles=2,
0682             maxOutliers=2,
0683         )
0684         ckfConfig = ckfConfig._replace(
0685             chi2CutOffMeasurement=15.0,
0686             chi2CutOffOutlier=25.0,
0687             numMeasurementsCutOff=2,
0688         )
0689     else:
0690         # fmt: off
0691         trackSelectorConfig = (
0692             TrackSelectorConfig(absEta=(None, 2.0), pt=(0.9 * u.GeV, None), nMeasurementsMin=9, maxHoles=2, maxOutliers=2, maxSharedHits=2),
0693             TrackSelectorConfig(absEta=(None, 2.6), pt=(0.4 * u.GeV, None), nMeasurementsMin=8, maxHoles=2, maxOutliers=2, maxSharedHits=2),
0694             TrackSelectorConfig(absEta=(None, 4.0), pt=(0.4 * u.GeV, None), nMeasurementsMin=7, maxHoles=2, maxOutliers=2, maxSharedHits=2),
0695         )
0696         # fmt: on
0697
0698     if args.odd:
0699         ckfConfig = ckfConfig._replace(
0700             pixelVolumes=[16, 17, 18],
0701             stripVolumes=[23, 24, 25],
0702             maxPixelHoles=1,
0703             maxStripHoles=2,
0704             constrainToVolumes=[
0705                 2,  # beam pipe
0706                 32,
0707                 4,  # beam pip gap
0708                 16,
0709                 17,
0710                 18,  # pixel
0711                 20,  # PST
0712                 23,
0713                 24,
0714                 25,  # short strip
0715                 26,
0716                 8,  # long strip gap
0717                 28,
0718                 29,
0719                 30,  # long strip
0720             ],
0721         )
0722     elif args.itk:
0723         ckfConfig = ckfConfig._replace(
0724             # ITk volumes from Noemi's plot
0725             pixelVolumes=[8, 9, 10, 13, 14, 15, 16, 18, 19, 20],
0726             stripVolumes=[22, 23, 24],
0727             maxPixelHoles=1,
0728             maxStripHoles=2,
0729         )
0730
0731     if args.output_detail == 1:
0732         writeDetail = dict(writeTrackSummary=False)
0733     elif args.output_detail == 2:
0734         writeDetail = dict(writeTrackStates=True)
0735     else:
0736         writeDetail = {}
0737
0738     if args.odd and args.output_detail != 1:
0739         writeCovMat = dict(writeCovMat=True)
0740     else:
0741         writeCovMat = {}
0742
0743     addCKFTracks(
0744         s,
0745         trackingGeometry,
0746         field,
0747         trackSelectorConfig=trackSelectorConfig,
0748         ckfConfig=ckfConfig,
0749         **writeDetail,
0750         **writeCovMat,
0751         outputDirRoot=outputDirLessRoot,
0752         outputDirCsv=outputDirLessCsv,
0753     )
0754
0755     if args.ambi_solver == "ML":
0756
0757         from acts.examples.reconstruction import (
0758             addAmbiguityResolutionML,
0759             AmbiguityResolutionMLConfig,
0760         )
0761
0762         addAmbiguityResolutionML(
0763             s,
0764             AmbiguityResolutionMLConfig(
0765                 maximumSharedHits=3, maximumIterations=1000000, nMeasurementsMin=7
0766             ),
0767             onnxModelFile=str(
0768                 actsDir
0769                 / "Examples/Scripts/Python/MLAmbiguityResolution/duplicateClassifier.onnx"
0770             ),
0771             outputDirRoot=outputDirLessRoot,
0772             outputDirCsv=outputDirLessCsv,
0773         )
0774
0775     elif args.ambi_solver == "scoring":
0776
0777         from acts.examples.reconstruction import (
0778             addScoreBasedAmbiguityResolution,
0779             ScoreBasedAmbiguityResolutionConfig,
0780         )
0781
0782         addScoreBasedAmbiguityResolution(
0783             s,
0784             ScoreBasedAmbiguityResolutionConfig(
0785                 minScore=0,
0786                 minScoreSharedTracks=1,
0787                 maxShared=2,
0788                 minUnshared=3,
0789                 maxSharedTracksPerMeasurement=2,
0790                 useAmbiguityScoring=False,
0791             ),
0792             ambiVolumeFile=args.ambi_config,
0793             **writeCovMat,
0794             outputDirRoot=outputDirLessRoot,
0795             outputDirCsv=outputDirLessCsv,
0796         )
0797
0798     elif args.ambi_solver == "greedy":
0799
0800         addAmbiguityResolution(
0801             s,
0802             AmbiguityResolutionConfig(
0803                 maximumSharedHits=3,
0804                 maximumIterations=10000 if args.itk else 1000000,
0805                 nMeasurementsMin=6 if args.itk else 7,
0806             ),
0807             **writeDetail,
0808             **writeCovMat,
0809             outputDirRoot=outputDirLessRoot,
0810             outputDirCsv=outputDirLessCsv,
0811         )
0812
0813     if args.vertexing:
0814         addVertexFitting(
0815             s,
0816             field,
0817             vertexFinder=VertexFinder.AMVF,
0818             outputDirRoot=outputDirLessRoot,
0819             outputDirCsv=outputDirLessCsv,
0820         )
0821
0822     return s
0823
0824
0825 def strToRange(s: str, optName: str, unit: float = 1.0):
0826     global logger
0827     try:
0828         range = [float(e) * unit if e != "" else None for e in s.split(":")]
0829     except ValueError:
0830         range = []
0831     if len(range) == 1:
0832         range.append(range[0])  # 100 -> 100:100
0833     if len(range) != 2:
0834         logger.fatal(f"bad option value: {optName} {s}")
0835         sys.exit(2)
0836     return range
0837
0838
0839 # Graciously taken from https://stackoverflow.com/a/60750535/4280680 (via seeding.py)
0840 class EnumAction(argparse.Action):
0841     """
0842     Argparse action for handling Enums
0843     """
0844
0845     def __init__(self, **kwargs):
0846         import enum
0847
0848         # Pop off the type value
0849         enum_type = kwargs.pop("enum", None)
0850
0851         # Ensure an Enum subclass is provided
0852         if enum_type is None:
0853             raise ValueError("type must be assigned an Enum when using EnumAction")
0854         if not issubclass(enum_type, enum.Enum):
0855             raise TypeError("type must be an Enum when using EnumAction")
0856
0857         # Generate choices from the Enum
0858         kwargs.setdefault("choices", tuple(e.name for e in enum_type))
0859
0860         super(EnumAction, self).__init__(**kwargs)
0861
0862         self._enum = enum_type
0863
0864     def __call__(self, parser, namespace, values, option_string=None):
0865         for e in self._enum:
0866             if e.name == values:
0867                 setattr(namespace, self.dest, e)
0868                 break
0869         else:
0870             raise ValueError("%s is not a validly enumerated algorithm." % values)
0871
0872
0873 # main program: parse arguments, setup sequence, and run the full chain
0874 full_chain(parse_args()).run()