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