Utilities/scripts/scoutingToRaw.py

0001 #!/bin/env python3
0002
0003 import struct
0004 import os,sys
0005 import json
0006 import shutil
0007
0008 os.umask(0)
0009
0010 #struct muon{
0011 #  uint32_t f;
0012 #  uint32_t s;
0013 #};
0014
0015 #struct block{
0016 #  uint32_t bx;
0017 #  uint32_t orbit;
0018 #  muon mu[16];
0019 #};
0020
0021 #class masks:
0022 #  phiext  = 0x3ff
0023 #  pt      = 0x1ff
0024 #  qual    = 0xf
0025 #  etaext  = 0x1ff
0026 #  etaextv = 0xff
0027 #  etaexts = 0x100
0028 #  iso     = 0x3
0029 #  chrg    = 0x1
0030 #  chrgv   = 0x1
0031 #  index   = 0x7f
0032 #  phi     = 0x3ff
0033 #  eta     = 0x1ff
0034 #  etav    = 0xff
0035 #  etas    = 0x100
0036 #  phiv    = 0x1ff
0037 #  phis    = 0x200
0038 #  sv      = 0x3
0039
0040 #class shifts:
0041 #  phiext  = 0
0042 #  pt      = 10
0043 #  qual    = 19
0044 #  etaext  = 23
0045 #  iso     = 0
0046 #  chrg    = 2
0047 #  chrgv   = 3
0048 #  index   = 4
0049 #  phi     = 11
0050 #  eta     = 21
0051 #  rsv     = 30
0052
0053 #class gmt_scales:
0054 #  pt_scale    = 0.5
0055 #  phi_scale   = 2.*M_PI/576.
0056 #  eta_scale   = 0.0870/8 #9th MS bit is sign
0057 #  phi_range   = M_PI
0058
0059
0060 #need to read this to find orbit ("event") boundary and calculate size per orbit
0061 class header_shifts:
0062   bxmatch    = 32;
0063   mAcount    = 16;
0064   orbitmatch = 8;
0065   mBcount    = 0;
0066
0067 class header_masks:
0068   bxmatch    = 0xff << header_shifts.bxmatch;
0069   mAcount    = 0xf  << header_shifts.mAcount;
0070   orbitmatch = 0xff << header_shifts.orbitmatch;
0071   mBcount    = 0xf
0072
0073
0074 #new V2 FRD file header (32 bytes)
0075 class frd_file_header_v2:
0076   ver_id = "RAW_0002".encode() # 64 (offset 0B)
0077   header_size = 32 #16 (offset 8B)
0078   data_type = 20 #16 (offset 10)
0079   event_count = 0 #32 (offset 12B)
0080   run_number = 0 #32 (offset 16B)
0081   lumisection = 0 #32 (offset 20B)
0082   file_size = 0 #64 (offset 24B)
0083
0084
0085 def parseMuonScoutingRawFile(infilepath, outdir, rn_override, maxorbits):
0086
0087   if infilepath != 'stdin':
0088     fin = open(infilepath,'rb')
0089   else:
0090     fin = sys.stdin.buffer
0091
0092   #sys.stdout.flush()
0093
0094   #orbit count per file
0095   orbitcount=0
0096   #total
0097   orbitcount_total=0
0098
0099   last_ls = 0
0100
0101   orbit_data = bytes()
0102   orbit_nr = 0
0103   orbit_size = 0
0104   flags = 0
0105   c_crc32c = 0
0106
0107   #ls = 1
0108   #event header (FRD format) const
0109   version = 6
0110
0111   #files
0112   fout = None
0113   if infilepath != 'stdin':
0114     fin = open(infilepath,'rb')
0115   else:
0116     fin = sys.stdin.buffer
0117
0118
0119   #write header before closing the file
0120   def update_header():
0121       nonlocal orbitcount
0122       nonlocal last_ls
0123       h = frd_file_header_v2()
0124       h.event_count = orbitcount
0125       h.run_number = rn_override
0126       h.lumisection = last_ls
0127       h.file_size = fout.tell()
0128       fout.seek(0, 0)
0129       fout.write(frd_file_header_v2.ver_id)
0130       fout.write(struct.pack('H',h.header_size))
0131       fout.write(struct.pack('H',h.data_type))
0132       fout.write(struct.pack('I',h.event_count))
0133       fout.write(struct.pack('I',h.run_number))
0134       fout.write(struct.pack('I',h.lumisection))
0135       fout.write(struct.pack('Q',h.file_size))
0136
0137       orbitcount = 0
0138       print(h.ver_id, h.header_size, h.data_type, h.event_count, h.lumisection, h.file_size)
0139
0140
0141   #write orbit when next one is detected or file is closed
0142   def write_orbit():
0143           nonlocal orbit_size
0144           nonlocal orbit_data
0145           if not orbit_size:
0146               return
0147
0148           #print(fout.tell(), struct.pack('H',version))
0149           fout.write(struct.pack('H',version)) #could be 8 bytes
0150           fout.write(struct.pack('H',flags)) #could be 8 bytes
0151           fout.write(struct.pack('I',rn_override)) #run
0152           #fout.write(struct.pack('I',ls)) #ls
0153           fout.write(struct.pack('I',last_ls)) #ls
0154           fout.write(struct.pack('I',orbit_nr)) #eid (orbit number, 32-bit)
0155           fout.write(struct.pack('I',orbit_size)) #payload size
0156           fout.write(struct.pack('I',c_crc32c)) #payload checksum (not used)
0157           fout.write(orbit_data)
0158
0159           orbit_data = bytes()
0160           orbit_size = 0
0161
0162   def writeout_close():
0163       write_orbit()
0164       update_header()
0165       fout.close()
0166       orbit_nr = 0
0167
0168   #read loop
0169   while True:
0170
0171           #check if exceeded max orbits specified
0172           if orbitcount_total > maxorbits:
0173               print(f"finish: {orbitcount_total-1}/{maxorbits} orbits")
0174               writeout_close()
0175
0176               if infilepath != 'stdin':
0177                   fin.close()
0178               sys.exit(0)
0179
0180           try:
0181               h_raw = fin.read(4)
0182               bxmatch = struct.unpack('B', h_raw[3:4])[0]
0183               mAcount = struct.unpack('B', h_raw[2:3])[0]
0184               orbitmatch = struct.unpack('B', h_raw[1:2])[0]
0185               mBcount = struct.unpack('B', h_raw[0:1])[0]
0186
0187               #print("bxmatch", bxmatch, "mA", mAcount, "orbitmatch", orbitmatch, "mB", mBcount)
0188
0189               bx_raw = fin.read(4)
0190               bx = struct.unpack('i', bx_raw)[0]
0191               #print("bx",bx)
0192               orbit_raw = fin.read(4)
0193               orbit = struct.unpack('i', orbit_raw)[0]
0194
0195               new_ls = orbit >> 18
0196
0197               if new_ls > last_ls:
0198               #open a new output file if crossing LS boundary or on first orbit
0199                   if last_ls:
0200                       write_orbit()
0201                       update_header()
0202                       fout.close()
0203                       orbitcount = 0
0204
0205                   last_ls = new_ls
0206                   fout = open(os.path.join(outdir, f'run{rn_override}_ls{str(new_ls).zfill(4)}_index000000.raw') ,'wb')
0207                   #empty file header, will be updated later
0208                   fout.write(frd_file_header_v2.ver_id)
0209 #                  fout.write(bytes(16))
0210                   fout.write(bytes(24))
0211
0212               read_len = 8*(mAcount+mBcount)
0213               mu_blk = fin.read(8*(mAcount+mBcount))
0214               if len(mu_blk) != read_len:
0215                   print('incomplete read')
0216                   sys.exit(1)
0217
0218               if not orbit_nr or orbit != orbit_nr:
0219                   #received new orbit, write previous one
0220                   if orbit_nr:
0221                       write_orbit()
0222
0223                   #should not decrease:
0224                   if orbit < orbit_nr:
0225                       orbit_count = -1
0226                       print("got smaller orbit than earlier!")
0227                       sys.exit(1)
0228
0229                   print("new orbit", orbit)
0230                   orbit_nr = orbit
0231
0232                   #per LS file counter:
0233                   orbitcount += 1
0234                   #total counter:
0235                   orbitcount_total += 1
0236
0237               #update orbit size and data variables
0238               orbit_size += 12 + read_len
0239               orbit_data += (h_raw + bx_raw + orbit_raw) + mu_blk
0240
0241           except Exception as ex:
0242               #reached premature end of the file?
0243               print(f"exception: {ex}")
0244               #writeout_close()
0245               #if infilepath != 'stdin':
0246               #    fin.close()
0247               sys.exit(1)
0248
0249           #print count," : ",version,run,lumi,eid,esize,crc32c,"override id/ls/run:",count,1,rn_override
0250           #lumi=1
0251
0252 if len(sys.argv) < 5:
0253   print("parameters: input file (or stdin), output directory, run number (use same as input file), orbits to write")
0254 else:
0255   parseMuonScoutingRawFile(sys.argv[1], sys.argv[2], int(sys.argv[3]), int(sys.argv[4]))
0256
0257
0258
0259