File indexing completed on 2024-10-19 04:58:01
0001
0002
0003 from argparse import ArgumentParser
0004 import logging
0005 import ROOT
0006
0007 class TFileContext(object):
0008 def __init__(self, *args):
0009 self.tfile = ROOT.TFile(*args)
0010 def __enter__(self):
0011 return self.tfile
0012 def __exit__(self, type, value, traceback):
0013 self.tfile.Close()
0014
0015 def display_results(t_data, style='twiki'):
0016 before = after = None
0017 if (style == 'twiki'):
0018 header_fmt = ' | '.join( ['{:^6s}'] + ['{:^9s}' ] * (len(t_data.keys())-1) )
0019 row_fmt = ' | '.join( ['{:6d}' ] + ['{:9.6f}'] * (len(t_data.keys())-1) )
0020 header_fmt = '| '+header_fmt+' |'
0021 row_fmt = '| '+row_fmt +' |'
0022 elif(style == 'latex'):
0023 before = '\\begin{tabular}{' + 'c'*len(t_data.keys()) + '}'
0024 after = '\\end{tabular}'
0025 header_fmt = ' & '.join( ['{:^6s}'] + ['{:^9s}' ] * (len(t_data.keys())-1) ) + ' \\\\\n\\hline'
0026 row_fmt = ' & '.join( ['{:6d}' ] + ['{:9.6f}'] * (len(t_data.keys())-1) ) + ' \\\\'
0027 elif(style == 'csv'):
0028 header_fmt = ', '.join( ['{:s}' ] * len(t_data.keys()) )
0029 row_fmt = ', '.join( ['{:d}' ] + ['{:f}' ] * (len(t_data.keys())-1) )
0030 else:
0031 raise RuntimeError('Unknown style "%s" for table'%(style))
0032
0033 if(before is not None): print(before)
0034 print( header_fmt.format(*t_data.keys()) )
0035 for i, run in enumerate(t_data['run']):
0036 print(row_fmt.format(run, *[t_data[k][i] for k in t_data.keys() if not k == 'run']))
0037 if(after is not None): print(after)
0038
0039
0040 def main():
0041 parser = ArgumentParser()
0042 parser.add_argument('fname' , metavar='FILE')
0043 parser.add_argument('-p', '--partition', default='BPIX', help='Tracker partition (e.g. BPIX, FPIX, BPIXLYR1, etc.). Default: %(default)s')
0044 parser.add_argument('-l', '--list-content' , action='store_true', dest='list_content', help='List the contents of file and exit')
0045 parser.add_argument( '--list-branches', action='store_true', help='List the branches of the tree and exit')
0046 parser.add_argument('-t', '--type' , default='barycentre', choices=('barycentre', 'beamspot'), type=str.lower, help='Default: %(default)s')
0047 parser.add_argument( '--label' , default=None, help='Additional label that is appended to the folder name (i.e. PixelBaryCentreAnalyzer by default)')
0048 parser.add_argument( '--quality' , action='store_true', help='Read results with the WithPixelQuality flag (default: %(default)s)')
0049 parser.add_argument('-s', '--style' , default='twiki', choices=('twiki', 'latex', 'csv'), type=str.lower, help='Table style for the output (default: %(default)s)')
0050 parser.add_argument( '--loglevel' , metavar='LEVEL', default='WARNING', help='Level for the python logging module. Can be either a mnemonic string like DEBUG, INFO or WARNING or an integer (lower means more verbose).')
0051
0052 args = parser.parse_args()
0053 loglevel = args.loglevel.upper() if not args.loglevel.isdigit() else int(args.loglevel)
0054 logging.basicConfig(format='%(levelname)s:%(module)s:%(funcName)s: %(message)s', level=loglevel)
0055 logging.debug('args: %s', args)
0056
0057
0058 folder = 'PixelBaryCentreAnalyzer' if not args.quality else 'PixelBaryCentreAnalyzerWithPixelQuality'
0059 tree_name = 'PixelBarycentre' if args.type == 'barycentre' else 'BeamSpot'
0060 if(args.label is not None):
0061 tree_name += '_'+args.label
0062 columns = ['run'] + [(args.partition if args.type=='barycentre' else 'BS')+'.'+coord for coord in ('x', 'y', 'z')]
0063
0064 with TFileContext(args.fname) as tfile:
0065 if(args.list_content):
0066 list_content(tfile)
0067 return 0
0068
0069 tfolder = tfile.Get(folder)
0070 if(not tfolder):
0071 raise KeyError('Folder "%s" not found in "%s"' %(folder, args.fname))
0072
0073 logging.debug('Opened folder "%s"', folder)
0074 tree = tfolder.Get(tree_name)
0075 if(not tree):
0076 logging.error('Tree "%s" not found; content of file "%s":', tree_name, args.fname)
0077 list_content(args.fname)
0078 raise KeyError(tree_name)
0079
0080 if(args.list_branches):
0081 list_branches(tree, folder=folder)
0082 return 0
0083
0084 rdf = ROOT.RDataFrame(tree)
0085 logging.info('Reading "%s"', tree_name)
0086 results = rdf.AsNumpy(columns)
0087
0088 display_results(results, style=args.style)
0089
0090 return 0
0091
0092
0093 def list_content(tfile):
0094 for tfolder in tfile.GetListOfKeys():
0095 for key in tfolder.ReadObj().GetListOfKeys():
0096 obj = key.ReadObj()
0097 print('%s\t%s/%s\t%d branches, %d entries' %(obj.ClassName(), tfolder.GetName(), key.GetName(), len(obj.GetListOfBranches()), obj.GetEntries()))
0098
0099
0100 def list_branches(tree, folder_name=''):
0101 branches = [b.GetName() for b in tree.GetListOfBranches()]
0102 logging.info('Branches (%s/%s): %d', (folder_name, tree.GetName(), len(branches)))
0103 print('\n'.join(branches))
0104
0105
0106 if __name__ == '__main__':
0107 exit(main())
