Example #1
0
def buildIndex(inputfilename,difxfile,antList):

	# Get metadata
	print ('Checking and indexing %s' % (inputfilename))
	(numfreqs, freqs) = parseDiFX.get_freqtable_info(inputfilename)
	(numtelescopes, telescopes) = parseDiFX.get_telescopetable_info(inputfilename)
	(numdatastreams, datastreams) = parseDiFX.get_datastreamtable_info(inputfilename)
	(numbaselines, baselines) = parseDiFX.get_baselinetable_info(inputfilename)
	if numfreqs == 0 or numtelescopes == 0 or numdatastreams == 0 or numbaselines == 0:
		print ("Couldn't parse input file " + inputfilename + " correctly")
		sys.exit(-1)

	# Read visibility records one by one
	indices={}
	peak_sec=0
	while True:
		offset = difxfile.tell()
		(vishdr,binhdr) = getVisibilityHeader(difxfile)
		if len(vishdr) <= 0:
			break

		# Visibility properties
		baseline = vishdr[0]
		seconds = vishdr[2]
		freqindex = vishdr[5]
		polpair = vishdr[6]
		uvw = vishdr[9:12]
		freq = freqs[freqindex]

		# Antenna order as in difx2mark4: ref=ant1="256*nr", rem=ant2="nr%256"
		ant2 = baseline % 256
		ant1 = (baseline-ant2)/256
		ant1name = telescopes[ant1-1].name.upper()
		ant2name = telescopes[ant2-1].name.upper()

		# Number of channels in this baseband
		nchan = freqs[freqindex].numchan / freqs[freqindex].specavg

		# Info string
		sband = 'U'
		if freq.lsb:
			sband = 'L'
		#key = ant1name + '-' + ant2name + '/' + polpair + '/' + str(freqindex) + '@' + str(freq.freq) + sband + ':' + str(seconds)
		key = ant1name + '-' + ant2name + '/' + str(freq.freq) + sband + '_' + polpair + ':' + str(seconds) + ':' + str(nchan)

		# Read the entire visibility data from disk
		rawvis = difxfile.read(8*nchan)
		if len(rawvis) < 8*nchan:
			break

		# Indexing
		if (ant1name in antList) or (ant2name in antList):
			indices[key] = offset
			#print key, uvw
		if seconds > peak_sec:
			peak_sec = seconds
			print ('.'),
			sys.stdout.flush()

	print ('\nIndexing complete, found %d records' % (len(indices)))
	return indices
def stitchVisibilityfile(basename, cfg):

    # Get settings
    target_bw = cfg['target_bw']
    target_nchan = cfg['target_nchan']
    stitch_basefreqs = cfg['stitch_basefreqs']
    stitch_endfreqs = cfg['stitch_endfreqs']
    stitch_antennas = cfg['stitch_antennas']
    stitch_oversamplenum = cfg['stitch_oversamplenum']
    stitch_oversampledenom = cfg['stitch_oversampledenom']
    blank_viz = numpy.fromstring('\0' * 8 * target_nchan, dtype='complex64')

    # All polarisation pairs 'RR', 'RL', ... to 'YL'
    if cfg['stitch_nstokes'] > 2:
        pols_list = ['R', 'L', 'X', 'Y']
        polpairs, tmp = setCrossProd(pols_list, pols_list)
    else:
        polpairs = ['RR', 'LL', 'XX', 'YY']  # parallel-hand only
        polpairs += ['RX', 'LY', 'XR', 'YL']  # potential parallel-hand
        polpairs += ['LX', 'RY', 'XL', 'YR'
                     ]  # potential parallel-hand (depends on polconvert.py)

    # Extract meta-infos from the DiFX .INPUT file
    if basename.endswith(('.difx', 'input')):
        basename = basename[:basename.rfind('.')]
    basename_pathless = basename
    basedir = os.getcwd()
    if basename.rfind('/') >= 0:
        basename_pathless = basename[(basename.rfind('/') + 1):]
        basedir = basename[:basename.rfind('/')]
    inputfile = basename + '.input'
    inputfile_cfg = parseDiFX.get_common_settings(inputfile)
    (numfreqs, freqs) = parseDiFX.get_freqtable_info(inputfile)
    (numtelescopes, telescopes) = parseDiFX.get_telescopetable_info(inputfile)
    (numdatastreams,
     datastreams) = parseDiFX.get_datastreamtable_info(inputfile)
    (numbaselines, baselines) = parseDiFX.get_baselinetable_info(inputfile)
    if 'difxfile' not in inputfile_cfg:
        parser.error("Couldn't parse COMMON SETTINGS of input file " +
                     inputfile + " correctly")
    if numfreqs == 0 or numtelescopes == 0 or numdatastreams == 0 or numbaselines == 0:
        parser.error("Couldn't parse input file " + inputfile + " correctly")

    # Internal frequency tables and remappings from old to new freq indices
    out_freqs = {}  # dict out['id']=Freq()
    in_rec_freqs = 0  # num of distinct recorded freqs
    n_max_freqs = 256
    freq_remaps = [-1] * n_max_freqs
    freq_remaps_isNew = [False] * n_max_freqs

    # Collect all recorded freqs listed in DATASTREAMS
    for d in datastreams:
        for fqidx in d.recfreqindex:
            i = findFreqObj(out_freqs, freqs[fqidx])
            if (i == None):
                id = inventNextKey(out_freqs)
                freq_remaps[fqidx] = id
                out_freqs[id] = copy.deepcopy(freqs[fqidx])
                in_rec_freqs += 1
                print(("Keeping recorded frequency  : index %2d/%2d : %s" %
                       (fqidx, id, freqs[fqidx].str())))

    # Collect all zoom bands listed in DATASTREAMS that already match the target bandwidth
    for d in datastreams:
        for fqidx in d.zoomfreqindex:
            zf = freqs[fqidx]
            if (zf.bandwidth != target_bw) or (
                (zf.numchan // zf.specavg) != target_nchan):
                # print ("Skipping zoom frequency     : index %d : %s" % (fqidx,zf.str()))
                continue
            i = findFreqObj(out_freqs, zf)
            if (i == None):
                id = inventNextKey(out_freqs)
                freq_remaps[fqidx] = id
                out_freqs[id] = copy.deepcopy(zf)
                print(("Keeping zoom frequency      : index %2d/%2d : %s" %
                       (fqidx, id, zf.str())))

    # Check stitch config: invent new zoom bands if necessary
    stitch_out_ids = []
    for fsky in stitch_basefreqs:
        match_existing = [(out_freqs[key].freq == fsky
                           and not out_freqs[key].lsb)
                          for key in out_freqs.keys()]
        exists = any(match_existing)
        if exists:
            # Re-use an existing matching zoom frequency
            i = match_existing.index(True)
            id = out_freqs.keys()[i]
            zf = out_freqs[id]
            stitch_out_ids.append(id)
            old_id = freq_remaps.index(id)
            print(("Re-using zoom for stitch    : index %2d/%2d : %s" %
                   (old_id, id, zf.str())))
        else:
            # Invent new zoom
            zf = parseDiFX.Freq()
            zf.freq = fsky
            zf.bandwidth = target_bw
            zf.numchan = target_nchan
            zf.specavg = 1
            zf.lsb = False  # TODO
            id = inventNextKey(out_freqs)
            out_freqs[id] = copy.deepcopy(zf)
            stitch_out_ids.append(id)
            print(("Creating new zoom frequency : index --/%2d : %s" %
                   (id, zf.str())))

    stitch_ids = [n for n in range(len(stitch_out_ids))]

    # Also map each to-be-stitched-multi-zoom into respective new single post-stitch zoom
    #for fi in range(in_rec_freqs,len(freqs)):
    for fi in range(0, len(freqs)):
        if (freqs[fi].bandwidth
                == target_bw) and (freqs[fi].numchan // freqs[fi].specavg
                                   == target_nchan):
            # Retain one-to-one map for existing matching zoom freqs
            freq_remaps_isNew[fi] = False
            continue
        lowedgefreq = freqs[fi].freq
        if freqs[fi].lsb:
            lowedgefreq -= freqs[fi].bandwidth
        stid = getGlueIndex(lowedgefreq, freqs[fi].bandwidth, cfg)
        #print "For freqindex", fi, ", stid is", stid, ", and stitch_out_ids for this freqid is", stitch_out_ids[stid]
        if (stid >= 0):
            #if freqs[fi].lsb:
            #        # Ignore LSB. All zoom outputs should be USB.
            #        print ("Ignore LSB %s" % (freqs[fi].str()))
            #        continue
            # Add to some many-to-one map of invented zoom freqs
            freq_remaps[fi] = stitch_out_ids[stid]
            freq_remaps_isNew[fi] = True
            print((
                "Map zoom %s to stitched single %12.6f--%12.6f : in fq#%2d -> stitch#%d -> out fq#%2d"
                % (freqs[fi].str(), stitch_basefreqs[stid],
                   stitch_endfreqs[stid], fi, stid, stitch_out_ids[stid])))

    # Read the DiFX .difx/DIFX_* file
    #glob_pattern = basename + '.difx/DIFX_*.s*.b*'
    glob_pattern = inputfile_cfg['difxfile'] + '/DIFX_*.s*.b*'
    difxfileslist = glob.glob(glob_pattern)
    if len(difxfileslist) <= 0:
        print(('Error: no visibility data file found in %s!' % (glob_pattern)))
        return
    #difxfilename = difxfileslist[0]
    for difxfilename in difxfileslist:
        global vis_hashtable
        vis_hashtable = {}
        global vis_hashtable_cleanupSec
        vis_hashtable_cleanupSec = 0

        difxfilename_pathless = difxfilename
        if difxfilename.rfind('/') >= 0:
            difxfilename_pathless = difxfilename[(difxfilename.rfind('/') +
                                                  1):]
        difxfile = open(difxfilename, 'rb')
        difxoutdir = basedir + '/' + basename_pathless + 'D2D.difx'
        difxoutname = difxoutdir + '/' + difxfilename_pathless
        try:
            os.mkdir(difxoutdir)
        except:
            pass
        if os.path.exists(difxoutname):
            print(
                ('Warning: %s already exists, skipping the processing of %s!' %
                 (difxoutname, difxfilename)))
            return
        difxout = open(difxoutname, 'wb')

        # Pull out antenna indices based on telescope names
        telNames = [t.name for t in telescopes]
        stAntIDs = []
        if len(stitch_antennas) == 1 and stitch_antennas[0] == "*":
            for i in range(len(telescopes)):
                stAntIDs.append(i + 1)
        else:
            for sa in stitch_antennas:
                id = [
                    n for n in range(len(telescopes))
                    if telescopes[n].name == sa
                ]
                if len(id) != 1:
                    print((
                        "Warning: got %d hits for telescope %s in .difx file telescope list of %s!"
                        % (len(id), sa, str(telNames))))
                    continue
                id = id[0] + 1  # index (0..N-1) into antenna nr (1..N)
                stAntIDs.append(id)

        # Affected baselines; data of these are stored in stitching memory, other baselines need no stitching
        baseline_list = []
        for sid in stAntIDs:
            # Cross
            for t in range(numtelescopes):
                tid = t + 1  # index (0..N-1) into antenna nr (1..N)
                if (tid != sid):
                    baseline_list.append(tid + sid * 256)
                    baseline_list.append(sid + tid * 256)
            # Auto
            baseline_list.append(sid + sid * 256)

        # Dictionary of per-polarizationpair working buffers into which spectra are stitched
        stitch_workbufs = {
            bline: {
                polkey:
                {bandkey: numpy.copy(blank_viz)
                 for bandkey in stitch_ids}
                for polkey in polpairs
            }
            for bline in baseline_list
        }
        stitch_chcounts = {
            bline: {
                polkey: {bandkey: 0
                         for bandkey in stitch_ids}
                for polkey in polpairs
            }
            for bline in baseline_list
        }
        stitch_freqbws = {
            bline: {
                polkey: {bandkey: []
                         for bandkey in stitch_ids}
                for polkey in polpairs
            }
            for bline in baseline_list
        }
        stitch_timestamps = {
            bline: {
                polkey: {bandkey: -1
                         for bandkey in stitch_ids}
                for polkey in polpairs
            }
            for bline in baseline_list
        }
        nstitched = 0
        ncopied = 0
        nskipped = 0
        nincomplete = 0

        # Parse each visibility entry
        seconds_prev = -1
        seconds_first = 0
        seconds_report_prev = 0
        while True:

            (vishdr, binhdr) = getVisibilityHeader(difxfile)
            if len(vishdr) <= 0:
                break

            # Visibility properties
            baseline = vishdr[0]
            mjd = vishdr[1]
            seconds = vishdr[2]
            freqindex = vishdr[5]
            polpair = vishdr[6]
            weight = vishdr[8]
            uvw = vishdr[9:12]
            # Antenna order as in difx2mark4: ref=ant1="256*nr", rem=ant2="nr%256"
            ant2 = baseline % 256
            ant1 = (baseline - ant2) // 256
            ant1name = telescopes[ant1 - 1].name
            ant2name = telescopes[ant2 - 1].name
            T = mjd + seconds / 86400.0
            baselinestr = '%s-%s' % (ant1name, ant2name)

            # Number of channels in this baseband
            nchan = freqs[freqindex].numchan // freqs[freqindex].specavg
            fsky = freqs[freqindex].freq
            bw = freqs[freqindex].bandwidth
            if freqs[freqindex].lsb:
                fsky -= bw

            # Read the entire visibility data from disk
            rawvis = difxfile.read(8 * nchan)
            if len(rawvis) != 8 * nchan:
                print('Short read! Stopping.')
                break

            if (seconds != seconds_prev):
                seconds_prev = seconds
                if cfg['verbose']:
                    print((('\n---- %d %12.7f ' + '-' * 115) % (mjd, seconds)))
                elif (seconds - seconds_report_prev) > 1.0:
                    if seconds_first <= 0:
                        seconds_first = mjd * 24 * 60 * 60 + seconds
                    dTstart = mjd * 24 * 60 * 60 + seconds - seconds_first
                    seconds_report_prev = seconds
                    print(("at %d %12.7f, %.3f seconds from start" %
                           (mjd, seconds, dTstart)))
                    print("\033[F\033[F")

            # Remap the frequency reference
            out_freqindex = freqindex
            if freq_remaps[freqindex] >= 0:
                # print ('Remapping input freq index %d to old/new output freq %d' % (freqindex,freq_remaps[freqindex]))
                out_freqindex = freq_remaps[freqindex]
                vishdr[5] = out_freqindex
                binhdr = parseDiFX.make_output_header_v1(vishdr)
            stid = getGlueIndex(fsky, bw, cfg)

            # Info string
            vis_info = '%s-%s/%d/%d(%d):sf<%d>:%.7f/%s  mjd:%12.8f nchan:%4d bw:%.7f uvw=%s' % (
                ant1name, ant2name, baseline, out_freqindex, freqindex, stid,
                fsky, polpair, T, nchan, bw, str(uvw))

            # Write out visibility to output file
            if baseline not in baseline_list and not stitch_antennas == "*":

                # Baseline for which no freq stitching was requested

                # Keep data if bandwidth matches
                if (bw == target_bw) and (nchan == target_nchan):

                    if vis_already_written(mjd, seconds, baseline,
                                           out_freqindex, polpair):
                        if cfg['verbose']:
                            print(('(copy): %s' % (vis_info)))
                        continue
                    if cfg['verbose']:
                        print(('copy  : %s' % (vis_info)))

                    assert (freq_remaps[freqindex] >= 0)
                    difxout.write(binhdr)
                    if cfg['target_chavg'] > 1:
                        rawvis = spectralAvgRaw(rawvis, cfg['target_chavg'])
                    difxout.write(rawvis)
                    ncopied += 1

                # Discard data if bandwidth mismatch. Happens e.g. for zoom bands on the non-stitch baselines,
                # can ignore these because covered by a different and complete zoomband that doesn't need stitching
                else:
                    nskipped += 1
                    if cfg['verbose']:
                        print(('skip  : %s' % (vis_info)))
                        pass
            else:

                # Baseline where freq stitching was requested
                resetStitch = False

                if (bw == target_bw):

                    # Visibility already at target bandwidth, just copy the data

                    if vis_already_written(mjd, seconds, baseline,
                                           out_freqindex, polpair):
                        if cfg['verbose']:
                            print(('(take): %s' % (vis_info)))
                        continue
                    if cfg['verbose']:
                        print(('take  : %s' % (vis_info)))

                    assert (freq_remaps[freqindex] >= 0)
                    difxout.write(binhdr)
                    if cfg['target_chavg'] > 1:
                        rawvis = spectralAvgRaw(rawvis, cfg['target_chavg'])
                    difxout.write(rawvis)
                    ncopied += 1

                    # Signal for further below: clear old working data, set new timestamp
                    resetStitch = True

                else:

                    # Need to stitch narrowband visibility into wideband context

                    # Ignore bands not requested to be stitched
                    if (stid < 0):
                        if cfg['verbose']:
                            print(('ignore: %s (not in stitch freqs list)' %
                                   (vis_info)))
                        continue

                    # Get the complex data
                    visdata = numpy.fromstring(rawvis, dtype='complex64')

                    # Do the oversampling cutout if needed
                    if stitch_oversamplenum > 1:
                        oversampleoffsetbw = (
                            bw - bw * (float(stitch_oversampledenom) /
                                       float(stitch_oversamplenum))) / 2.0
                        oversamplelengthchannels = int(
                            (float(stitch_oversampledenom) /
                             float(stitch_oversamplenum)) * nchan + 0.5)
                        oversampleoffsetchannels = (
                            nchan - oversamplelengthchannels) // 2
                        visdata = visdata[
                            oversampleoffsetchannels:oversampleoffsetchannels +
                            oversamplelengthchannels]
                        fsky += oversampleoffsetbw
                        bw -= 2 * oversampleoffsetbw
                        nchan = oversamplelengthchannels

                    # Determine the region into which this visibility data would be inserted
                    choffset = int(
                        ((fsky - stitch_basefreqs[stid]) / target_bw) *
                        target_nchan)
                    if (choffset + nchan) > target_nchan:
                        # Sometimes a wide recorded band falls into the stitching frequency region,
                        # but then may not fit fully into the stitching freq region.
                        # Options are:
                        # 1) discard and hope the corresponding zoom band data will appear soon
                        #if cfg['verbose']:
                        #        print ('ignore: %s (too wide, maybe a recorded instead of zoom freq)' % (vis_info))
                        #continue
                        # 2) use the maximum possible amount of the data
                        nchan_new = target_nchan - choffset
                        if (nchan_new < 1):
                            print((
                                'Unexpected, after cropping a too-wide band ended up with %d remaining channels'
                                % (nchan_new)))
                            continue
                        visdata = visdata[:nchan_new]
                        bw = bw * nchan_new / nchan
                        nchan = nchan_new
                        if cfg['verbose']:
                            print(('crop  : %s to %d chan' %
                                   (vis_info, nchan_new)))

                    if choffset < 0:
                        print("This shouldn't happen!")
                        continue

                    # Keep track of timestamp
                    if (stitch_timestamps[baseline][polpair][stid] < 0):
                        stitch_timestamps[baseline][polpair][stid] = seconds

                    # Detect a jump in timestamp
                    if (seconds != stitch_timestamps[baseline][polpair][stid]):
                        # We assume the input visibilities are stored in time-increasing order,
                        # so a timestamp change indicates all data of previous time should've been fully read by now (no older data follows)
                        ncurr = stitch_chcounts[baseline][polpair][stid]
                        if (ncurr > 0):
                            T_old = mjd + stitch_timestamps[baseline][polpair][
                                stid] / 86400.0
                            nincomplete += 1
                            msg = '%s-%s/%d/%d:%.7f/%s mjd:%12.8f only %d of %d channels' % (
                                ant1name, ant2name, baseline, stid,
                                stitch_basefreqs[stid], polpair, T_old, ncurr,
                                target_nchan)
                            if cfg['verbose']:
                                print((
                                    'Warning: discarded an incomplete stitch: %s'
                                    % (msg)))
                        assert (ncurr < target_nchan)
                        stitch_timestamps[baseline][polpair][stid] = seconds
                        stitch_chcounts[baseline][polpair][stid] = 0
                        stitch_freqbws[baseline][polpair][stid] = []
                        stitch_workbufs[baseline][polpair][stid].fill(0)
                        #print (('---- %d %.7f ' + '-'*80 + '\n') % (mjd,seconds))

                    # Insert the new visibility data into the correct working buffer at the correct location
                    stitch_workbufs[baseline][polpair][stid][choffset:(
                        choffset + nchan)] = visdata
                    stitch_chcounts[baseline][polpair][stid] += nchan
                    stitch_freqbws[baseline][polpair][stid].append(bw)

                    #print (' glue bline %s %s band %d : ch %4d ... %4d : %d new chs : stitched %d ch total' % (baselinestr,polpair,stid,choffset,choffset+nchan-1,nchan,stitch_chcounts[baseline][polpair][stid]))
                    #print choffset, nchan, choffset+nchan, stitch_chcounts[baseline][polpair][stid], (100.0*stitch_chcounts[baseline][polpair][stid])/target_nchan

                    # Assembled the full bandwidth now?
                    if stitch_chcounts[baseline][polpair][stid] >= target_nchan:
                        TS = mjd + stitch_timestamps[baseline][polpair][
                            stid] / 86400.0
                        bwsum = numpy.sum(
                            stitch_freqbws[baseline][polpair][stid])
                        vis_info = '%s-%s/%d/%d(%d):sf<%d>:%.7f/%s  mjd:%12.8f nchan:%4d bw:%.7f uvw=%s' % (
                            ant1name, ant2name, baseline, out_freqindex,
                            freqindex, stid, fsky, polpair, TS,
                            stitch_chcounts[baseline][polpair][stid], bwsum,
                            str(uvw))

                        if vis_already_written(mjd, seconds, baseline,
                                               out_freqindex, polpair):
                            if cfg['verbose']:
                                print(('(stch): %s' % (vis_info)))
                            continue
                        if cfg['verbose']:
                            print(('stitch: %s' % (vis_info)))

                        # Double-check the bandwidth
                        if numpy.abs(target_bw - bwsum) > 100:
                            print((
                                'Warning: stitched bands gave %.6f MHz bandwidth, differs from target %.6f MHz!'
                                % (bwsum, target_bw)))

                        # Write header and assembled data
                        difxout.write(binhdr)
                        if cfg['target_chavg'] > 1:
                            vis = spectralAvgNumpy(
                                stitch_workbufs[baseline][polpair][stid],
                                cfg['target_chavg'])
                        else:
                            vis = stitch_workbufs[baseline][polpair][stid]
                        vis.tofile(difxout)
                        nstitched += 1

                        # Reset
                        resetStitch = True

                if resetStitch and (stid >= 0):
                    stitch_timestamps[baseline][polpair][stid] = seconds
                    stitch_chcounts[baseline][polpair][stid] = 0
                    stitch_freqbws[baseline][polpair][stid] = []
                    stitch_workbufs[baseline][polpair][stid].fill(0)

            #if nstitched > 100: break  # quick-exit for debug

    ## Generate new .input

    # New FREQ table
    new_freqs = []
    for of in out_freqs.keys():  # dict into list
        out_freqs[of].specavg = out_freqs[of].specavg * cfg['target_chavg']
        new_freqs.append(out_freqs[of])

    # New DATASTREAM table
    new_datastreams = []
    for ds in datastreams:
        newds = copy.deepcopy(ds)
        ds_specific_remaps = [freq_remaps[zfi] for zfi in ds.zoomfreqindex]

        # Determine number of pols in this datastream (for example ALMA has two single-pol datastreams)
        npol = len(list(set(newds.recbandpol)))

        # Retain all recorded freqs and bands
        # newds.recfreqindex = newds.recfreqindex
        # newds.recfreqpols = newds.recfreqpols
        # newds.recbandindex = newds.recbandindex
        # newds.recbandpol = newds.recbandpol
        newds.recfreqindex = [freq_remaps[rfi] for rfi in newds.recfreqindex]

        # Retain all bandwidth-matching zoom freqs
        newds.zoomfreqindex = [
            freq_remaps[zfi] for zfi in ds.zoomfreqindex
            if (freq_remaps[zfi] >= 0 and not freq_remaps_isNew[zfi])
        ]
        newds.zoomfreqpols = [
            ds.zoomfreqpols[ds.zoomfreqindex.index(zfi)]
            for zfi in ds.zoomfreqindex
            if (freq_remaps[zfi] >= 0 and not freq_remaps_isNew[zfi])
        ]

        # Add all stitched invented freqs
        # TODO: This is a quick hack: find ouy why the commented code wasn't working!
        #newds.zoomfreqindex += [nzfi for nzfi in stitch_out_ids if nzfi in ds_specific_remaps]
        #newds.zoomfreqpols += [npol for nzfi in stitch_out_ids if nzfi in ds_specific_remaps]
        newds.zoomfreqindex += [nzfi for nzfi in stitch_out_ids]
        newds.zoomfreqpols += [npol for nzfi in stitch_out_ids]

        newds.zoomfreqindex = list(
            set(newds.zoomfreqindex)
        )  # keep uniques only, TODO: should shorten zoomfreqpols list equally!

        # Translate freqs into bands
        newds.nrecband = npol * len(newds.recfreqindex)
        newds.nzoomband = npol * len(newds.zoomfreqindex)
        newds.zoombandindex = [int(n / npol) for n in range(newds.nzoomband)]
        #TODO: The following is necessary if a datastream had no zoom bands to begin with, but I'm not sure it will always work
        if len(ds.zoombandpol) == 0:
            ds.zoombandpol = newds.recbandpol[:npol]
        newds.zoombandpol = ds.zoombandpol[:npol] * (newds.nzoomband // npol)

        # Update the counts
        newds.nzoomfreq = len(newds.zoomfreqindex)
        new_datastreams.append(newds)

        if cfg['verbose']:
            print(("DS%d : rec freqs      : %s" %
                   (datastreams.index(ds), str(newds.recfreqindex))))
            print(("      rec freq pols  : %s" % (str(newds.recfreqpols))))
            print(("      rec bands      : %s" % (str(newds.recbandindex))))
            print(("      rec band pols  : %s" % (str(newds.recbandpol))))
            print(("      zoom freqs     : %s" % (str(newds.zoomfreqindex))))
            print(("      zoom freq pols : %s" % (str(newds.zoomfreqpols))))
            print(("      zoom bands     : %s" % (str(newds.zoombandindex))))
            print(("      zoom band pols : %s" % (str(newds.zoombandpol))))

    # New BASELINE table
    new_baselines = []
    remapped_freqs = [
        n for n in range(len(freq_remaps)) if freq_remaps[n] >= 0
    ]
    for b in baselines:

        newbl = copy.deepcopy(b)

        nds1 = new_datastreams[newbl.dsaindex]
        nds2 = new_datastreams[newbl.dsbindex]
        npol1 = len(list(set(nds1.recbandpol)))
        npol2 = len(list(set(nds2.recbandpol)))
        max_stokes = npol1 * npol2

        # Keep only freqs common to both datastreams, and that match the target bandwidth
        allfreqs1 = nds1.recfreqindex + nds1.zoomfreqindex
        allfreqs2 = nds2.recfreqindex + nds2.zoomfreqindex
        common_freqs = list(set(allfreqs1) & set(allfreqs2))
        common_freqs = [
            fnr for fnr in common_freqs
            if (out_freqs[fnr].bandwidth == target_bw)
        ]

        newbl.nfreq = len(common_freqs)
        newbl.dsabandindex = []
        newbl.dsbbandindex = []
        newbl.freqpols = []

        if cfg['verbose']:
            print(("Baseline DS%d x DS%d" % (newbl.dsaindex, newbl.dsbindex)))
            print(
                ("     stream %d freqs %s" % (newbl.dsaindex, str(allfreqs1))))
            print(
                ("     stream %d freqs %s" % (newbl.dsbindex, str(allfreqs2))))
            print(("     common freqs %s" % (str(common_freqs))))

        for f in common_freqs:
            npol = 2
            assert (f in allfreqs1)
            assert (f in allfreqs2)
            i1 = allfreqs1.index(f)
            i2 = allfreqs2.index(f)
            if max_stokes == 1:
                newbl.dsabandindex.append([npol1 * i1])
                newbl.dsbbandindex.append([npol2 * i2])
            elif max_stokes == 2:  # one of XL XR / YL YR / LL LR / RL RR
                if npol1 == 1:
                    newbl.dsabandindex.append([npol1 * i1, npol1 * i1])
                    newbl.dsbbandindex.append([npol2 * i2 + 0, npol2 * i2 + 1])
                else:
                    newbl.dsabandindex.append([npol1 * i1 + 0, npol1 * i1 + 1])
                    newbl.dsbbandindex.append([npol2 * i2, npol2 * i2])
            elif max_stokes == 4:
                if cfg['stitch_nstokes'] == 4:
                    newbl.dsabandindex.append([
                        npol1 * i1 + 0, npol1 * i1 + 0, npol1 * i1 + 1,
                        npol1 * i1 + 1
                    ])
                    newbl.dsbbandindex.append([
                        npol2 * i2 + 0, npol2 * i2 + 1, npol2 * i2 + 0,
                        npol2 * i2 + 1
                    ])
                elif cfg['stitch_nstokes'] == 2:
                    # one XX YY / XL YR / LL RR
                    newbl.dsabandindex.append([npol1 * i1 + 0, npol1 * i1 + 1])
                    newbl.dsbbandindex.append([npol2 * i2 + 0, npol2 * i2 + 1])
                else:
                    print(("Warning: unexpected nstokes of %d" %
                           (cfg['stitch_nstokes'])))
                    newbl.dsabandindex.append([npol1 * i1])
                    newbl.dsbbandindex.append([npol2 * i2])

            newbl.freqpols.append(len(newbl.dsabandindex[-1]))

            assert (len(newbl.dsabandindex[-1]) == len(newbl.dsbbandindex[-1]))
            if cfg['verbose']:
                print(("        now freq %d = (ds1 frq %d, ds2 frq %d)" %
                       (f, i1, 2)))
                print(("            num pols = (%d, %d)" % (npol1, npol2)))
                print(("            dsa[end] = %s" %
                       (str(newbl.dsabandindex[-1]))))
                print(("            dsb[end] = %s" %
                       (str(newbl.dsbbandindex[-1]))))
                print(("        has %d freqpols" % (newbl.freqpols[-1])))
                #print newbl.dsabandindex[-1], newbl.dsbbandindex[-1]

        new_baselines.append(newbl)

    # Read original .input without parsing
    fin = open(inputfile, "r")
    in_lines = fin.readlines()
    fin.close()

    # Replace OUTPUT FILENAME entry
    for l in in_lines:
        if l[:16] == "OUTPUT FILENAME:":
            i = in_lines.index(l)
            in_lines[i] = "%-20s%s\n" % ("OUTPUT FILENAME:", difxoutdir)

    # Generate new reference .input for later, re-use parts of original .input
    outputinputfile = basename_pathless + 'D2D.input'
    fout = open(outputinputfile, "w")
    idx = [i for i, elem in enumerate(in_lines) if "FREQ TABLE" in elem]
    for n in range(idx[0]):
        fout.write(in_lines[n])
    parseDiFX.put_freqtable_info(fout, new_freqs)
    idx1 = [i for i, elem in enumerate(in_lines) if "TELESCOPE TABLE" in elem]
    idx2 = [i for i, elem in enumerate(in_lines) if "DATASTREAM TABLE" in elem]
    for n in range(idx1[0], idx2[0]):
        fout.write(in_lines[n])
    parseDiFX.put_datastreamtable_info(fout, new_datastreams)
    parseDiFX.put_baselinetable_info(fout, new_baselines)
    idx = [i for i, elem in enumerate(in_lines) if "DATA TABLE" in elem]
    for trailing in in_lines[idx[0]:]:
        fout.write(trailing)
    fout.close()

    # Finished
    print('\nDone! Final statistics:')
    print(('    vis. copied        : %d' % (ncopied)))
    print(('    vis. skipped       : %d' % (nskipped)))
    print(('    vis. stitched      : %d' % (nstitched)))
    print(('    incomplete stitch  : %d' % (nincomplete)))
    print('\nOutput files:')
    print(('    visbility data     : %s' % (difxoutname)))
    print(('    changes for .input : %s' % (outputinputfile)))
    print(' ')
Example #3
0
def patchDiFX(basename_dst,basename_src,antList):

	# Derive file names
	(pathless_basename_dst, basename_dst) = getBasename(basename_dst)
	(pathless_basename_src, basename_src) = getBasename(basename_src)

	# Open DiFX files
	difxdstfile = glob.glob(basename_dst + '.difx/DIFX_*.s*.b*')[0]
	fdst = open(difxdstfile, 'r')
	difxsrcfile = glob.glob(basename_src + '.difx/DIFX_*.s*.b*')[0]
	fsrc = open(difxsrcfile, 'r')
	difxoutdir = pathless_basename_dst + '_antreplaced.difx'
	difxoutfile = pathless_basename_dst + '_antreplaced.difx/' + difxdstfile[difxdstfile.rfind('/')+1:]
	try:
		os.mkdir(difxoutdir)
	except:
		pass
	fout = open(difxoutfile, 'w')		
	print ('Destination  : %s' % (difxdstfile))
	print ('Source       : %s' % (difxsrcfile))
	print ('Merged into  : %s' % (difxoutfile))

	# Make index of visibilities in source dataset
	src_index = buildIndex(basename_src+'.input',fsrc,antList)
	fsrc.seek(0)

	# Extract meta-infos from the DiFX .INPUT file
	inputfilename = basename_dst + '.input'
	(numfreqs, freqs) = parseDiFX.get_freqtable_info(inputfilename)
	(numtelescopes, telescopes) = parseDiFX.get_telescopetable_info(inputfilename)
	(numdatastreams, datastreams) = parseDiFX.get_datastreamtable_info(inputfilename)
	(numbaselines, baselines) = parseDiFX.get_baselinetable_info(inputfilename)
	if numfreqs == 0 or numtelescopes == 0 or numdatastreams == 0 or numbaselines == 0:
		print ("Couldn't parse input file " + inputfilename + " correctly")
		sys.exit(-1)
	newinputfile = pathless_basename_dst + '_antreplaced.input'
	try:
		shutil.copyfile(inputfilename, newinputfile)
	except Exception as e:
		print ('Error copying %s to %s : %s' % (inputfilename,newinputfile,str(e)))

	# Copy/replace visibility data
	Ncopied = 0
	Nreplaced = 0
	Nnotfound = 0
	Vnotfound = []
	peak_sec = 0
	print ('Replacing visibilities to %s...' % (str(antList)))
	while True:

		(vishdr,binhdr) = getVisibilityHeader(fdst)
		if len(vishdr) <= 0:
			break

		# Visibility properties
		baseline = vishdr[0]
		seconds = vishdr[2]
		freqindex = vishdr[5]
		polpair = vishdr[6]
		freq = freqs[freqindex]

		# Antenna order as in difx2mark4: ref=ant1="256*nr", rem=ant2="nr%256"
		ant2 = baseline % 256
		ant1 = (baseline-ant2)/256
		ant1name = telescopes[ant1-1].name.upper()
		ant2name = telescopes[ant2-1].name.upper()

		# Number of channels in this baseband
		nchan = freqs[freqindex].numchan / freqs[freqindex].specavg

		# Info string
		sband = 'U'
		if freq.lsb:
			sband = 'L'
		#key = ant1name + '-' + ant2name + '/' + polpair + '/' + str(freqindex) + '@' + str(freq.freq) + sband + ':' + str(seconds)
		key = ant1name + '-' + ant2name + '/' + str(freq.freq) + sband + '_' + polpair + ':' + str(seconds) + ':' + str(nchan)

		if seconds > peak_sec:
			peak_sec = seconds
			print ('.'),
			sys.stdout.flush()

		# Read the entire visibility data from disk
		rawvis = fdst.read(8*nchan)
		if len(rawvis) < 8*nchan:
			continue

		# Replace with data from source if key matches
		# note: vis record header from 'src' is thrown away; use Freq ID, Ant ID etc from 'dst' file since numerical IDs can mismatch between src/dst!
		if (ant1name in antList) or (ant2name in antList):
			if key in src_index:
				offset = src_index[key]
				fsrc.seek(offset)
				(vishdr2,binhdr2) = getVisibilityHeader(fsrc)
				rawvis = fsrc.read(8*nchan)
				if len(rawvis) < 8*nchan:
					print ('Error: could not get %s with %d nchannels from source file!' % (key,nchan))
					continue
				Nreplaced += 1
			else:
				Nnotfound += 1
				brief_key = ant1name + '-' + ant2name + '/' + str(freq.freq) + sband + '_' + polpair
				if not (brief_key in Vnotfound):
					Vnotfound.append(brief_key)
				continue
		else:
			Ncopied += 1

		fout.write(binhdr)
		fout.write(rawvis)

	# Finished
	fout.close()
	print('')
	print('Vis. copied     : %d' % (Ncopied))
	print('Vis. replaced   : %d' % (Nreplaced))
	print('Vis. not in src : %d' % (Nnotfound))
	if Nnotfound > 0:
		print('        details : %s' % (str(Vnotfound)))

	# TODO:
	print ('Generated visibilities and %s, but please manually edit its OUTPUT FILENAME line to fix the path (TODO: automatic path fix)' % (newinputfile))
Example #4
0
def stitchVisibilityfile(basename, cfg, writeMetaOnly=False):
    """Copy visibilities from given base DiFX fileset into new file, while doing frequency stitching during the process"""

    # Get settings
    target_bw = cfg['target_bw']
    target_nchan = cfg['target_nchan']
    stitch_basefreqs = cfg['stitch_basefreqs']
    stitch_endfreqs = cfg['stitch_endfreqs']
    stitch_antennas = cfg['stitch_antennas']
    blank_viz = numpy.fromstring('\0' * 8 * target_nchan, dtype='complex64')

    # All polarisation pairs 'RR', 'RL', ... to 'YL'
    if cfg['stitch_nstokes'] > 2:
        pols_list = ['R', 'L', 'X', 'Y']
        polpairs, tmp = setCrossProd(pols_list, pols_list)
    else:
        polpairs = ['RR', 'LL', 'XX', 'YY']  # parallel-hand only
        polpairs += ['RX', 'LY', 'XR', 'YL']  # potential parallel-hand
        polpairs += ['LX', 'RY', 'XL', 'YR'
                     ]  # potential parallel-hand (depends on polconvert.py)

    # Extract meta-infos from the DiFX .INPUT file
    if basename.endswith(('.difx', 'input')):
        basename = basename[:basename.rfind('.')]
    basename_pathless = basename
    if basename.rfind('/') >= 0:
        basename_pathless = basename[(basename.rfind('/') + 1):]
    inputfile = basename + '.input'
    inputfile_cfg = parseDiFX.get_common_settings(inputfile)
    (numfreqs, freqs) = parseDiFX.get_freqtable_info(inputfile)
    (numtelescopes, telescopes) = parseDiFX.get_telescopetable_info(inputfile)
    (numdatastreams,
     datastreams) = parseDiFX.get_datastreamtable_info(inputfile)
    (numbaselines, baselines) = parseDiFX.get_baselinetable_info(inputfile)
    if 'difxfile' not in inputfile_cfg:
        parser.error("Couldn't parse COMMON SETTINGS of input file " +
                     inputfile + " correctly")
    if numfreqs == 0 or numtelescopes == 0 or numdatastreams == 0 or numbaselines == 0:
        parser.error("Couldn't parse input file " + inputfile + " correctly")

    # Internal frequency tables and remappings from old to new freq indices
    out_freqs = {}  # dict out['id']=Freq()
    in_rec_freqs = 0  # num of distinct recorded freqs
    n_max_freqs = 256
    freq_remaps = [-1] * n_max_freqs
    freq_remaps_isNew = [False] * n_max_freqs

    # Collect all recorded freqs listed in DATASTREAMS
    for d in datastreams:
        for fqidx in d.recfreqindex:
            id = findFreqObj(out_freqs, freqs[fqidx])
            if (id == None):
                id = inventNextKey(out_freqs)
                freq_remaps[fqidx] = id
                out_freqs[id] = copy.deepcopy(freqs[fqidx])
                in_rec_freqs += 1
                print("Keeping recorded frequency  : index %2d/%2d : %s" %
                      (fqidx, id, freqs[fqidx].str()))
            else:
                if (freq_remaps[fqidx] < 0):
                    existing_fq = id
                    # print ('Warning: DiFX produced duplicate FREQ entries! Telescope %d/%s, remapping freq[%d]-->existing freq[%d]' % (d.telescopeindex,telescopes[d.telescopeindex].name,fqidx,existing_fq))
                    out_freqs[id] = out_freqs[existing_fq]
                    freq_remaps[fqidx] = existing_fq
                    print(
                        "Remapping duplicate freq    : index %2d/%2d : %s --> index %2d : %s"
                        % (fqidx, id, freqs[fqidx].str(), existing_fq,
                           out_freqs[existing_fq].str().strip()))

    # Collect all zoom bands listed in DATASTREAMS that already match the target bandwidth
    for d in datastreams:
        for fqidx in d.zoomfreqindex:
            zf = freqs[fqidx]
            if (zf.bandwidth != target_bw) or (
                (zf.numchan / zf.specavg) != target_nchan):
                # print ("Skipping zoom frequency     : index %d : %s" % (fqidx,zf.str().strip()))
                continue
            i = findFreqObj(out_freqs, zf)
            if (i == None):
                id = inventNextKey(out_freqs)
                freq_remaps[fqidx] = id
                out_freqs[id] = copy.deepcopy(zf)
                print("Keeping zoom frequency      : index %2d/%2d : %s" %
                      (fqidx, id, zf.str()))

    # Create per-antenna list of old frequency IDs (recs first then zooms)
    antenna_freqids = {}
    for d in datastreams:
        ant = telescopes[d.telescopeindex].name
        fqs = d.recfreqindex
        if ant not in antenna_freqids:
            antenna_freqids[ant] = list(fqs)
        else:
            antenna_freqids[ant] += listDiff(fqs, antenna_freqids[ant])
    for ant in antenna_freqids:
        # When antenna has adhoc phase info make sure it covers all IFs
        if ant not in cfg['adhoc_phasors']:
            continue
        Nrec = len(antenna_freqids[ant])
        Ndef = len(cfg['adhoc_phasors'][ant])
        if Ndef != Nrec:
            print(
                "Error: config file section [phases_deg] provides phases for '%s' for %d of %d IFs!"
                % (ant, Ndef, Nrec))
            sys.exit(1)
    for d in datastreams:
        ant = telescopes[d.telescopeindex].name
        fqs = d.zoomfreqindex
        antenna_freqids[ant] += listDiff(fqs, antenna_freqids[ant])

    # Determine how much averaging post-FFT was done in DiFX itself
    common_difx_avgfactor = -1
    for key in out_freqs:
        common_difx_avgfactor = out_freqs[key].specavg
    if common_difx_avgfactor < 1:
        print(
            "Warning: may have failed to detect original .input 'CHANS TO AVG' factor! Assuming factor 1!"
        )
        common_difx_avgfactor = 1

    # Check stitch config: re-use existing USB zoom bands or freqs where possible, invent new zoom bands if necessary
    stitch_out_ids = []
    for fsky in stitch_basefreqs:
        match_existing = [(out_freqs[key].freq == fsky
                           and out_freqs[key].bandwidth == target_bw
                           and out_freqs[key].numchan / out_freqs[key].specavg
                           == target_nchan and not out_freqs[key].lsb)
                          for key in out_freqs.keys()]
        exists = any(match_existing)
        if exists:
            # Re-use an existing matching zoom frequency
            i = match_existing.index(True)
            id = out_freqs.keys()[i]
            zf = out_freqs[id]
            zf.numchan = target_nchan * common_difx_avgfactor
            zf.specavg = common_difx_avgfactor
            stitch_out_ids.append(id)
            old_id = freq_remaps.index(id)
            print("Re-using zoom for stitch    : index %2d/%2d : %s" %
                  (old_id, id, zf.str()))
        else:
            # Invent new zoom
            zf = parseDiFX.Freq()
            zf.freq = fsky
            zf.bandwidth = target_bw
            zf.numchan = target_nchan * common_difx_avgfactor
            zf.specavg = common_difx_avgfactor
            zf.lsb = False
            id = inventNextKey(out_freqs)
            out_freqs[id] = copy.deepcopy(zf)
            stitch_out_ids.append(id)
            print("Creating new zoom frequency : index --/%2d : %s" %
                  (id, zf.str()))

    stitch_ids = [n for n in range(len(stitch_out_ids))]

    # Also map each to-be-stitched-multi-zoom into respective new single post-stitch zoom
    for fi in range(in_rec_freqs, len(freqs)):
        if (freqs[fi].bandwidth
                == target_bw) and (freqs[fi].numchan / freqs[fi].specavg
                                   == target_nchan):
            # Retain one-to-one map for existing matching zoom freqs
            freq_remaps_isNew[fi] = False
            continue
        stid = getGlueIndex(freqs[fi].freq, cfg)
        if (stid >= 0):
            if freqs[fi].lsb:
                # Ignore LSB. All zoom outputs should be USB.
                print("Ignore LSB %s" % (freqs[fi].str()))
                continue
            # Add to many-to-one map of invented zoom freqs
            freq_remaps[fi] = stitch_out_ids[stid]
            freq_remaps_isNew[fi] = True
            print(
                "Map zoom %s to stitched single %12.6f--%12.6f : in fq#%2d -> stitch#%d -> out fq#%2d"
                % (freqs[fi].str(), stitch_basefreqs[stid],
                   stitch_endfreqs[stid], fi, stid, stitch_out_ids[stid]))
        else:
            print("No map for index %d : %s" % (fi, freqs[fi].str()))

    # Propagate any averaging to be done by difx2difx.py into the frequency table
    if cfg['extra_chavg'] > 1:
        for of in out_freqs.keys():
            out_freqs[of].specavg = out_freqs[of].specavg * cfg['extra_chavg']

    # Propagate ad-hoc phase table infos from rec freq (parent) to any zoom (child)
    for ant in antenna_freqids:
        if ant not in cfg['adhoc_phasors']:
            continue
        if len(antenna_freqids[ant]) <= len(cfg['adhoc_phasors'][ant]):
            continue
        # Duplicate .cfg adhoc phase entries (rec bands) to e.g. zoom bands
        Ndefined = len(cfg['adhoc_phasors'][ant])
        Ndefneeded = len(antenna_freqids[ant])
        freqs_w_defined_phases = antenna_freqids[ant][0:Ndefined]
        cfg['adhoc_phasors'][ant] = numpy.append(
            cfg['adhoc_phasors'][ant], [1 + 0j] * (Ndefneeded - Ndefined))
        for n in range(Ndefined, Ndefneeded):
            fq = antenna_freqids[ant][n]
            for m in range(Ndefined):
                recfq = freqs_w_defined_phases[m]
                if isParentFreq(freqs[recfq], freqs[fq]):
                    cfg['adhoc_phasors'][ant][n] = cfg['adhoc_phasors'][ant][m]
                    print(
                        'Adhoc phase %s phase offset : zoom %s gets config file entry %d phase of rec %s'
                        % (ant, freqs[fq].str(), m, freqs[recfq].str()))

    # Read the DiFX .difx/DIFX_* file
    #glob_pattern = basename + '.difx/DIFX_*.s*.b*'
    glob_pattern = inputfile_cfg['difxfile'] + '/DIFX_*.s*.b*'
    difxfileslist = glob.glob(glob_pattern)
    if len(difxfileslist) <= 0:
        print('Error: no visibility data file found in %s!' % (glob_pattern))
        return
    difxfilename = difxfileslist[0]
    difxfilename_pathless = difxfilename
    if difxfilename.rfind('/') >= 0:
        difxfilename_pathless = difxfilename[(difxfilename.rfind('/') + 1):]
    difxfile = open(difxfilename, 'r')

    # Output DiFX file
    difxoutdir = basename_pathless + 'D2D.difx'
    difxoutname = difxoutdir + '/' + difxfilename_pathless
    if not writeMetaOnly:
        # New output directory and file
        try:
            os.mkdir(difxoutdir)
        except:
            pass
        if os.path.exists(difxoutname):
            print(
                'Warning: %s already exists, skipping the processing of %s!' %
                (difxoutname, difxfilename))
            return
        difxout = open(difxoutname, 'w')

        # Copy PCAL files if any
        for pcalfile in glob.glob(inputfile_cfg['difxfile'] + '/PCAL_*'):
            print('Copying %s' % (pcalfile))
            shutil.copy(pcalfile, difxoutdir + '/')
    else:
        difxout = None

    # Pull out antenna indices based on telescope names
    telNames = [t.name for t in telescopes]
    stAntIDs = []
    for sa in stitch_antennas:
        id = [n for n in range(len(telescopes)) if telescopes[n].name == sa]
        if len(id) != 1:
            print(
                "Warning: got %d hits for telescope %s in .difx file telescope list of %s!"
                % (len(id), sa, str(telNames)))
            continue
        id = id[0] + 1  # index (0..N-1) into antenna nr (1..N)
        stAntIDs.append(id)

    # Affected baselines; data of these are stored in stitching memory, other baselines need no stitching
    baseline_list = []
    for sid in stAntIDs:
        # Cross
        for t in range(numtelescopes):
            tid = t + 1  # index (0..N-1) into antenna nr (1..N)
            if (tid != sid):
                baseline_list.append(tid + sid * 256)
                baseline_list.append(sid + tid * 256)
        # Auto
        baseline_list.append(sid + sid * 256)

    # Dictionary of per-polarizationpair working buffers into which spectra are stitched
    stitch_workbufs = {
        bline: {
            polkey: {bandkey: numpy.copy(blank_viz)
                     for bandkey in stitch_ids}
            for polkey in polpairs
        }
        for bline in baseline_list
    }
    stitch_chcounts = {
        bline: {
            polkey: {bandkey: 0
                     for bandkey in stitch_ids}
            for polkey in polpairs
        }
        for bline in baseline_list
    }
    stitch_freqbws = {
        bline: {
            polkey: {bandkey: []
                     for bandkey in stitch_ids}
            for polkey in polpairs
        }
        for bline in baseline_list
    }
    stitch_timestamps = {
        bline: {
            polkey: {bandkey: -1
                     for bandkey in stitch_ids}
            for polkey in polpairs
        }
        for bline in baseline_list
    }
    nstitched = 0
    ncopied = 0
    nskipped = 0
    nincomplete = 0

    # Parse each visibility entry
    seconds_prev = -1
    seconds_first = 0
    seconds_report_prev = 0
    while not writeMetaOnly:

        (vishdr, binhdr) = getVisibilityHeader(difxfile)
        if len(vishdr) <= 0:
            break

        # Visibility properties
        baseline = vishdr[0]
        mjd = vishdr[1]
        seconds = vishdr[2]
        freqindex = vishdr[5]
        polpair = vishdr[6]
        weight = vishdr[8]
        uvw = vishdr[9:12]
        # Antenna order as in difx2mark4: ref=ant1="256*nr", rem=ant2="nr%256"
        ant2 = baseline % 256
        ant1 = (baseline - ant2) / 256
        ant1name = telescopes[ant1 - 1].name
        ant2name = telescopes[ant2 - 1].name
        T = mjd + seconds / 86400.0
        baselinestr = '%s-%s' % (ant1name, ant2name)

        # Stokes check:
        if polpair not in polpairs:
            print(
                'Error: unexpected pol pair %s, perhaps increase stitch_nstokes (currently %d) in config file?'
                % (polpair, cfg['stitch_nstokes']))
            sys.exit(-1)

        # Number of channels in this baseband
        nchan = freqs[freqindex].numchan / freqs[freqindex].specavg
        fsky = freqs[freqindex].freq
        bw = freqs[freqindex].bandwidth

        # Read the entire visibility data from disk
        rawvis = difxfile.read(8 * nchan)
        if len(rawvis) != 8 * nchan:
            print('Short read! Stopping.')
            break

        if (seconds != seconds_prev):
            seconds_prev = seconds
            if cfg['verbose']:
                print(('\n---- %d %12.7f ' + '-' * 115) % (mjd, seconds))
            elif (seconds - seconds_report_prev) > 1.0:
                if seconds_first <= 0:
                    seconds_first = mjd * 24 * 60 * 60 + seconds
                dTstart = mjd * 24 * 60 * 60 + seconds - seconds_first
                seconds_report_prev = seconds
                print("at %d %12.7f, %.3f seconds from start" %
                      (mjd, seconds, dTstart))
                print("\033[F\033[F")

        if cfg['dbg_drop_antennas'] and ant1name in cfg[
                'dbg_drop_antennas'] or ant2name in cfg['dbg_drop_antennas']:
            nskipped += 1
            continue

        # Remap the frequency reference
        out_freqindex = freqindex
        if freq_remaps[freqindex] >= 0:
            # print ('Remapping input freq index %d to old/new output freq %d' % (freqindex,freq_remaps[freqindex]))
            out_freqindex = freq_remaps[freqindex]
            vishdr[5] = out_freqindex
            binhdr = parseDiFX.make_output_header_v1(vishdr)

        # Destination stitch band index
        if not freqs[freqindex].lsb:
            stid = getGlueIndex(fsky, cfg)
        else:
            stid = -1

        # Info string
        vis_info = '%s-%s/%d/%d(%d):sf<%d>:%.7f/%s  mjd:%12.8f nchan:%4d bw:%.7f uvw=%s' % (
            ant1name, ant2name, baseline, out_freqindex, freqindex, stid, fsky,
            polpair, T, nchan, bw, str(uvw))

        # Adjust phase of adhoc per-IF phases in config file
        if hasPhaseAdj(ant1name, ant2name, freqindex, antenna_freqids, cfg):
            # (ToDO: always for all data, or only when visibility data will actually be used?)
            rawvis = spectralPhaseAdj(rawvis, ant1name, ant2name, freqindex,
                                      antenna_freqids, cfg)

        # Write out visibility to output file
        if baseline not in baseline_list:

            # Baseline for which no freq stitching was requested

            # Keep data if bandwidth matches
            if (bw == target_bw) and (nchan == target_nchan):

                if wasVisAlreadyWritten(mjd, seconds, baseline, out_freqindex,
                                        polpair):
                    if cfg['verbose']:
                        print('(copy): %s' % (vis_info))
                    continue
                if cfg['verbose']:
                    print('copy  : %s' % (vis_info))

                assert (freq_remaps[freqindex] >= 0)
                difxout.write(binhdr)
                if cfg['extra_chavg'] > 1:
                    rawvis = spectralAvgRaw(rawvis, cfg['extra_chavg'])
                difxout.write(rawvis)
                ncopied += 1

            # Discard data if bandwidth mismatch. Happens e.g. for zoom bands on the non-stitch baselines,
            # can ignore these because covered by a different and complete zoomband that doesn't need stitching
            else:
                nskipped += 1
                if cfg['verbose']:
                    print('skip  : %s' % (vis_info))
                    pass
        else:

            # Baseline where freq stitching was requested
            resetStitch = False

            if (bw == target_bw) and freqs[freqindex].lsb:

                # Visibility already at target bandwidth, but LSB i.e. not the expected zoom (all-USB)
                # As the full-bw zoom ought to exist or is going to be stitched together, we skip the LSB x LSB
                # visibility here. This avoids duplicates (since e.g. 86188 LSB x LSB rec 32 MHz == 86156 USB x USB zoom 32 MHz)

                nskipped += 1
                if cfg['verbose']:
                    print('skip L: %s' % (vis_info))

            elif (bw == target_bw):

                # Visibility already at target bandwidth, just copy the data

                if wasVisAlreadyWritten(mjd, seconds, baseline, out_freqindex,
                                        polpair):
                    if cfg['verbose']:
                        print('(take): %s' % (vis_info))
                    continue
                if cfg['verbose']:
                    print('take  : %s' % (vis_info))

                assert (freq_remaps[freqindex] >= 0)
                difxout.write(binhdr)
                if cfg['extra_chavg'] > 1:
                    rawvis = spectralAvgRaw(rawvis, cfg['extra_chavg'])
                difxout.write(rawvis)
                ncopied += 1

                # Signal for further below: clear old working data, set new timestamp
                resetStitch = True

            else:

                # Need to stitch narrowband visibility into wideband context

                # Ignore bands not requested to be stitched
                if (stid < 0):
                    if cfg['verbose']:
                        print('ignore: %s (not in stitch freqs list)' %
                              (vis_info))
                    continue

                # Get the complex data
                visdata = numpy.fromstring(rawvis, dtype='complex64')

                # Determine the region into which this visibility data would be inserted
                choffset = int(((fsky - stitch_basefreqs[stid]) / target_bw) *
                               target_nchan)
                if (choffset + nchan) > target_nchan:
                    # Sometimes a wide recorded band falls into the stitching frequency region,
                    # but then may not fit fully into the stitching freq region.
                    # Options are:
                    # 1) discard and hope the corresponding zoom band data will appear soon
                    #if cfg['verbose']:
                    #	print ('ignore: %s (too wide, maybe a recorded instead of zoom freq)' % (vis_info))
                    #continue
                    # 2) use the maximum possible amount of the data
                    nchan_new = target_nchan - choffset
                    if (nchan_new < 1):
                        print(
                            'Unexpected, after cropping a too-wide band ended up with %d remaining channels'
                            % (nchan_new))
                        continue
                    visdata = visdata[:nchan_new]
                    bw = bw * nchan_new / nchan
                    nchan = nchan_new
                    if cfg['verbose']:
                        print('crop  : %s to %d chan' % (vis_info, nchan_new))

                # Keep track of timestamp
                if (stitch_timestamps[baseline][polpair][stid] < 0):
                    stitch_timestamps[baseline][polpair][stid] = seconds

                # Detect a jump in timestamp
                if (seconds != stitch_timestamps[baseline][polpair][stid]):
                    # We assume the input visibilities are stored in time-increasing order,
                    # so a timestamp change indicates all data of previous time should've been fully read by now (no older data follows)
                    ncurr = stitch_chcounts[baseline][polpair][stid]
                    if (ncurr > 0):
                        T_old = mjd + stitch_timestamps[baseline][polpair][
                            stid] / 86400.0
                        nincomplete += 1
                        msg = '%s-%s/%d/%d:%.7f/%s mjd:%12.8f only %d of %d channels' % (
                            ant1name, ant2name, baseline, stid,
                            stitch_basefreqs[stid], polpair, T_old, ncurr,
                            target_nchan)
                        if cfg['verbose']:
                            print(
                                'Warning: discarded an incomplete stitch: %s' %
                                (msg))
                    assert (ncurr <= target_nchan)
                    stitch_timestamps[baseline][polpair][stid] = seconds
                    stitch_chcounts[baseline][polpair][stid] = 0
                    stitch_freqbws[baseline][polpair][stid] = []
                    stitch_workbufs[baseline][polpair][stid].fill(0)
                    #print (('---- %d %.7f ' + '-'*80 + '\n') % (mjd,seconds))

                # Insert the new visibility data into the correct working buffer at the correct location
                stitch_workbufs[baseline][polpair][stid][choffset:(
                    choffset + nchan)] = visdata
                stitch_chcounts[baseline][polpair][stid] += nchan
                stitch_freqbws[baseline][polpair][stid].append(bw)

                #print (' glue bline %s %s band %d : ch %4d ... %4d : %d new chs : stitched %d ch total' % (baselinestr,polpair,stid,choffset,choffset+nchan-1,nchan,stitch_chcounts[baseline][polpair][stid]))
                #print choffset, nchan, choffset+nchan, stitch_chcounts[baseline][polpair][stid], (100.0*stitch_chcounts[baseline][polpair][stid])/target_nchan

                # Assembled the full bandwidth now?
                if stitch_chcounts[baseline][polpair][stid] >= target_nchan:
                    TS = mjd + stitch_timestamps[baseline][polpair][
                        stid] / 86400.0
                    bwsum = numpy.sum(stitch_freqbws[baseline][polpair][stid])
                    fsky = out_freqs[out_freqindex].freq
                    vis_info = '%s-%s/%d/%d(%d):sf<%d>:%.7f/%s  mjd:%12.8f nchan:%4d bw:%.7f uvw=%s' % (
                        ant1name, ant2name, baseline, out_freqindex, freqindex,
                        stid, fsky, polpair, TS,
                        stitch_chcounts[baseline][polpair][stid], bwsum,
                        str(uvw))

                    if wasVisAlreadyWritten(mjd, seconds, baseline,
                                            out_freqindex, polpair):
                        if cfg['verbose']:
                            print('(stch dup): %s' % (vis_info))
                        continue

                    # Double-check the bandwidth
                    if numpy.abs(target_bw - bwsum) > 100:
                        print(
                            'Warning: stitched bands gave %.6f MHz bandwidth, differs from target %.6f MHz!'
                            % (bwsum, target_bw))

                    # Write header and assembled data
                    difxout.write(binhdr)
                    if cfg['extra_chavg'] > 1:
                        vis = spectralAvgNumpy(
                            stitch_workbufs[baseline][polpair][stid],
                            cfg['extra_chavg'])
                    else:
                        vis = stitch_workbufs[baseline][polpair][stid]

                    if cfg['verbose']:
                        print('stitch: %s' % (vis_info))

                    vis.tofile(difxout)
                    nstitched += 1

                    # Reset
                    resetStitch = True

            if resetStitch and (stid >= 0):
                stitch_timestamps[baseline][polpair][stid] = seconds
                stitch_chcounts[baseline][polpair][stid] = 0
                stitch_freqbws[baseline][polpair][stid] = []
                stitch_workbufs[baseline][polpair][stid].fill(0)

        #if nstitched > 100: break  # quick-exit for debug

    ## Generate new .input

    # New FREQ table
    new_freqs = copy.deepcopy(out_freqs)

    # New DATASTREAM table
    new_datastreams = []
    for ds in datastreams:

        # Copy the existing values, then update/replace
        newds = copy.deepcopy(ds)

        # Retain all recorded freqs and bands
        newds.recfreqindex = [freq_remaps[rfi] for rfi in newds.recfreqindex]
        newds.recfreqpols = [p for p in ds.recfreqpols]
        newds.nrecfreq = len(newds.recfreqindex)
        if False:
            # Sort the bands to have L R L R L R L R rather than L L L L R R R R,
            # makes baseline table re-creation easier
            # However, messes up PCAL file parsing since for some reason parser
            # compares PCAL file pols&freqs against 'recbandpol'&'recfreq[recbandindex]'
            # which then are in unexpected order
            newds.recbandpol = [
                x for _, x in sorted(zip(newds.recbandindex, newds.recbandpol),
                                     key=lambda pair: pair[0])
            ]
            newds.recbandindex.sort()
        assert (len(newds.recfreqindex) == len(newds.recfreqpols))
        assert (newds.nrecband == sum(newds.recfreqpols))

        # Retain bandwidth-matching existing zoom freqs & bands
        # and also add all stitched/invented zoom freqs & bands
        newds.zoomfreqindex, newds.zoomfreqpols, newds.zoombandindex, newds.zoombandpol = [], [], [], []
        #ds_old_zooms = [zfi for zfi in ds.zoomfreqindex if (freq_remaps[zfi]>=0 and not freq_remaps_isNew[zfi])]
        #ds_new_zooms = [zfi for zfi in ds.zoomfreqindex if (freq_remaps[zfi]>=0 and freq_remaps_isNew[zfi])]
        all_zooms = [
            zfi for zfi in ds.zoomfreqindex if (freq_remaps[zfi] >= 0)
        ]
        for zf in all_zooms:
            if freq_remaps[zf] not in newds.zoomfreqindex:
                pols = getPolsForFreq(ds, zf)  # from original DS & Freq
                if len(pols) < 1: continue
                newds.zoomfreqindex.append(freq_remaps[zf])
                newds.zoomfreqpols.append(len(pols))
                n = newds.zoomfreqindex.index(freq_remaps[zf])
                newds.zoombandindex += [n] * len(pols)
                newds.zoombandpol += pols
        newds.nzoomfreq = len(newds.zoomfreqpols)
        newds.nzoomband = sum(newds.zoomfreqpols)

        # Store
        new_datastreams.append(newds)

        if cfg['verbose']:
            print("DS%d : rec freqs      : %s" %
                  (datastreams.index(ds), str(newds.recfreqindex)))
            print("      rec freq pols  : %s" % (str(newds.recfreqpols)))
            print("      rec bands      : %s" % (str(newds.recbandindex)))
            print("      rec band pols  : %s" % (str(newds.recbandpol)))
            print("      zoom freqs     : %s" % (str(newds.zoomfreqindex)))
            print("      zoom freq pols : %s" % (str(newds.zoomfreqpols)))
            print("      zoom bands     : %s" % (str(newds.zoombandindex)))
            print("      zoom band pols : %s" % (str(newds.zoombandpol)))
            print(
                "      counts         : %d rec fq, %d zoom fq, %d rec band, %d zoom band"
                % (newds.nrecfreq, newds.nzoomfreq, newds.nrecband,
                   newds.nzoomband))

    # New BASELINE table
    full_freq_remaps = [
        freq_remaps[n] if freq_remaps[n] >= 0 else n
        for n in range(len(freq_remaps))
    ]
    #print ('Full remap table:',str(len(full_freq_remaps)), ' entries, ', str(full_freq_remaps))
    new_baselines = []
    for b in baselines:
        newbl = copy.deepcopy(b)
        ods1 = datastreams[b.dsaindex]
        ods2 = datastreams[b.dsbindex]
        nds1 = new_datastreams[newbl.dsaindex]
        nds2 = new_datastreams[newbl.dsbindex]
        ant1 = telescopes[nds1.telescopeindex].name
        ant2 = telescopes[nds2.telescopeindex].name
        if cfg['verbose']:
            print("Baseline %s x %s" % (ant1, ant2))
            print("     stream %d x stream %d" %
                  (newbl.dsaindex, newbl.dsbindex))

        newbl.dsabandindex = []
        newbl.dsbbandindex = []
        newbl.freqpols = []
        copied_ids = []
        for i in range(len(b.dsabandindex)):
            # For each baseline "frequency" entry there are two lists of bands e.g. [0,1,2,3] x [0,2,1,3]
            # that were correlated against each other in DiFX.
            # Datastream and Freq were updated above, hence the band indexing has changed.
            # Re-map the referenced band numbers/indices in the lists from old to new indices.
            # Also, look up Freq of the respective band indices and skip narrow pre-stitch zooms from Baseline.
            bl_bands_A = b.dsabandindex[i]
            bl_bands_B = b.dsbbandindex[i]
            bl_new_bands_A, bl_new_bands_B = [], []
            for (bl_band_A, bl_band_B) in zip(bl_bands_A, bl_bands_B):
                oldFqA, polA = getFreqPolOfBand(ods1, bl_band_A)
                oldFqB, polB = getFreqPolOfBand(ods2, bl_band_B)
                if oldFqA < 0 or oldFqB < 0:
                    print(
                        "Warning: unexpectedly could not locate old input band information!"
                    )
                    continue
                newFqA = full_freq_remaps[oldFqA]
                newFqB = full_freq_remaps[oldFqB]
                assert (freqs[oldFqA].bandwidth == freqs[oldFqB].bandwidth)
                if (newFqA not in new_freqs) or (newFqB not in new_freqs):
                    # print ("     corr %2d : fq %2d(old:%d) x %2d(old:%d) %s%s : skip, lies outside of output bands" % (i,newFqA,oldFqA,newFqB,oldFqB,polA,polB))
                    continue
                if (new_freqs[newFqA].bandwidth != target_bw) or (
                        new_freqs[newFqB].bandwidth != target_bw):
                    print(
                        "     corr %2d : fq %2d(old:%d) x %2d(old:%d) %s%s : skip, wrong bandwidth in %s x %s"
                        % (i, newFqA, oldFqA, newFqB, oldFqB, polA, polB,
                           new_freqs[newFqA].str(), new_freqs[newFqB].str()))
                    continue
                id = '%d_%d_%s%s' % (newFqA, newFqB, polA, polB)
                if id not in copied_ids:
                    newBandA = getBandIndexOfFreqPol(nds1, newFqA, polA)
                    newBandB = getBandIndexOfFreqPol(nds2, newFqB, polB)
                    if newBandA < 0 or newBandB < 0:
                        continue
                    bl_new_bands_A.append(newBandA)
                    bl_new_bands_B.append(newBandB)
                    copied_ids.append(id)
                    if cfg['verbose']:
                        print(
                            "     corr %2d : bands %2d x %2d %s%s : fq %2d x %2d --> new fq %2d x %2d : copied, %s"
                            % (i, bl_band_A, bl_band_B, polA, polB, oldFqA,
                               oldFqB, newFqA, newFqB, id))
                        #print ("                fq old %s x %s" % (freqs[oldFqA].str().strip(),freqs[oldFqB].str().strip()))
                        #print ("                fq new %s x %s" % (new_freqs[newFqA].str().strip(),new_freqs[newFqB].str().strip()))
                else:
                    if cfg['verbose']:
                        print(
                            "     corr %2d : bands %2d x %2d %s%s : fq %2d x %2d --> new fq %2d x %2d : skip, sub-stitch"
                            % (i, bl_band_A, bl_band_B, polA, polB, oldFqA,
                               oldFqB, newFqA, newFqB))
            if len(bl_new_bands_A) > 0:
                newbl.dsabandindex.append(bl_new_bands_A)
                newbl.dsbbandindex.append(bl_new_bands_B)
                newbl.freqpols.append(len(bl_new_bands_A))

        newbl.nfreq = len(newbl.dsabandindex)
        new_baselines.append(newbl)
    print('\n')

    # Read original .input without parsing
    fin = open(inputfile, "r")
    in_lines = fin.readlines()
    fin.close()

    # Replace OUTPUT FILENAME entry
    for l in in_lines:
        if l[:16] == "OUTPUT FILENAME:":
            i = in_lines.index(l)
            in_lines[i] = "%-20s%s\n" % ("OUTPUT FILENAME:", difxoutdir)

    # Replace CALC FILENAME and CORE CONF FILENAME and point them to a local, renamed copy of these files
    for l in in_lines:
        if l[:19] == 'CORE CONF FILENAME:':
            i = in_lines.index(l)
            orig_threadfile = l[20:].strip()
            in_lines[i] = "%-20s%s\n" % ("CORE CONF FILENAME:",
                                         basename_pathless + 'D2D.threads')
            try:
                shutil.copyfile(orig_threadfile,
                                basename_pathless + 'D2D.threads')
            except:
                print('Note: could not copy %s ' % (orig_threadfile))
        elif l[:14] == "CALC FILENAME:":
            i = in_lines.index(l)
            orig_calcfile = l[20:].strip()
            in_lines[i] = "%-20s%s\n" % ("CALC FILENAME:",
                                         basename_pathless + 'D2D.calc')
            try:
                shutil.copyfile(orig_calcfile, basename_pathless + 'D2D.calc')
            except:
                print('Note: could not copy %s ' % (orig_calcfile))

    # Generate new reference .input for later, re-use parts of original .input
    outputinputfile = basename_pathless + 'D2D.input'
    fout = open(outputinputfile, "w")
    idx = [i for i, elem in enumerate(in_lines) if "FREQ TABLE" in elem]
    for n in range(idx[0]):
        fout.write(in_lines[n])
    parseDiFX.put_freqtable_info(fout, new_freqs)
    idx1 = [i for i, elem in enumerate(in_lines) if "TELESCOPE TABLE" in elem]
    idx2 = [i for i, elem in enumerate(in_lines) if "DATASTREAM TABLE" in elem]
    for n in range(idx1[0], idx2[0]):
        fout.write(in_lines[n])
    parseDiFX.put_datastreamtable_info(fout, new_datastreams)
    parseDiFX.put_baselinetable_info(fout, new_baselines)
    idx = [i for i, elem in enumerate(in_lines) if "DATA TABLE" in elem]
    for trailing in in_lines[idx[0]:]:
        fout.write(trailing)
    fout.close()

    # Go through the local copy .calc file and point references therein to other local copies
    for line in fileinput.input(basename_pathless + 'D2D.calc', inplace=True):
        # fileinput module: in this loop STDOUT is redirected by the module into the file
        if 'IM FILENAME' in line:
            orig_imfile = line[15:].strip()
            copied_imfile = basename_pathless + 'D2D.im'
            shutil.copyfile(orig_imfile, copied_imfile)
            line = '%-20s%s' % ('IM FILENAME:', copied_imfile)
        elif 'FLAG FILENAME' in line:
            orig_flagfile = line[15:].strip()
            copied_flagfile = basename_pathless + 'D2D.flag'
            shutil.copyfile(orig_flagfile, copied_flagfile)
            line = '%-20s%s' % ('FLAG FILENAME:', copied_flagfile)
        print(line.strip())

    # Finished
    if not writeMetaOnly:
        print('\nDone! Final statistics:')
        print('    vis. copied        : %d' % (ncopied))
        print('    vis. skipped       : %d' % (nskipped))
        print('    vis. stitched      : %d' % (nstitched))
        print('    incomplete stitch  : %d' % (nincomplete))
        print('\nOutput files:')
        print('    visbility data     : %s' % (difxoutname))
        print('    changes for .input : %s' % (outputinputfile))
        print(' ')
    else:
        print('\nDone! Metadata-only mode.')
        print('\nOutput files:')
        print('    changes for .input : %s' % (outputinputfile))
Example #5
0
    parser.error(
        "If you scrunch baselines, you must allow all of them (targetbaseline < 0)"
    )

if scrunchbaselines and scrunchautocorrs:
    parser.error("You can only scrunch baselines, or autocorrs, not both!")

if scrunchautocorrs and targetbaseline >= 0:
    parser.error(
        "If you scrunch autocorrs, you must allow all of them (targetbaseline < 0)"
    )

(numconfigs, configs) = parseDiFX.get_configtable_info(inputfile)
(numfreqs, freqs) = parseDiFX.get_freqtable_info(inputfile)
(numtelescopes, telescopes) = parseDiFX.get_telescopetable_info(inputfile)
(numdatastreams, datastreams) = parseDiFX.get_datastreamtable_info(inputfile)
(numbaselines, baselines) = parseDiFX.get_baselinetable_info(inputfile)

if numfreqs == 0 or numtelescopes == 0 or numdatastreams == 0 or numbaselines == 0 or numconfigs == 0:
    parser.error("Couldn't parse input file " + inputfile + " correctly")

inttime = configs[0].inttime
chans = []
amp = []
phase = []
for i in range(numfreqs):
    amp.append([])
    phase.append([])
    for j in range(maxchannels):
        amp[i].append([])
        phase[i].append([])
def filterVisibilityfile(basename,targetAnts):

	# Extract meta-infos from the DiFX .INPUT file
	if basename.endswith(('.difx','.input','.calc')):
		basename = basename[:basename.rfind('.')]
	pathless_basename = basename
	if basename.rfind('/')>=0:
		pathless_basename = basename[basename.rfind('/')+1:]
	inputfile = basename + '.input'
	(numfreqs, freqs) = parseDiFX.get_freqtable_info(inputfile)
	(numtelescopes, telescopes) = parseDiFX.get_telescopetable_info(inputfile)
	(numdatastreams, datastreams) = parseDiFX.get_datastreamtable_info(inputfile)
	(numbaselines, baselines) = parseDiFX.get_baselinetable_info(inputfile)
	if numfreqs == 0 or numtelescopes == 0 or numdatastreams == 0 or numbaselines == 0:
		parser.error("Couldn't parse input file " + inputfile + " correctly")

	# Read the DiFX .difx/DIFX_* file
	difxfileslist = glob.glob(basename + '.difx/DIFX_*.s*.b*')
	difxfilename = difxfileslist[0]
	difxfile = open(difxfilename, 'r')
	difxoutdir = pathless_basename + 'filtered.difx'

	# Stop early if telescope has no data in this scan
	telescopenames = [t.name for t in telescopes]
	if not any(target in telescopenames for target in targetAnts):
		print ('Telescope(s) %s not among in stations %s of %s. Nothing to do!' % (str(targetAnts),str(telescopenames),basename)) 
		#shutil.copytree(basename + '.difx', difxoutdir)
		#return

	# Prepare output
	try:
		os.mkdir(difxoutdir)
	except:
		pass
	difxoutname = difxoutdir+'/'+difxfilename[difxfilename.rfind('/')+1:]
	difxout = open(difxoutname, 'w')
	(vishdr,binhdr) = getVisibilityHeader(difxfile)

	# Parse each visibility entry
	nremoved = 0
	npassed = 0
	while len(vishdr) > 0:

		# Visibility properties
		baseline = vishdr[0]
		freqindex = vishdr[5]
		polpair = vishdr[6]

		# Antenna order as in difx2mark4: ref=ant1="256*nr", rem=ant2="nr%256"
		# since if using opposite order (like in DiFX python utils) get "missing autocorr" complaint from difx2mark4
		# and final pols look wrong; with the below order the baseline pols look correctly swapped (in fourfit, fplot)
		ant2 = baseline % 256
		ant1 = (baseline-ant2)/256

		ant1name = telescopes[ant1-1].name
		ant2name = telescopes[ant2-1].name
		seconds = vishdr[2]

		# Number of channels in this baseband
		nchan = freqs[freqindex].numchan / freqs[freqindex].specavg

		# Read the entire visibility data from disk
		rawvis = difxfile.read(8*nchan)
		if len(rawvis) < 8*nchan:
			break

		# Check if these data should be discarded or can be copied
		if (ant1 == ant2) and (rawvis == '\0'*8*nchan):
			nremoved += 1
		else:
			difxout.write(binhdr)
			difxout.write(rawvis)
			npassed += 1

		(vishdr,binhdr) = getVisibilityHeader(difxfile)

	difxout.close()

	# Finished
	print ('\nDone! Final statistics:')
	print ('    vis. passed through : %d' % (npassed))
	print ('    vis. removed        : %d' % (nremoved))
	print ('\nOutput file:')
	print ('    visbility data     : %s' % (difxoutname))
	print (' ')
def mergeDiFX(basename, antRemove, antKeep):

    # Derive file names
    if basename.endswith(('.difx', '.input', '.calc')):
        basename = basename[:basename.rfind('.')]
    pathless_basename = basename
    if pathless_basename.rfind('/') >= 0:
        pathless_basename = pathless_basename[pathless_basename.rfind('/') +
                                              1:]
    difxoutdir = pathless_basename + '_stripped.difx'

    # Open DiFX files
    difxfilename = glob.glob(basename + '.difx/DIFX_*.s*.b*')[0]
    difxfile = open(difxfilename, 'r')
    try:
        os.mkdir(difxoutdir)
    except:
        pass
    difxoutname = difxoutdir + '/' + difxfilename[difxfilename.rfind('/') + 1:]
    difxout = open(difxoutname, 'w')
    print('Source       : %s' % (difxfilename))
    print('Destination  : %s' % (difxoutname))

    # Extract meta-infos from the DiFX .INPUT file
    inputfile = basename + '.input'
    (numfreqs, freqs) = parseDiFX.get_freqtable_info(inputfile)
    (numtelescopes, telescopes) = parseDiFX.get_telescopetable_info(inputfile)
    (numdatastreams,
     datastreams) = parseDiFX.get_datastreamtable_info(inputfile)
    (numbaselines, baselines) = parseDiFX.get_baselinetable_info(inputfile)
    if numfreqs == 0 or numtelescopes == 0 or numdatastreams == 0 or numbaselines == 0:
        parser.error("Couldn't parse input file " + inputfile + " correctly")

    # Copy/strip
    Ncopied = 0
    Ndropped = 0
    while True:

        (vishdr, binhdr) = getVisibilityHeader(difxfile)
        if len(vishdr) <= 0:
            break

        # Visibility properties
        baseline = vishdr[0]
        seconds = vishdr[2]
        freqindex = vishdr[5]
        polpair = vishdr[6]

        # Antenna order as in difx2mark4: ref=ant1="256*nr", rem=ant2="nr%256"
        ant2 = baseline % 256
        ant1 = (baseline - ant2) / 256
        ant1name = telescopes[ant1 - 1].name.upper()
        ant2name = telescopes[ant2 - 1].name.upper()

        # Number of channels in this baseband
        nchan = freqs[freqindex].numchan / freqs[freqindex].specavg

        # Read the entire visibility data from disk
        rawvis = difxfile.read(8 * nchan)
        if len(rawvis) < 8 * nchan:
            continue

        # Copy to output or discard?
        copy = True
        if (len(antRemove) > 0) and (ant1name in antRemove) or (ant2name
                                                                in antRemove):
            copy = False
        elif (len(antKeep) > 0) and not ((ant1name in antKeep) or
                                         (ant2name in antKeep)):
            copy = False
        #print ('  %s-%s : keep=%s' % (ant1name,ant2name,str(copy)))
        if not copy:
            Ndropped += 1
        else:
            difxout.write(binhdr)
            difxout.write(rawvis)
            Ncopied += 1

    difxout.close()

    # Finished
    print('Vis. dropped : %d' % (Ndropped))
    print('Vis. copied  : %d' % (Ncopied))