Ejemplo n.º 1
0
def get_dspec(vis=None, savespec=True, specfile=None, bl='', uvrange='', field='', scan='', datacolumn='data',
              domedian=False, timeran=None, spw=None, timebin='0s', regridfreq=False, fillnan=None, verbose=False):
    # from split_cli import split_cli as split
    if vis.endswith('/'):
        vis = vis[:-1]
    msfile = vis
    if not spw:
        spw = ''
    if not timeran:
        timeran = ''
    if not bl:
        bl = ''
    if domedian:
        if not uvrange:
            uvrange = '0.2~0.8km'
        bl = ''
    else:
        uvrange = ''
    # Open the ms and plot dynamic spectrum
    if verbose:
        print('Splitting selected data...')
    vis_spl = './tmpms.splitted'
    if os.path.exists(vis_spl):
        os.system('rm -rf ' + vis_spl)

    # split(vis=msfile, outputvis=vis_spl, timerange=timeran, antenna=bl, field=field, scan=scan, spw=spw,
    #       uvrange=uvrange, timebin=timebin, datacolumn=datacolumn)

    ms.open(msfile, nomodify=False)
    ms.split(outputms=vis_spl, whichcol=datacolumn, time=timeran, spw=spw, baseline=bl, field=field, scan=scan,
             uvrange=uvrange, timebin=timebin)
    ms.close()
    if verbose:
        print('Regridding into a single spectral window...')
        # print('Reading data spw by spw')

    try:
        tb.open(vis_spl + '/POLARIZATION')
        corrtype = tb.getcell('CORR_TYPE', 0)
        pols = [stokesenum[p] for p in corrtype]
        tb.close()
    except:
        pols = []

    if regridfreq:
        ms.open(vis_spl, nomodify=False)
        ms.cvel(outframe='LSRK', mode='frequency', interp='nearest')
        ms.selectinit(datadescid=0, reset=True)
        data = ms.getdata(['amplitude', 'time', 'axis_info'], ifraxis=True)
        specamp = data['amplitude']
        freq = data['axis_info']['freq_axis']['chan_freq']

    else:
        ms.open(vis_spl)
        ms.selectinit(datadescid=0, reset=True)
        spwinfo = ms.getspectralwindowinfo()
        specamp = []
        freq = []
        time = []
        for descid in range(len(spwinfo.keys())):
            ms.selectinit(datadescid=0, reset=True)
            ms.selectinit(datadescid=descid)
            data = ms.getdata(['amplitude', 'time', 'axis_info'], ifraxis=True)
            specamp_ = data['amplitude']
            freq_ = data['axis_info']['freq_axis']['chan_freq']
            time_ = data['time']
            if fillnan is not None:
                flag_ = ms.getdata(['flag', 'time', 'axis_info'], ifraxis=True)['flag']
                if type(fillnan) in [int, float, long]:
                    specamp_[flag_] = float(fillnan)
                else:
                    specamp_[flag_] = 0.0
            specamp.append(specamp_)
            freq.append(freq_)
            time.append(time_)
        specamp = np.concatenate(specamp, axis=1)
        freq = np.concatenate(freq, axis=0)
        ms.selectinit(datadescid=0, reset=True)
    ms.close()
    os.system('rm -rf ' + vis_spl)
    (npol, nfreq, nbl, ntim) = specamp.shape
    freq = freq.reshape(nfreq)

    if verbose:
        print('npol, nfreq, nbl, ntime:', (npol, nfreq, nbl, ntim))
    spec = np.swapaxes(specamp, 2, 1)

    tim = data['time']

    if domedian:
        if verbose:
            print('doing median of all the baselines')
        # mask zero values before median
        spec_masked = np.ma.masked_where(spec < 1e-9, spec)
        spec_med = np.ma.filled(np.ma.median(spec_masked, axis=1), fill_value=0.)
        nbl = 1
        ospec = spec_med.reshape((npol, nbl, nfreq, ntim))
    else:
        ospec = spec
    # Save the dynamic spectral data
    if savespec:
        if not specfile:
            specfile = msfile + '.dspec.npz'
        if os.path.exists(specfile):
            os.system('rm -rf ' + specfile)
        np.savez(specfile, spec=ospec, tim=tim, freq=freq,
                 timeran=timeran, spw=spw, bl=bl, uvrange=uvrange, pol=pols)
        if verbose:
            print('Median dynamic spectrum saved as: ' + specfile)

    return {'spec': ospec, 'tim': tim, 'freq': freq, 'timeran': timeran, 'spw': spw, 'bl': bl, 'uvrange': uvrange,
            'pol': pols}
Ejemplo n.º 2
0
def get_dspec(vis=None,
              savespec=True,
              specfile=None,
              bl='',
              uvrange='',
              field='',
              scan='',
              datacolumn='data',
              domedian=False,
              timeran=None,
              spw=None,
              timebin='0s',
              verbose=False):
    from split_cli import split_cli as split

    msfile = vis
    if not spw:
        spw = ''
    if not timeran:
        timeran = ''
    if not bl:
        bl = ''
    if domedian:
        if not uvrange:
            uvrange = '0.2~0.8km'
    else:
        uvrange = ''
    # Open the ms and plot dynamic spectrum
    if verbose:
        print 'Splitting selected data...'
    vis_spl = './tmpms.splitted'
    if os.path.exists(vis_spl):
        os.system('rm -rf ' + vis_spl)

    ms.open(msfile, nomodify=False)
    ms.split(outputms=vis_spl,
             whichcol=datacolumn,
             time=timeran,
             spw=spw,
             baseline=bl,
             field=field,
             scan=scan,
             uvrange=uvrange,
             timebin=timebin)
    ms.close()
    # split(vis=msfile, outputvis=vis_spl, timerange=timeran, antenna=bl, field=field, scan=scan, spw=spw,
    #       uvrange=uvrange, timebin=timebin, datacolumn=datacolumn)
    ms.open(vis_spl, nomodify=False)
    if verbose:
        print 'Regridding into a single spectral window...'
        # print 'Reading data spw by spw'
    ms.cvel(outframe='LSRK', mode='frequency', interp='nearest')
    ms.selectinit(datadescid=0, reset=True)
    data = ms.getdata(['amplitude', 'time', 'axis_info'], ifraxis=True)
    ms.close()
    os.system('rm -rf ' + vis_spl)
    specamp = data['amplitude']
    (npol, nfreq, nbl, ntim) = specamp.shape
    if verbose:
        print 'npol, nfreq, nbl, ntime:', data['amplitude'].shape
    spec = np.swapaxes(specamp, 2, 1)
    freq = data['axis_info']['freq_axis']['chan_freq'].reshape(nfreq)
    tim = data['time']

    if domedian:
        if verbose:
            print('doing median of all the baselines')
        # mask zero values before median
        spec_masked = np.ma.masked_where(spec < 1e-9, spec)
        spec_med = np.ma.filled(np.ma.median(spec_masked, axis=1),
                                fill_value=0.)
        nbl = 1
        ospec = spec_med.reshape((npol, nbl, nfreq, ntim))
    else:
        ospec = spec
    # Save the dynamic spectral data
    if savespec:
        if not specfile:
            specfile = msfile + '.dspec.npz'
        if os.path.exists(specfile):
            os.system('rm -rf ' + specfile)
        np.savez(specfile,
                 spec=ospec,
                 tim=tim,
                 freq=freq,
                 timeran=timeran,
                 spw=spw,
                 bl=bl,
                 uvrange=uvrange)
        if verbose:
            print 'Median dynamic spectrum saved as: ' + specfile

    return {
        'spec': ospec,
        'tim': tim,
        'freq': freq,
        'timeran': timeran,
        'spw': spw,
        'bl': bl,
        'uvrange': uvrange
    }
Ejemplo n.º 3
0
def get_dspec(vis=None, savespec=True, specfile=None, bl='', uvrange='', field='', scan='', datacolumn='data',
              domedian=False, timeran=None, spw=None, timebin='0s', regridfreq=False, fillnan=None, verbose=False,
              usetbtool=False):
    # from split_cli import split_cli as split
    if vis.endswith('/'):
        vis = vis[:-1]
    msfile = vis
    if not spw:
        spw = ''
    if not timeran:
        timeran = ''
    if domedian:
        if not uvrange:
            uvrange = '0.2~0.8km'
        # bl = ''
    else:
        uvrange = ''
    if not bl:
        bl = ''
    else:
        uvrange = ''
    # Open the ms and plot dynamic spectrum
    if verbose:
        print('Splitting selected data...')

    if usetbtool:
        try:
            tb.open(vis + '/POLARIZATION')
            corrtype = tb.getcell('CORR_TYPE', 0)
            pols = [stokesenum[p] for p in corrtype]
            tb.close()
        except:
            pols = []

        antmask = []
        if uvrange is not '' or bl is not '':
            ms.open(vis)
            ms.selectinit(datadescid=0)
            mdata = ms.metadata()
            antlist = mdata.antennaids()
            mdata.done()
            staql = {'uvdist': uvrange, 'baseline': bl, 'spw': spw, 'field': field, 'scan': scan, 'timerange': timeran}
            ### todo the selection only works for uvrange and bl. To make the selection of other items works,
            ## I need to make mask for other items.
            a = ms.msselect(staql)
            mdata = ms.metadata()
            baselines = mdata.baselines()
            for lidx, l in enumerate(antlist):
                antmask.append(baselines[l][antlist[lidx:]])
            antmask = np.hstack(antmask)
            mdata.done()
            ms.close()

        tb.open(vis)
        spwtb = tbtool()
        spwtb.open(vis + '/SPECTRAL_WINDOW')
        ptb = tbtool()
        ptb.open(vis + '/POLARIZATION')

        ms.open(vis)
        spwlist = []
        mdata = ms.metadata()
        nspw = mdata.nspw()
        nbl = mdata.nbaselines() + mdata.nantennas()
        nscans = mdata.nscans()
        spw_nfrq = []  # List of number of frequencies in each spw
        for i in range(nspw):
            spw_nfrq.append(mdata.nchan(i))
        spw_nfrq = np.array(spw_nfrq)
        nf = np.sum(spw_nfrq)
        smry = mdata.summary()
        scan_ntimes = []  # List of number of times in each scan
        for iscan in range(nscans):
            scan_ntimes.append(
                smry['observationID=0']['arrayID=0']['scan=' + str(iscan)]['fieldID=0']['nrows'] / nspw / nbl)
        scan_ntimes = np.array(scan_ntimes)
        scan_ntimes_integer = scan_ntimes.astype(np.int)
        if len(np.where(scan_ntimes % scan_ntimes_integer != 0)[0]) != 0:
            # if True:
            scan_ntimes = []  # List of number of times in each scan
            for iscan in range(nscans):
                scan_ntimes.append(
                    len(smry['observationID=0']['arrayID=0']['scan=' + str(iscan)]['fieldID=0'].keys()) - 6)
            scan_ntimes = np.array(scan_ntimes)
        else:
            scan_ntimes = scan_ntimes_integer

        nt = np.sum(scan_ntimes)
        times = tb.getcol('TIME')
        if times[nbl] - times[0] != 0:
            # This is frequency/scan sort order
            order = 'f'
        elif times[nbl * nspw - 1] - times[0] != 0:
            # This is time sort order
            order = 't'
        npol = ptb.getcol('NUM_CORR', 0, 1)[0]
        ptb.close()
        freq = np.zeros(nf, float)
        times = np.zeros(nt, float)
        if order == 't':
            specamp = np.zeros((npol, nf, nbl, nt), np.complex)
            for j in range(nt):
                fptr = 0
                # Loop over spw
                for i in range(nspw):
                    # Get channel frequencies for this spw (annoyingly comes out as shape (nf, 1)
                    cfrq = spwtb.getcol('CHAN_FREQ', i, 1)[:, 0]
                    if j == 0:
                        # Only need this the first time through
                        spwlist += [i] * len(cfrq)
                    if i == 0:
                        times[j] = tb.getcol('TIME', nbl * (i + nspw * j), 1)  # Get the time
                    spec_ = tb.getcol('DATA', nbl * (i + nspw * j), nbl)  # Get complex data for this spw
                    flag = tb.getcol('FLAG', nbl * (i + nspw * j), nbl)  # Get flags for this spw
                    nfrq = len(cfrq)
                    # Apply flags
                    if type(fillnan) in [int, float]:
                        spec_[flag] = float(fillnan)
                    else:
                        spec_[flag] = 0.0
                    # Insert data for this spw into larger array
                    specamp[:, fptr:fptr + nfrq, :, j] = spec_
                    freq[fptr:fptr + nfrq] = cfrq
                    fptr += nfrq
        else:
            specf = np.zeros((npol, nf, nt, nbl), np.complex)  # Array indexes are swapped
            iptr = 0
            for j in range(nscans):
                # Loop over scans
                for i in range(nspw):
                    # Loop over spectral windows
                    s = scan_ntimes[j]
                    f = spw_nfrq[i]
                    s1 = np.sum(scan_ntimes[:j])  # Start time index
                    s2 = np.sum(scan_ntimes[:j + 1])  # End time index
                    f1 = np.sum(spw_nfrq[:i])  # Start freq index
                    f2 = np.sum(spw_nfrq[:i + 1])  # End freq index
                    spec_ = tb.getcol('DATA', iptr, nbl * s)
                    flag = tb.getcol('FLAG', iptr, nbl * s)
                    if j == 0:
                        cfrq = spwtb.getcol('CHAN_FREQ', i, 1)[:, 0]
                        freq[f1:f2] = cfrq
                        spwlist += [i] * len(cfrq)
                    times[s1:s2] = tb.getcol('TIME', iptr, nbl * s).reshape(s, nbl)[:, 0]  # Get the times
                    iptr += nbl * s
                    # Apply flags
                    if type(fillnan) in [int, float]:
                        spec_[flag] = float(fillnan)
                    else:
                        spec_[flag] = 0.0
                    # Insert data for this spw into larger array
                    specf[:, f1:f2, s1:s2] = spec_.reshape(npol, f, s, nbl)
                    # Swap the array indexes back to the desired order
            specamp = np.swapaxes(specf, 2, 3)
        tb.close()
        spwtb.close()
        ms.close()
        if len(antmask) > 0:
            specamp = specamp[:, :, np.where(antmask)[0], :]
        (npol, nfreq, nbl, ntim) = specamp.shape
        tim = times
    else:
        # Open the ms and plot dynamic spectrum
        if verbose:
            print('Splitting selected data...')
        vis_spl = './tmpms.splitted'
        if os.path.exists(vis_spl):
            os.system('rm -rf ' + vis_spl)

        # split(vis=msfile, outputvis=vis_spl, timerange=timeran, antenna=bl, field=field, scan=scan, spw=spw,
        #       uvrange=uvrange, timebin=timebin, datacolumn=datacolumn)

        try:
            from split_cli import split_cli as split
            split(vis=msfile, outputvis=vis_spl, datacolumn=datacolumn, timerange=timeran, spw=spw, antenna=bl,
                  field=field,
                  scan=scan, uvrange=uvrange, timebin=timebin)
        except:
            ms.open(msfile, nomodify=True)
            ms.split(outputms=vis_spl, whichcol=datacolumn, time=timeran, spw=spw, baseline=bl, field=field, scan=scan,
                     uvrange=uvrange, timebin=timebin)
            ms.close()

        if verbose:
            print('Regridding into a single spectral window...')
            # print('Reading data spw by spw')

        try:
            tb.open(vis_spl + '/POLARIZATION')
            corrtype = tb.getcell('CORR_TYPE', 0)
            pols = [stokesenum[p] for p in corrtype]
            tb.close()
        except:
            pols = []

        if regridfreq:
            ms.open(vis_spl, nomodify=False)
            ms.cvel(outframe='LSRK', mode='frequency', interp='nearest')
            ms.selectinit(datadescid=0, reset=True)
            data = ms.getdata(['amplitude', 'time', 'axis_info'], ifraxis=True)
            specamp = data['amplitude']
            freq = data['axis_info']['freq_axis']['chan_freq']
        else:
            ms.open(vis_spl)
            ms.selectinit(datadescid=0, reset=True)
            spwinfo = ms.getspectralwindowinfo()
            specamp = []
            freq = []
            time = []
            for descid in range(len(spwinfo.keys())):
                ms.selectinit(datadescid=0, reset=True)
                ms.selectinit(datadescid=descid)
                data = ms.getdata(['amplitude', 'time', 'axis_info'], ifraxis=True)
                specamp_ = data['amplitude']
                freq_ = data['axis_info']['freq_axis']['chan_freq']
                time_ = data['time']
                if fillnan is not None:
                    flag_ = ms.getdata(['flag', 'time', 'axis_info'], ifraxis=True)['flag']
                    if type(fillnan) in [int, float, long]:
                        specamp_[flag_] = float(fillnan)
                    else:
                        specamp_[flag_] = 0.0
                specamp.append(specamp_)
                freq.append(freq_)
                time.append(time_)
            specamp = np.concatenate(specamp, axis=1)
            freq = np.concatenate(freq, axis=0)
            ms.selectinit(datadescid=0, reset=True)
        ms.close()
        os.system('rm -rf ' + vis_spl)
        (npol, nfreq, nbl, ntim) = specamp.shape
        freq = freq.reshape(nfreq)

        tim = data['time']

    if verbose:
        print('npol, nfreq, nbl, ntime:', (npol, nfreq, nbl, ntim))
    spec = np.swapaxes(specamp, 2, 1)

    if domedian:
        if verbose:
            print('doing median of all the baselines')
        # mask zero values before median
        # spec_masked = np.ma.masked_where(spec < 1e-9, spec)
        # spec_masked2 = np.ma.masked_invalid(spec)
        # spec_masked = np.ma.masked_array(spec, mask=np.logical_or(spec_masked.mask, spec_masked2.mask))
        # spec_med = np.ma.filled(np.ma.median(spec_masked, axis=1), fill_value=0.)
        spec = np.abs(spec)
        spec_med = np.nanmedian(spec, axis=1)
        nbl = 1
        ospec = spec_med.reshape((npol, nbl, nfreq, ntim))
    else:
        ospec = spec
    # Save the dynamic spectral data
    if savespec:
        if not specfile:
            specfile = msfile + '.dspec.npz'
        if os.path.exists(specfile):
            os.system('rm -rf ' + specfile)
        np.savez(specfile, spec=ospec, tim=tim, freq=freq,
                 timeran=timeran, spw=spw, bl=bl, uvrange=uvrange, pol=pols)
        if verbose:
            print('Median dynamic spectrum saved as: ' + specfile)

    return {'spec': ospec, 'tim': tim, 'freq': freq, 'timeran': timeran, 'spw': spw, 'bl': bl, 'uvrange': uvrange,
            'pol': pols}