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 CsvMuonSpacePointReader, CsvMuonSegmentReader
0008 from acts.examples.ActsExamplesPythonBindingsRoot import RootMuonSpacePointReader
0009
0010 u = acts.UnitConstants
0011 rnd = acts.examples.RandomNumbers(seed=42)
0012
0013
0014 def runHoughFromRoot(inFile: str, nEvents: int):
0015     s = acts.examples.Sequencer(
0016         events=nEvents, numThreads=1, logLevel=acts.logging.VERBOSE
0017     )
0018
0019     readerCfg = RootMuonSpacePointReader.Config()
0020     readerCfg.filePath = inFile
0021     readerCfg.outputSpacePoints = "MuonSpacePoints"
0022     readerCfg.treeName = "MuonSpacePoints"
0023
0024     evReader = RootMuonSpacePointReader(
0025         readerCfg,
0026         acts.logging.Level.VERBOSE,
0027     )
0028     s.addReader(evReader)
0029
0030     s.run()
0031
0032
0033 def runHoughFromCsv(inDir: str, nEvents: int, etaOnly: bool):
0034     # create temporary file with pixel SPs and run the seeding
0035
0036     s = acts.examples.Sequencer(
0037         events=nEvents, numThreads=1, logLevel=acts.logging.VERBOSE
0038     )
0039
0040     # Read input space points from input csv files
0041     evReader = CsvMuonSpacePointReader(
0042         inputStem="SpacePoints",
0043         inputDir=os.path.dirname(inDir),
0044         outputSpacePoints="MuonSpacePoints",
0045         level=acts.logging.VERBOSE,
0046     )
0047
0048     truthReader = CsvMuonSegmentReader(
0049         inputStem="MuonTruthSegment",
0050         inputDir=os.path.dirname(inDir),
0051         outputSegments="MuonTruthSegments",
0052         level=acts.logging.VERBOSE,
0053     )
0054
0055     # add csv reader
0056     s.addReader(evReader)
0057     s.addReader(truthReader)
0058     ### Add the hough seeder algorithm
0059     seeder = MuonHoughSeeder(
0060         inSpacePoints=evReader.config.outputSpacePoints,
0061         inTruthSegments=truthReader.config.outputSegments,
0062         outHoughMax="MuonHoughSeeds",
0063         extendWithPhi=not etaOnly,
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     p.add_argument(
0087         "--eta-only",
0088         action="store_true",
0089         help="return Eta Hough maxima without applying the phi extension",
0090     )
0091
0092     args = p.parse_args()
0093     if args.isCSV:
0094         runHoughFromCsv(inDir=args.input, nEvents=args.nEvents, etaOnly=args.eta_only)
0095     else:
0096         runHoughFromRoot(args.input, nEvents=args.nEvents)