RFI Signals

#Takes a UV object and injects randomly distributed RFI signals into the data array of object

#applyRFI should take in a dict with keys narrowband,scatter,burst
    #scatter: mu,sigma, max_times, max_freq, points
    #narrowband: mu, sigma, times, freq
    #burst: mu, sigma, ???

RFIdict = {

    'narrowband' : {

        'mu' : 1e6,
        'sigma' : 1e9,
        'freq' : 900,
        'max_times' : 60

    },
    'scatter' : {

        'mu' : 1e6,
        'sigma' : 1e9,
        'max_times' : 60,
        'max_freqs' : 1024,
        'points' : 30

    },
    'burst' : {

        'mu' : 1e6,
        'sigma' : 1e9,
        'freq' : 500,
        'max_times' : 30

    }
}


def scatterRFI(baseline_data, RFIdict):

    max_times = RFIdict.get('scatter').get('max_times')
    max_freq = RFIdict.get('scatter').get('max_freqs')
    mu = RFIdict.get('scatter').get('mu')
    sigma = RFIdict.get('scatter').get('sigma')
    points = RFIdict.get('scatter').get('points')

    #Defines range of baseline timestamps an frequencies
    times = np.random.randint(0,max_times,size=points)
    freq = np.random.randint(0,max_freq,size=points)

    #Loops for each point to be adjusted
    #print times
    for i in range(points):

        #Increasing randomly distributed RFI signals
        baseline_data[times[i],freq[i]] += np.random.normal(loc=mu,scale=sigma)

    return baseline_data

def narrowbandRFI(baseline_data, RFIdict):

    #get freq of narrowband signal
    freq = RFIdict.get('narrowband').get('freq')
    max_times = RFIdict.get('narrowband').get('max_times')
    mu = RFIdict.get('narrowband').get('mu')
    sigma = RFIdict.get('narrowband').get('sigma')

    #iterate through ALL times and apply narrowband RFI to baseline
    for i in range(max_times):

        #Increasing randomly distributed RFI signals
        baseline_data[i,freq] += np.random.normal(loc=mu,scale=sigma)

    return baseline_data

def burstRFI(baseline_data, RFIdict):

    freq = RFIdict.get('burst').get('freq')
    max_times = RFIdict.get('burst').get('max_times')
    mu = RFIdict.get('burst').get('mu')
    sigma = RFIdict.get('burst').get('sigma')
    times = np.random.randint(0,max_times)
    time_offset = np.random.randint(0,max_times)
    max_taper = (times/2) if (times >= 2) else (1)
    taper_length = np.random.randint(0,max_taper)
    center_strength = np.random.randint(1e0,1e2)

    #iterate through ALL times and apply burst RFI to baseline
    for i in range(times):
        if(i < taper_length):

            baseline_data[i + time_offset,freq] = (center_strength * ((i+1) / float(taper_length +1)))
            #print (i+1), (taper_length), baseline_data[i + time_offset,freq]
            #print center_strength, (i+1), (taper_length+1), (center_strength * ((i+1) / (taper_length+1)))
        elif(times - i <= taper_length):
            baseline_data[i + time_offset,freq] = (center_strength * ((times - i+1) / float(taper_length+1)))
        else:
            baseline_data[i + time_offset,freq] = center_strength

    return baseline_data

def applyRFI(uv_object, RFIdict):

    #Looping through all antenna pairs
    for cur_pair in uv_object.get_antpairs():
        #Getting baseline data for antenna pair
        baseline_data = uv_object.get_data(cur_pair, force_copy = 'true')
        #Zero
        baseline_data = np.zeros_like(baseline_data)

        if RFIdict.has_key('scatter'):
            baseline_data = scatterRFI(baseline_data, RFIdict)

        if RFIdict.has_key('narrowband'):
            baseline_data = narrowbandRFI(baseline_data, RFIdict)

        if RFIdict.has_key('burst'):
            baseline_data = burstRFI(baseline_data, RFIdict)

        antpair_indices = uv_object.antpair2ind(cur_pair[0],cur_pair[1])
        uv_object.data_array[antpair_indices,0,:,0] = copy.copy(baseline_data)

    return uv_object

new_RFI1 = applyRFI(UV1, RFIdict)

#new_RFI1.write_miriad('RFIrandom20')

plt.figure(figsize=(20,20))
plt.imshow(
    np.abs(
        new_RFI1.data_array[np.where(
            new_RFI1.baseline_array == new_RFI1.antnums_to_baseline(0,2)
        ),0,:,0][0]), aspect='auto', norm = matplotlib.colors.LogNorm(1e0,1e2))
plt.colorbar()
plt.xlabel('Frequency')
plt.ylabel('Time (Integrations)')