DRAO vdif conversion script (quite a hack)

drao_vdif.py

from baseband import vdif
from baseband.vdif.header import VDIFHeader0, HeaderParser, VDIFHeader
import numpy as np
import astropy.units as u


files = list('{:07d}.dat'.format(no) for no in range(100000, 100001))


class DRAOVDIFHeader(VDIFHeader0):

    _header_parser = VDIFHeader0._header_parser + HeaderParser(
        (('link', (3, 16, 4)),
         ('slot', (3, 20, 6)),
         ('eud2', (5, 0, 32))))

    def __new__(cls, words, edv=None, verify=True, **kwargs):
        return object.__new__(cls)

    def verify(self):
        pass

    @classmethod
    def fromfile(cls, fh, edv=0, verify=False):
        self = super(DRAOVDIFHeader, cls).fromfile(fh, edv=0, verify=False)
        # Correct wrong bps
        self.mutable = True
        self['bits_per_sample'] = 3
        return self


if __name__ == '__main__':
    nfiles = len(files)
    fhs = []
    fils = []
    headers = []
    eud2s = []
    eud2_mins = []
    eud2_maxs = []
    links = []
    slots = []
    pols = []
    fh = np.memmap('drao/' + files[0], dtype=np.uint32, mode='r')
    fhs.append(fh)
    fil = fh.reshape(-1, 5032 // 4)
    fils.append(fil)
    header = DRAOVDIFHeader(fil.T[:8], edv=0, verify=False)
    headers.append(header)
    eud2s.append(header['eud2'].copy())
    links.append(header['link'])
    slots.append(header['slot'])
    pols.append(header['station_id'])
    header0 = DRAOVDIFHeader(fil[eud2s[0].argmin()], edv=0, verify=False)
    eud2_mins.append(header0['eud2'])
    eud2_maxs.append(eud2s[0].max())

    time0 = header0.get_time(framerate=625*u.Hz)
    output_header = VDIFHeader.fromvalues(edv=0, bps=4, complex_data=True,
                                          nchan=1024, samples_per_frame=1,
                                          station='Pe')
    fw = vdif.open('try.vdif', 'ws', header=output_header, nthread=2,
                   frames_per_second=390625)
    output_header.set_time(time0, framerate=400.*u.MHz/1024)

    file_max = 0
    eud2_sel = eud2_mins[0]
    while True:
        if file_max == nfiles - 1 and eud2_sel > eud2_maxs[-1]:
            break

        if eud2_sel > eud2_mins[-1] and file_max < nfiles-1:
            file_max += 1
            fh = np.memmap('drao/' + files[file_max], dtype='<u4',
                           mode='r')
            fhs.append(fh)
            fil = fh.reshape(-1, 5032 // 4)
            fils.append(fil)
            header = DRAOVDIFHeader(fil.T[:8], edv=0, verify=False)
            headers.append(header)
            eud2s.append(header['eud2'].copy())
            links.append(header['link'])
            slots.append(header['slot'])
            pols.append(header['station_id'])
            eud2_mins.append(eud2s[-1].min())
            eud2_maxs.append(eud2s[-1].max())

        out = np.zeros((625, 2, 1024), dtype=np.complex64)
        for eud2, fil, slot, link, pol in zip(eud2s, fils, slots, links, pols):
            select = np.where(eud2 == eud2_sel)[0]
            print(eud2_sel, len(select))
            if len(select) > 1:
                for i in select:
                    payload = vdif.VDIFPayload(fil[i, 8:],
                                               nchan=header0.nchan,
                                               bps=4, complex_data=True)
                    freq_ids = slot[i] + 16 * link[i] + 128 * np.arange(8)
                    out[:, pol[i], freq_ids] = payload.data
        fw.write(out)

        eud2_sel += 625
        if eud2_sel > eud2_maxs[0]:
            fh = fhs.pop(0)
            del fh
            fils.pop(0)
            headers.pop(0)
            eud2s.pop(0)
            links.pop(0)
            slots.pop(0)
            pols.pop(0)
            eud2_mins.pop(0)
            eud2_maxs.pop(0)
            if file_max == nfiles - 1:
                break