示例#1
0
def _get_source_parameters(aa, timestamp, srcs):
    """
    Given an aipy AntennaArray object, an observation time, and aipy.src 
    object, return all of the parameters needed for a simulation.
    """

    # Set the time for the array
    aa.set_unixtime(timestamp)

    # Compute the source parameters
    srcs_tp = []
    srcs_mt = []
    srcs_jy = []
    srcs_fq = []
    for name in srcs:
        ## Update the source's coordinates
        src = srcs[name]
        src.compute(aa)

        ## Get parameters
        top = src.get_crds(crdsys='top', ncrd=3)  # topo. coords.
        mat = src.map  # equitorial -> topo. rotation matrix
        jys = src.get_jys()  # F_nu
        frq = aa.get_afreqs()  # nu

        ## Fix the lowest frequencies to avoid problems with the flux blowing up
        ## at nu = 0 Hz by replacing flux values below 1 MHz with the flux at
        ## 1 MHz
        Jyat1MHz = jys[numpy.where(
            numpy.abs(frq - 0.001) == numpy.abs(frq - 0.001).min())]
        jys = numpy.where(frq >= 0.001, jys, Jyat1MHz)

        ## Filter out sources that are below the horizon or have no flux
        srcAzAlt = aipycoord.top2azalt(top)
        if srcAzAlt[1] <= 0 or jys.sum() <= 0:
            continue

        ## Save values into the source arrays
        srcs_tp.append(top)
        srcs_mt.append(mat)
        srcs_jy.append(jys)
        srcs_fq.append(frq)

    # Return the values as a dictionary
    return {
        'topo': srcs_tp,
        'trans': srcs_mt,
        'flux': srcs_jy,
        'freq': srcs_fq
    }
示例#2
0
文件: deconv.py 项目: lwa-project/lsl
def clean_sources(aa,
                  dataDict,
                  aipyImg,
                  srcs,
                  input_image=None,
                  size=80,
                  res=0.50,
                  wres=0.10,
                  pol='XX',
                  chan=None,
                  gain=0.1,
                  max_iter=150,
                  sigma=2.0,
                  verbose=True,
                  plot=False):
    """
    Given a AIPY antenna array instance, a data dictionary, an AIPY ImgW 
    instance filled with data, and a dictionary of sources, return the CLEAN
    components and the residuals map.  This function uses a CLEAN-like method
    that computes the array beam for each peak in the flux.  Thus the CLEAN 
    loop becomes: 
      1.  Find the peak flux in the residual image
      2.  Compute the systems response to a point source at that location
      3.  Remove the scaled porition of this beam from the residuals
      4.  Go to 1.
    
    This function differs from clean() in that it only cleans localized 
    regions around each source rather than the whole image.  This is
    intended to help the mem() function along.
    
    CLEAN tuning parameters:
      * gain - CLEAN loop gain (default 0.1)
      * max_iter - Maximum number of iterations (default 150)
      * sigma - Threshold in sigma to stop cleaning (default 2.0)
    """

    # Sort out the channels to work on
    if chan is None:
        chan = range(dataDict.freq.size)

    # Get a grid of right ascensions and dec values for the image we are working with
    xyz = aipyImg.get_eq(0.0, aa.lat, center=(size, size))
    RA, dec = eq2radec(xyz)
    RA += aa.sidereal_time()
    RA %= (2 * numpy.pi)
    top = aipyImg.get_top(center=(size, size))
    az, alt = top2azalt(top)

    # Get the list of baselines to generate visibilites for
    baselines = dataDict.baselines

    # Get the actual image out of the ImgW instance
    if input_image is None:
        img = aipyImg.image(center=(size, size))
    else:
        img = input_image * 1.0

    # Setup the arrays to hold the point sources and the residual.
    cleaned = numpy.zeros_like(img)
    working = numpy.zeros_like(img)
    working += img

    # Setup the dictionary that will hold the beams as they are computed
    prevBeam = {}

    # Estimate the zenith beam response
    psfSrc = {
        'z':
        RadioFixedBody(aa.sidereal_time(),
                       aa.lat,
                       jys=1.0,
                       index=0,
                       epoch=aa.date)
    }
    psfDict = build_sim_data(aa,
                             psfSrc,
                             jd=aa.get_jultime(),
                             pols=[
                                 pol,
                             ],
                             chan=chan,
                             baselines=baselines,
                             flat_response=True)
    psf = utils.build_gridded_image(psfDict,
                                    size=size,
                                    res=res,
                                    wres=wres,
                                    chan=chan,
                                    pol=pol,
                                    verbose=verbose)
    psf = psf.image(center=(size, size))
    psf /= psf.max()

    # Fit a Guassian to the zenith beam response and use that for the restore beam
    beamCutout = psf[size // 2:3 * size // 2, size // 2:3 * size // 2]
    beamCutout = numpy.where(beamCutout > 0.0, beamCutout, 0.0)
    h, cx, cy, sx, sy = _fit_gaussian(beamCutout)
    gauGen = gaussian2d(1.0, size / 2 + cx, size / 2 + cy, sx, sy)
    FWHM = int(round((sx + sy) / 2.0 * 2.0 * numpy.sqrt(2.0 * numpy.log(2.0))))
    beamClean = psf * 0.0
    for i in range(beamClean.shape[0]):
        for j in range(beamClean.shape[1]):
            beamClean[i, j] = gauGen(i, j)
    beamClean /= beamClean.sum()
    convMask = xyz.mask[0, :, :]

    # Go!
    if plot:
        import pylab
        from matplotlib import pyplot as plt

        pylab.ion()

    for name, src in srcs.items():
        # Make sure the source is up
        src.compute(aa)
        if verbose:
            print('Source: %s @ %s degrees elevation' % (name, src.alt))
        if src.alt <= 10 * numpy.pi / 180.0:
            continue

        # Locate the approximate position of the source
        srcDist = (src.ra - RA)**2 + (src.dec - dec)**2
        srcPeak = numpy.where(srcDist == srcDist.min())

        # Define the clean box - this is fixed at 2*FWHM in width on each side
        rx0 = max([0, srcPeak[0][0] - FWHM // 2])
        rx1 = min([rx0 + FWHM + 1, img.shape[0]])
        ry0 = max([0, srcPeak[1][0] - FWHM // 2])
        ry1 = min([ry0 + FWHM + 1, img.shape[1]])

        # Define the background box - this lies outside the clean box and serves
        # as a reference for the background
        X, Y = numpy.indices(working.shape)
        R = numpy.sqrt((X - srcPeak[0][0])**2 + (Y - srcPeak[1][0])**2)
        bpad = 3
        background = numpy.where((R <= FWHM + bpad) & (R > FWHM))
        while len(background[0]) == 0 and bpad < img.shape[0]:
            bpad += 1
            background = numpy.where((R <= FWHM + bpad) & (R > FWHM))

        px0 = min(background[0]) - 1
        px1 = max(background[0]) + 2
        py0 = min(background[1]) - 1
        py1 = max(background[1]) + 2

        exitStatus = 'iteration'
        for i in range(max_iter):
            # Find the location of the peak in the flux density
            peak = numpy.where(working[rx0:rx1,
                                       ry0:ry1] == working[rx0:rx1,
                                                           ry0:ry1].max())
            peak_x = peak[0][0] + rx0
            peak_y = peak[1][0] + ry0
            peakV = working[peak_x, peak_y]

            # Optimize the location
            try:
                peakParams = _fit_gaussian(
                    working[peak_x - FWHM // 2:peak_x + FWHM // 2 + 1,
                            peak_y - FWHM // 2:peak_y + FWHM // 2 + 1])
            except IndexError:
                peakParams = [peakV, FWHM // 2, FWHM // 2]
            peakVO = peakParams[0]
            peak_xO = peak_x - FWHM // 2 + peakParams[1]
            peak_yO = peak_y - FWHM // 2 + peakParams[2]

            # Quantize to try and keep the computation down without over-simplifiying things
            subpixelationLevel = 5
            peak_xO = round(
                peak_xO * subpixelationLevel) / float(subpixelationLevel)
            peak_yO = round(
                peak_yO * subpixelationLevel) / float(subpixelationLevel)

            # Pixel coordinates to right ascension, dec.
            try:
                peakRA = _interpolate(RA, peak_xO, peak_yO)
            except IndexError:
                peak_xO, peak_yO = peak_x, peak_y
                peakRA = RA[peak_x, peak_y]
            try:
                peakDec = _interpolate(dec, peak_xO, peak_yO)
            except IndexError:
                peakDec = dec[peak_x, peak_y]

            # Pixel coordinates to az, el
            try:
                peakAz = _interpolate(az, peak_xO, peak_yO)
            except IndexError:
                peak_xO, peak_yO = peak_x, peak_y
                peakAz = az[peak_x, peak_y]
            try:
                peakEl = _interpolate(alt, peak_x, peak_y)
            except IndexError:
                peakEl = alt[peak_x, peak_y]

            if verbose:
                currRA = deg_to_hms(peakRA * 180 / numpy.pi)
                currDec = deg_to_dms(peakDec * 180 / numpy.pi)
                currAz = deg_to_dms(peakAz * 180 / numpy.pi)
                currEl = deg_to_dms(peakEl * 180 / numpy.pi)

                print(
                    "%s - Iteration %i:  Log peak of %.3f at row: %i, column: %i"
                    % (name, i + 1, numpy.log10(peakV), peak_x, peak_y))
                print("               -> RA: %s, Dec: %s" % (currRA, currDec))
                print("               -> az: %s, el: %s" % (currAz, currEl))

            # Check for the exit criteria
            if peakV < 0:
                exitStatus = 'peak value is negative'

                break

            # Find the beam index and see if we need to compute the beam or not
            beamIndex = (int(peak_xO * subpixelationLevel),
                         int(peak_yO * subpixelationLevel))
            try:
                beam = prevBeam[beamIndex]

            except KeyError:
                if verbose:
                    print("               -> Computing beam(s)")

                beamSrc = {
                    'Beam':
                    RadioFixedBody(peakRA,
                                   peakDec,
                                   jys=1.0,
                                   index=0,
                                   epoch=aa.date)
                }
                beamDict = build_sim_data(aa,
                                          beamSrc,
                                          jd=aa.get_jultime(),
                                          pols=[
                                              pol,
                                          ],
                                          chan=chan,
                                          baselines=baselines,
                                          flat_response=True)
                beam = utils.build_gridded_image(beamDict,
                                                 size=size,
                                                 res=res,
                                                 wres=wres,
                                                 chan=chan,
                                                 pol=pol,
                                                 verbose=verbose)
                beam = beam.image(center=(size, size))
                beam /= beam.max()
                if verbose:
                    print("                  ", beam.mean(), beam.min(),
                          beam.max(), beam.sum())

                prevBeam[beamIndex] = beam
                if verbose:
                    print("               -> Beam cache contains %i entries" %
                          len(prevBeam.keys()))

            # Calculate how much signal needs to be removed...
            toRemove = gain * peakV * beam
            working -= toRemove
            asum = 0.0
            for l in range(int(peak_xO), int(peak_xO) + 2):
                if l > peak_xO:
                    side1 = (peak_xO + 0.5) - (l - 0.5)
                else:
                    side1 = (l + 0.5) - (peak_xO - 0.5)

                for m in range(int(peak_yO), int(peak_yO) + 2):
                    if m > peak_yO:
                        side2 = (peak_yO + 0.5) - (m - 0.5)
                    else:
                        side2 = (m + 0.5) - (peak_yO - 0.5)

                    area = side1 * side2
                    asum += area
                    #print('II', l, m, area, asum)
                    cleaned[l, m] += gain * area * peakV

            if plot:
                try:
                    pylab.subplot(2, 2, 1)
                    pylab.imshow((working + toRemove)[px0:px1, py0:py1],
                                 origin='lower',
                                 interpolation='nearest')
                    pylab.title('Before')

                    pylab.subplot(2, 2, 2)
                    pylab.imshow(working[px0:px1, py0:py1],
                                 origin='lower',
                                 interpolation='nearest')
                    pylab.title('After')

                    pylab.subplot(2, 2, 3)
                    pylab.imshow(toRemove[px0:px1, py0:py1],
                                 origin='lower',
                                 interpolation='nearest')
                    pylab.title('Removed')

                    pylab.subplot(2, 2, 4)
                    pylab.imshow(convolve(cleaned, beamClean,
                                          mode='same')[px0:px1, py0:py1],
                                 origin='lower',
                                 interpolation='nearest')
                    pylab.title('CLEAN Comps.')
                except:
                    pass

                try:
                    st.set_text('%s @ %i' % (name, i + 1))
                except NameError:
                    st = pylab.suptitle('%s @ %i' % (name, i + 1))
                pylab.draw()

            if numpy.abs(
                    numpy.max(working[rx0:rx1, ry0:ry1]) -
                    numpy.median(working[background])) / rStd(
                        working[background]) <= sigma:
                exitStatus = 'peak is less than %.3f-sigma' % sigma

                break

        # Summary
        print("Exited after %i iterations with status '%s'" %
              (i + 1, exitStatus))

    # Restore
    conv = convolve(cleaned, beamClean, mode='same')
    conv = numpy.ma.array(conv, mask=convMask)
    conv *= ((img - working).max() / conv.max())

    if plot:
        # Make an image for comparison purposes if we are verbose
        fig = plt.figure()
        ax1 = fig.add_subplot(2, 2, 1)
        ax2 = fig.add_subplot(2, 2, 2)
        ax3 = fig.add_subplot(2, 2, 3)
        ax4 = fig.add_subplot(2, 2, 4)

        c = ax1.imshow(img,
                       extent=(1, -1, -1, 1),
                       origin='lower',
                       interpolation='nearest')
        fig.colorbar(c, ax=ax1)
        ax1.set_title('Input')

        d = ax2.imshow(conv,
                       extent=(1, -1, -1, 1),
                       origin='lower',
                       interpolation='nearest')
        fig.colorbar(d, ax=ax2)
        ax2.set_title('CLEAN Comps.')

        e = ax3.imshow(working,
                       extent=(1, -1, -1, 1),
                       origin='lower',
                       interpolation='nearest')
        fig.colorbar(e, ax=ax3)
        ax3.set_title('Residuals')

        f = ax4.imshow(conv + working,
                       extent=(1, -1, -1, 1),
                       origin='lower',
                       interpolation='nearest')
        fig.colorbar(f, ax=ax4)
        ax4.set_title('Final')

        plt.show()

    if plot:
        pylab.ioff()

    # Return
    return conv, working