Scripts/Python/seeding.py

0001 #!/usr/bin/env python3
0002 from pathlib import Path
0003 from enum import Enum
0004 import argparse
0005
0006 import acts
0007 import acts.examples
0008 from acts.examples.odd import getOpenDataDetector
0009
0010 u = acts.UnitConstants
0011
0012
0013 # Graciously taken from https://stackoverflow.com/a/60750535/4280680
0014 class EnumAction(argparse.Action):
0015     """
0016     Argparse action for handling Enums
0017     """
0018
0019     def __init__(self, **kwargs):
0020         # Pop off the type value
0021         enum_type = kwargs.pop("enum", None)
0022
0023         # Ensure an Enum subclass is provided
0024         if enum_type is None:
0025             raise ValueError("type must be assigned an Enum when using EnumAction")
0026         if not issubclass(enum_type, Enum):
0027             raise TypeError("type must be an Enum when using EnumAction")
0028
0029         # Generate choices from the Enum
0030         kwargs.setdefault("choices", tuple(e.name for e in enum_type))
0031
0032         super(EnumAction, self).__init__(**kwargs)
0033
0034         self._enum = enum_type
0035
0036     def __call__(self, parser, namespace, values, option_string=None):
0037         for e in self._enum:
0038             if e.name == values:
0039                 setattr(namespace, self.dest, e)
0040                 break
0041         else:
0042             raise ValueError("%s is not a validly enumerated algorithm." % values)
0043
0044
0045 from acts.examples.reconstruction import SeedingAlgorithm
0046
0047
0048 def runSeeding(
0049     trackingGeometry,
0050     field,
0051     outputDir,
0052     s=None,
0053     seedingAlgorithm=SeedingAlgorithm.Default,
0054 ):
0055     from acts.examples.simulation import (
0056         addParticleGun,
0057         EtaConfig,
0058         PhiConfig,
0059         ParticleConfig,
0060         ParticleSelectorConfig,
0061         addFatras,
0062         addDigitization,
0063     )
0064
0065     s = s or acts.examples.Sequencer(
0066         events=100, numThreads=-1, logLevel=acts.logging.INFO
0067     )
0068     rnd = acts.examples.RandomNumbers(seed=42)
0069     outputDir = Path(outputDir)
0070
0071     addParticleGun(
0072         s,
0073         EtaConfig(-2.0, 2.0),
0074         ParticleConfig(4, acts.PdgParticle.eMuon, True),
0075         PhiConfig(0.0, 360.0 * u.degree),
0076         multiplicity=2,
0077         outputDirCsv=outputDir / "csv",
0078         outputDirRoot=outputDir,
0079         rnd=rnd,
0080     )
0081
0082     addFatras(
0083         s,
0084         trackingGeometry,
0085         field,
0086         outputDirCsv=outputDir / "csv",
0087         outputDirRoot=outputDir,
0088         rnd=rnd,
0089         preSelectParticles=None,
0090         postSelectParticles=ParticleSelectorConfig(
0091             pt=(1.0 * u.GeV, None),
0092             eta=(-2.5, 2.5),
0093             hits=(9, None),
0094             removeNeutral=True,
0095         ),
0096     )
0097
0098     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0099     addDigitization(
0100         s,
0101         trackingGeometry,
0102         field,
0103         digiConfigFile=srcdir
0104         / "Examples/Algorithms/Digitization/share/default-smearing-config-generic.json",
0105         rnd=rnd,
0106     )
0107
0108     from acts.examples.reconstruction import (
0109         addSeeding,
0110         SeedFinderConfigArg,
0111         SeedFinderOptionsArg,
0112     )
0113
0114     addSeeding(
0115         s,
0116         trackingGeometry,
0117         field,
0118         SeedFinderConfigArg(
0119             r=(None, 200 * u.mm),  # rMin=default, 33mm
0120             deltaR=(1 * u.mm, 60 * u.mm),
0121             collisionRegion=(-250 * u.mm, 250 * u.mm),
0122             z=(-2000 * u.mm, 2000 * u.mm),
0123             maxSeedsPerSpM=1,
0124             sigmaScattering=50,
0125             radLengthPerSeed=0.1,
0126             minPt=500 * u.MeV,
0127             impactMax=3 * u.mm,
0128         ),
0129         SeedFinderOptionsArg(
0130             bFieldInZ=2 * u.T,
0131         ),
0132         acts.logging.VERBOSE,
0133         seedingAlgorithm=seedingAlgorithm,
0134         geoSelectionConfigFile=srcdir
0135         / "Examples/Algorithms/TrackFinding/share/geoSelection-genericDetector.json",
0136         outputDirRoot=outputDir,
0137     )
0138     return s
0139
0140
0141 if "__main__" == __name__:
0142     p = argparse.ArgumentParser(
0143         description="Example script to run seed finding",
0144     )
0145
0146     p.add_argument(
0147         "--algorithm",
0148         action=EnumAction,
0149         enum=SeedingAlgorithm,
0150         default=SeedingAlgorithm.Default,
0151         help="Select the seeding algorithm to use",
0152     )
0153
0154     args = p.parse_args()
0155     # detector = getOpenDataDetector()
0156     detector = acts.examples.GenericDetector()
0157     trackingGeometry = detector.trackingGeometry()
0158
0159     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0160
0161     runSeeding(
0162         trackingGeometry, field, outputDir=Path.cwd(), seedingAlgorithm=args.algorithm
0163     ).run()