Scripts/Python/muon_hough.py

0001 #!/usr/bin/env python3
0002 import os
0003 import argparse
0004
0005 import acts
0006
0007 from acts.examples import (
0008     CsvMuonSpacePointReader,
0009     CsvMuonSegmentReader,
0010     MuonHoughSeeder,
0011     RootMuonSpacePointReader,
0012 )
0013
0014 u = acts.UnitConstants
0015 rnd = acts.examples.RandomNumbers(seed=42)
0016
0017
0018 def runHoughFromRoot(inFile: str, nEvents: int):
0019     s = acts.examples.Sequencer(
0020         events=nEvents, numThreads=1, logLevel=acts.logging.VERBOSE
0021     )
0022
0023     # Read input space points from input csv files
0024     evReader = RootMuonSpacePointReader(
0025         filePath=inFile,
0026         outputSpacePoints="MuonSpacePoints",
0027         level=acts.logging.VERBOSE,
0028     )
0029     s.addReader(evReader)
0030
0031     s.run()
0032
0033
0034 def runHoughFromCsv(inDir: str, nEvents: int):
0035     # create temporary file with pixel SPs and run the seeding
0036
0037     s = acts.examples.Sequencer(
0038         events=nEvents, numThreads=1, logLevel=acts.logging.VERBOSE
0039     )
0040
0041     # Read input space points from input csv files
0042     evReader = CsvMuonSpacePointReader(
0043         inputStem="SpacePoints",
0044         inputDir=os.path.dirname(inDir),
0045         outputSpacePoints="MuonSpacePoints",
0046         level=acts.logging.VERBOSE,
0047     )
0048
0049     truthReader = CsvMuonSegmentReader(
0050         inputStem="MuonTruthSegment",
0051         inputDir=os.path.dirname(inDir),
0052         outputSegments="MuonTruthSegments",
0053         level=acts.logging.VERBOSE,
0054     )
0055
0056     # add csv reader
0057     s.addReader(evReader)
0058     s.addReader(truthReader)
0059     ### Add the hough seeder algorithm
0060     seeder = MuonHoughSeeder(
0061         inSpacePoints=evReader.config.outputSpacePoints,
0062         inTruthSegments=truthReader.config.outputSegments,
0063         outHoughMax="MuonHoughSeeds",
0064         level=acts.logging.VERBOSE,
0065     )
0066
0067     s.addAlgorithm(seeder)
0068     s.run()
0069
0070
0071 if "__main__" == __name__:
0072     p = argparse.ArgumentParser(
0073         description="Example script to run ITk seed finding based on CSV space points",
0074     )
0075     p.add_argument(
0076         "--input",
0077         help="Path to the script's input. By default it's assumed that a ROOT n-tuple is parsed. Otherwise, it's also possible to parse a CSV directory",
0078     )
0079     p.add_argument(
0080         "--isCSV",
0081         default=False,
0082         action="store_true",
0083         help="Flag toggling that the input is a CSV directory",
0084     )
0085     p.add_argument("--nEvents", default=100, help="Number of events to run", type=int)
0086
0087     args = p.parse_args()
0088     if args.isCSV:
0089         runHoughFromCsv(inFile=args.input, nEvents=args.nEvents)
0090     else:
0091         runHoughFromRoot(args.input, nEvents=args.nEvents)