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