epic-lfhcal-tbana/OldStructure/dt5202_dataconverter.py

0001 #!/usr/bin/python3
0002
0003 import numpy as np
0004 import pandas as pd
0005 from datetime import datetime
0006 import argparse
0007
0008 import uproot
0009
0010
0011 def dt5202_text(file_name, mapping = None):
0012         with open(file_name, 'r') as file:
0013                 tstamps = []
0014                 trgids = []
0015
0016                 boards = []
0017                 channels = []
0018                 LGs = []
0019                 HGs = []
0020
0021                 boards_evt = []
0022                 channels_evt = []
0023                 LGs_evt = []
0024                 HGs_evt = []
0025
0026                 layers = []
0027                 assemblies = []
0028                 layer_chs = []
0029                 layer_rows = []
0030                 layer_cols = []
0031
0032                 layer_evt = []
0033                 assembly_evt = []
0034                 layer_ch_evt = []
0035                 layer_row_evt = []
0036                 layer_col_evt = []
0037
0038
0039                 for line in file:
0040
0041                         if line.startswith('//'):
0042                                 pass
0043                         elif 'Tstamp_us' in line:
0044                                 pass
0045                         else:
0046                                 tokens = line.split()
0047                                 if len(tokens) == 6:
0048                                         if len(tstamps) > 0:
0049                                                 boards.append(boards_evt)
0050                                                 channels.append(channels_evt)
0051                                                 LGs.append(LGs_evt)
0052                                                 HGs.append(HGs_evt)
0053                                                 layers.append(layer_evt)
0054                                                 assemblies.append(assembly_evt)
0055                                                 layer_chs.append(layer_ch_evt)
0056                                                 layer_rows.append(layer_row_evt)
0057                                                 layer_cols.append(layer_col_evt)
0058
0059                                                 boards_evt = []
0060                                                 channels_evt = []
0061                                                 LGs_evt = []
0062                                                 HGs_evt = []
0063                                                 layer_evt = []
0064                                                 assembly_evt = []
0065                                                 layer_ch_evt = []
0066                                                 layer_row_evt = []
0067                                                 layer_col_evt = []
0068
0069
0070                                         tstamps.append(float(tokens[0]))
0071                                         trgids.append(int(tokens[1]))
0072                                         boards_evt.append(int(tokens[2]))
0073                                         channels_evt.append(int(tokens[3]))
0074                                         LGs_evt.append(int(tokens[4]))
0075                                         HGs_evt.append(int(tokens[5]))
0076
0077                                         if mapping:
0078                                             try:
0079                                                 layer, assembly, layer_ch, layer_row, layer_col = mapping[(int(tokens[2]), int(tokens[3]))]
0080                                                 layer_evt.append(layer)
0081                                                 assembly_evt.append(assembly)
0082                                                 layer_ch_evt.append(layer_ch)
0083                                                 layer_row_evt.append(layer_row)
0084                                                 layer_col_evt.append(layer_col)
0085                                             except KeyError:
0086                                                 print(f'No mapping for board {int(tokens[2])} channel {int(tokens[3])}')
0087                                                 layer_evt.append(-1)
0088                                                 assembly_evt.append(-1)
0089                                                 layer_ch_evt.append(-1)
0090                                                 layer_row_evt.append(-1)
0091                                                 layer_col_evt.append(-1)
0092                                         else:
0093                                             layer_evt.append(-1)
0094                                             assembly_evt.append(-1)
0095                                             layer_ch_evt.append(-1)
0096                                             layer_row_evt.append(-1)
0097                                             layer_col_evt.append(-1)
0098
0099                                 elif len(tokens) == 4:
0100                                         boards_evt.append(int(tokens[0]))
0101                                         channels_evt.append(int(tokens[1]))
0102                                         LGs_evt.append(int(tokens[2]))
0103                                         HGs_evt.append(int(tokens[3]))
0104
0105                                         if mapping:
0106                                             try:
0107                                                 layer, assembly, layer_ch, layer_row, layer_col = mapping[(int(tokens[0]), int(tokens[1]))]
0108                                                 layer_evt.append(layer)
0109                                                 assembly_evt.append(assembly)
0110                                                 layer_ch_evt.append(layer_ch)
0111                                                 layer_row_evt.append(layer_row)
0112                                                 layer_col_evt.append(layer_col)
0113                                             except KeyError:
0114                                                 print(f'No mapping for board {int(tokens[0])} channel {int(tokens[1])}')
0115                                                 layer_evt.append(-1)
0116                                                 assembly_evt.append(-1)
0117                                                 layer_ch_evt.append(-1)
0118                                                 layer_row_evt.append(-1)
0119                                                 layer_col_evt.append(-1)
0120                                         else:
0121                                             layer_evt.append(-1)
0122                                             assembly_evt.append(-1)
0123                                             layer_ch_evt.append(-1)
0124                                             layer_row_evt.append(-1)
0125                                             layer_col_evt.append(-1)
0126
0127
0128
0129                 boards.append(boards_evt)
0130                 channels.append(channels_evt)
0131                 LGs.append(LGs_evt)
0132                 HGs.append(HGs_evt)
0133                 layers.append(layer_evt)
0134                 assemblies.append(assembly_evt)
0135                 layer_chs.append(layer_ch_evt)
0136                 layer_rows.append(layer_row_evt)
0137                 layer_cols.append(layer_col_evt)
0138
0139
0140         return  tstamps, trgids, boards, channels, LGs, HGs, layers, assemblies, layer_chs, layer_rows, layer_cols
0141
0142
0143
0144 def flatten(tstamps, trgids, boards, channels, LGs, HGs, layers, assemblies, layer_chs, layer_rows, layer_cols):
0145         tstamps_flat = []
0146         trgids_flat = []
0147
0148         boards_flat = []
0149         channels_flat = []
0150         LGs_flat = []
0151         HGs_flat = []
0152
0153         layers_flat = []
0154         assemblies_flat = []
0155         layer_chs_flat = []
0156         layer_rows_flat = []
0157         layer_cols_flat = []
0158
0159         for i in range(len(tstamps)):
0160                 for j in range(len(channels[i])):
0161                         tstamps_flat.append(tstamps[i])
0162                         trgids_flat.append(trgids[i])
0163
0164                         boards_flat.append(boards[i][j])
0165                         channels_flat.append(channels[i][j])
0166                         LGs_flat.append(LGs[i][j])
0167                         HGs_flat.append(HGs[i][j])
0168
0169                         layers_flat.append(layers[i][j])
0170                         assemblies_flat.append(assemblies[i][j])
0171                         layer_chs_flat.append(layer_chs[i][j])
0172                         layer_rows_flat.append(layer_rows[i][j])
0173                         layer_cols_flat.append(layer_cols[i][j])
0174
0175         return np.array(tstamps_flat), np.array(trgids_flat), np.array(boards_flat), np.array(channels_flat), np.array(LGs_flat), np.array(HGs_flat), np.array(layers_flat), np.array(assemblies_flat), np.array(layer_chs_flat), np.array(layer_rows_flat), np.array(layer_cols_flat)
0176
0177
0178 # def flatten(tstamps, trgids, boards, channels, LGs, HGs):
0179 #         tstamps_flat = []
0180 #         trgids_flat = []
0181
0182 #         boards_flat = []
0183 #         channels_flat = []
0184 #         LGs_flat = []
0185 #         HGs_flat = []
0186
0187 #         for i in range(len(tstamps)):
0188 #                 for j in range(len(channels[i])):
0189 #                         tstamps_flat.append(tstamps[i])
0190 #                         trgids_flat.append(trgids[i])
0191
0192 #                         boards_flat.append(boards[i][j])
0193 #                         channels_flat.append(channels[i][j])
0194 #                         LGs_flat.append(LGs[i][j])
0195 #                         HGs_flat.append(HGs[i][j])
0196
0197 #         return np.array(tstamps_flat), np.array(trgids_flat), np.array(boards_flat), np.array(channels_flat), np.array(LGs_flat), np.array(HGs_flat)
0198
0199 def save_as_numpy(tstamps, trgids, boards, channels, LGs, HGs, layers, assemblies, layer_chs, layer_rows, layer_cols, file_out):
0200         tstamps_flat, trgids_flat, boards_flat, channels_flat, LGs_flat, HGs_flat, layers_flat, assemblies_flat, layer_chs_flat, layer_rows_flat, layer_cols_flat = flatten(tstamps, trgids, boards, channels, LGs, HGs, layers, assemblies, layer_chs, layer_rows, layer_cols)
0201
0202         data = np.array([tstamps_flat, trgids_flat, boards_flat, channels_flat, LGs_flat, HGs_flat,layers_flat, assemblies_flat, layer_chs_flat, layer_rows_flat, layer_cols_flat])
0203         np.save(file_out, data)
0204
0205
0206 def save_as_root(tstamps, trgids, boards, channels, LGs, HGs, layers, assemblies, layer_chs, layer_rows, layer_cols, file_out):
0207     file = uproot.recreate(file_out)
0208
0209     file["tree"] = {"t_stamp": tstamps,
0210                     "trgid": trgids,
0211                     "board": boards,
0212                     "channel": channels,
0213                     "LG": LGs,
0214                     "HG": HGs,
0215                     "layer": layers,
0216                     "assembly": assemblies,
0217                     "layer_ch": layer_chs,
0218                     "layer_row": layer_rows,
0219                     "layer_col": layer_cols}
0220     file.close()
0221
0222 def save_as_dataframe(tstamps, trgids, boards, channels, LGs, HGs, layers, assemblies, layer_chs, layer_rows, layer_cols, file_out):
0223         tstamps_flat, trgids_flat, boards_flat, channels_flat, LGs_flat, HGs_flat, layers_flat, assemblies_flat, layer_chs_flat, layer_rows_flat, layer_cols_flat = flatten(tstamps, trgids, boards, channels, LGs, HGs, layers, assemblies, layer_chs, layer_rows, layer_cols)
0224
0225         df = pd.DataFrame({'t_stamp': tstamps_flat,
0226                    'trgid': trgids_flat,
0227                    'board': boards_flat,
0228                    'channel': channels_flat,
0229                    'LG': LGs_flat,
0230                    'HG': HGs_flat,
0231                    'layer': layers_flat,
0232                    'assembly': assemblies_flat,
0233                    'layer_ch': layer_chs_flat,
0234                    'layer_row': layer_rows_flat,
0235                    'layer_col': layer_cols_flat})
0236
0237         df.to_pickle(file_out)
0238
0239 def is_valid_mapping(parser, mapping):
0240     if not mapping:
0241         return False
0242     try:
0243         with open(mapping, 'r') as file:
0244             mapping = {}
0245             for line in file:
0246                 if line.startswith('#'):
0247                     continue
0248                 if len(line.strip()) == 0:
0249                     continue
0250                 tokens = line.split()
0251                 if len(tokens) != 7:
0252                     print(line)
0253                     parser.error('Mapping file must have 7 columns')
0254                 for t in tokens:
0255                     if not t.isdigit():
0256                         parser.error('All columns in mapping file must be integers')
0257                 caen_board = int(tokens[0])
0258                 caen_ch = int(tokens[1])
0259                 layer = int(tokens[2])
0260                 assembly = int(tokens[3])
0261                 layer_ch = int(tokens[4])
0262                 layer_row = int(tokens[5])
0263                 layer_col = int(tokens[6])
0264
0265                 mapping[(caen_board, caen_ch)] = (layer, assembly, layer_ch, layer_row, layer_col)
0266
0267         return mapping
0268     except FileNotFoundError:
0269         parser.error('Mapping file not found')
0270     except Exception as e:
0271         parser.error(f'Error reading mapping file: {e}')
0272     return False
0273
0274 if __name__ == '__main__':
0275     parser = argparse.ArgumentParser()
0276     parser.add_argument('-m', '--mapping', type = lambda x:is_valid_mapping(parser, x),help='mapping file')
0277     parser.add_argument('-r', '--root', default=False, action='store_true', help='output as root file')
0278     parser.add_argument('-p', '--pandas', default=False, action='store_true', help='output as pickled pandas dataframe')
0279     parser.add_argument('-n', '--numpy',  default=False, action='store_true', help='output as numpy file')
0280     parser.add_argument('file', type=str, nargs='+', help='Input file(s)')
0281
0282     args = parser.parse_args()
0283
0284     if args.mapping:
0285         print(f'Using mapping file: {args.mapping}')
0286     else:
0287         print('No mapping file specified')
0288
0289
0290     if not (args.root or args.pandas or args.numpy):
0291         parser.error('No output format specified, add -r, -p or -n')
0292
0293     if not args.file:
0294         parser.error('No input file specified')
0295
0296     for f in args.file:
0297         if not f.endswith('.txt'):
0298             parser.error('Input file must be a .txt file')
0299
0300     for f in args.file:
0301         print(f'{f}: ', end = '')
0302         tstamps, trgids, boards, channels, LGs, HGs, layers, assemblies, layer_chs, layer_rows, layer_cols = dt5202_text(f, args.mapping)
0303         if args.root:
0304             file_out_root = f.replace('.txt', '.root')
0305             save_as_root(tstamps, trgids, boards, channels, LGs, HGs, layers, assemblies, layer_chs, layer_rows, layer_cols, file_out_root)
0306             print(f'root ', end = '')
0307         if args.pandas:
0308             file_out_pkl = f.replace('.txt', '.pkl')
0309             save_as_dataframe(tstamps, trgids, boards, channels, LGs, HGs, layers, assemblies, layer_chs, layer_rows, layer_cols, file_out_pkl)
0310             print(f'pandas ', end = '')
0311         if args.numpy:
0312             file_out_npy = f.replace('.txt', '.npy')
0313             save_as_numpy(tstamps, trgids, boards, channels, LGs, HGs, layers, assemblies, layer_chs, layer_rows, layer_cols, file_out_npy)
0314             print(f'numpy ', end = '')
0315
0316     print('done.')
0317
0318