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.GridTriplet,
0054 ):
0055     from acts.examples.simulation import (
0056         addParticleGun,
0057         EtaConfig,
0058         PhiConfig,
0059         ParticleConfig,
0060         addFatras,
0061         addDigitization,
0062         ParticleSelectorConfig,
0063         addDigiParticleSelection,
0064     )
0065
0066     s = s or acts.examples.Sequencer(
0067         events=100, numThreads=-1, logLevel=acts.logging.INFO
0068     )
0069     rnd = acts.examples.RandomNumbers(seed=42)
0070     outputDir = Path(outputDir)
0071
0072     addParticleGun(
0073         s,
0074         EtaConfig(-2.0, 2.0),
0075         ParticleConfig(4, acts.PdgParticle.eMuon, True),
0076         PhiConfig(0.0, 360.0 * u.degree),
0077         multiplicity=2,
0078         outputDirCsv=outputDir / "csv",
0079         outputDirRoot=outputDir,
0080         rnd=rnd,
0081     )
0082
0083     addFatras(
0084         s,
0085         trackingGeometry,
0086         field,
0087         outputDirCsv=outputDir / "csv",
0088         outputDirRoot=outputDir,
0089         rnd=rnd,
0090     )
0091
0092     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0093     addDigitization(
0094         s,
0095         trackingGeometry,
0096         field,
0097         digiConfigFile=srcdir / "Examples/Configs/generic-digi-smearing-config.json",
0098         rnd=rnd,
0099     )
0100
0101     addDigiParticleSelection(
0102         s,
0103         ParticleSelectorConfig(
0104             pt=(1.0 * u.GeV, None),
0105             eta=(-2.5, 2.5),
0106             measurements=(9, None),
0107             removeNeutral=True,
0108         ),
0109     )
0110
0111     from acts.examples.reconstruction import (
0112         addSeeding,
0113         SeedFinderConfigArg,
0114         SeedFinderOptionsArg,
0115     )
0116
0117     addSeeding(
0118         s,
0119         trackingGeometry,
0120         field,
0121         SeedFinderConfigArg(
0122             r=(None, 200 * u.mm),  # rMin=default, 33mm
0123             deltaR=(1 * u.mm, 300 * u.mm),
0124             collisionRegion=(-250 * u.mm, 250 * u.mm),
0125             z=(-2000 * u.mm, 2000 * u.mm),
0126             maxSeedsPerSpM=1,
0127             sigmaScattering=50,
0128             radLengthPerSeed=0.1,
0129             minPt=500 * u.MeV,
0130             impactMax=3 * u.mm,
0131         ),
0132         SeedFinderOptionsArg(
0133             bFieldInZ=2 * u.T,
0134         ),
0135         acts.logging.VERBOSE,
0136         seedingAlgorithm=seedingAlgorithm,
0137         geoSelectionConfigFile=srcdir / "Examples/Configs/generic-seeding-config.json",
0138         outputDirRoot=outputDir,
0139     )
0140     return s
0141
0142
0143 if "__main__" == __name__:
0144     p = argparse.ArgumentParser(
0145         description="Example script to run seed finding",
0146     )
0147
0148     p.add_argument(
0149         "--algorithm",
0150         action=EnumAction,
0151         enum=SeedingAlgorithm,
0152         default=SeedingAlgorithm.GridTriplet,
0153         help="Select the seeding algorithm to use",
0154     )
0155
0156     args = p.parse_args()
0157     # detector = getOpenDataDetector()
0158     detector = acts.examples.GenericDetector()
0159     trackingGeometry = detector.trackingGeometry()
0160
0161     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0162
0163     runSeeding(
0164         trackingGeometry, field, outputDir=Path.cwd(), seedingAlgorithm=args.algorithm
0165     ).run()