Exemple #1
0
    def test_add_noise(self):
        """Test that we can add baseline noise to a data dictionary without
        error"""

        # Setup
        lwa1 = lwa_common.lwa1
        antennas = lwa1.antennas[0:20]
        freqs = numpy.arange(30e6, 50e6, 1e6)
        aa = vis.build_sim_array(lwa1, antennas, freqs)

        # Build the data dictionary
        out = vis.build_sim_data(aa, vis.SOURCES)

        # Add in the noise
        na = vis.add_baseline_noise(out, 15e3, 0.061)

        #
        # Single-channel test
        #

        # Setup
        lwa1 = lwa_common.lwa1
        antennas = lwa1.antennas[0:20]
        freqs = numpy.array([
            30e6,
        ])
        aa = vis.build_sim_array(lwa1, antennas, freqs)

        # Build the data dictionary
        out = vis.build_sim_data(aa, vis.SOURCES)

        # Add in the noise
        na = vis.add_baseline_noise(out, 15e3, 0.061, bandwidth=1e6)

        #
        # VisibilityData test
        #

        # Setup
        lwa1 = lwa_common.lwa1
        antennas = lwa1.antennas[0:20]
        freqs = numpy.arange(30e6, 50e6, 1e6)
        aa = vis.build_sim_array(lwa1, antennas, freqs)

        # Build the data dictionary
        out = vis.build_sim_data(aa, vis.SOURCES)
        out2 = VisibilityData()
        out2.append(out)

        # Add in the noise
        na2 = vis.add_baseline_noise(out2, 15e3, 0.061)
Exemple #2
0
    def test_build_data(self):
        """Test building simulated visibility data"""

        # Setup
        lwa1 = lwa_common.lwa1
        antennas = lwa1.antennas[0:20]
        freqs = numpy.arange(30e6, 50e6, 1e6)
        aa = vis.build_sim_array(lwa1, antennas, freqs)

        # Build the data dictionary
        out = vis.build_sim_data(aa, vis.SOURCES)

        # Do a check of frequencies
        for fa, fq in zip(out.freq, freqs):
            self.assertAlmostEqual(fa, fq, 6)

        # Do a check to make sure that the polarizations
        self.assertEqual(out.npol, 4)
        self.assertTrue('XX' in out.pols)
        self.assertTrue('XY' in out.pols)
        self.assertTrue('YX' in out.pols)
        self.assertTrue('YY' in out.pols)

        # Try a simulation on a single baselines
        aa.sim(0, 1)

        #
        # Single-channel test
        #

        # Setup
        lwa1 = lwa_common.lwa1
        antennas = lwa1.antennas[0:20]
        freqs = numpy.array([
            30e6,
        ])
        aa = vis.build_sim_array(lwa1, antennas, freqs)

        # Build the data dictionary
        out = vis.build_sim_data(aa, vis.SOURCES)

        # Do a check of frequencies
        for fa, fq in zip(out.freq, freqs):
            self.assertAlmostEqual(fa, fq, 6)

        # Do a check to make sure that the polarizations
        self.assertEqual(out.npol, 4)
        self.assertTrue('XX' in out.pols)
        self.assertTrue('XY' in out.pols)
        self.assertTrue('YX' in out.pols)
        self.assertTrue('YY' in out.pols)
Exemple #3
0
    def test_shift_data(self):
        """Test that we can shift the uvw coordinates of a data dictionary 
        without error"""

        # Setup
        lwa1 = lwa_common.lwa1
        antennas = lwa1.antennas[0:20]
        freqs = numpy.arange(30e6, 50e6, 1e6)
        aa = vis.build_sim_array(lwa1, antennas, freqs)

        # Build the data dictionary
        out = vis.build_sim_data(aa, vis.SOURCES)

        # Shift
        sft = vis.shift_data(out, aa)

        #
        # Single-channel test
        #

        # Setup
        lwa1 = lwa_common.lwa1
        antennas = lwa1.antennas[0:20]
        freqs = numpy.array([
            30e6,
        ])
        aa = vis.build_sim_array(lwa1, antennas, freqs)

        # Build the data dictionary
        out = vis.build_sim_data(aa, vis.SOURCES)

        # Shift
        sft = vis.shift_data(out, aa)

        #
        # VisibilityData test
        #

        # Setup
        lwa1 = lwa_common.lwa1
        antennas = lwa1.antennas[0:20]
        freqs = numpy.arange(30e6, 50e6, 1e6)
        aa = vis.build_sim_array(lwa1, antennas, freqs)

        # Build the data dictionary
        out = vis.build_sim_data(aa, vis.SOURCES)
        out2 = VisibilityData()
        out2.append(out)

        # Shift
        sft2 = vis.shift_data(out2, aa)
Exemple #4
0
    def test_build_data_res(self):
        """Test building simulated visibility data with resolved sources"""

        # Setup
        lwa1 = lwa_common.lwa1
        antennas = lwa1.antennas[0:20]
        freqs = numpy.arange(30e6, 50e6, 1e6)
        aa = vis.build_sim_array(lwa1, antennas, freqs)

        # Build the data dictionary
        out = vis.build_sim_data(aa, vis.SOURCES, resolve_src=True)

        # Do a check of keys
        # Do a check of frequencies
        numpy.testing.assert_allclose(out.freq, freqs)

        # Do a check to make sure that the polarizations
        self.assertEqual(out.npol, 4)
        self.assertTrue('XX' in out.pols)
        self.assertTrue('XY' in out.pols)
        self.assertTrue('YX' in out.pols)
        self.assertTrue('YY' in out.pols)

        #
        # Single-channel test
        #

        # Setup
        lwa1 = lwa_common.lwa1
        antennas = lwa1.antennas[0:20]
        freqs = numpy.array([
            30e6,
        ])
        aa = vis.build_sim_array(lwa1, antennas, freqs)

        # Build the data dictionary
        out = vis.build_sim_data(aa, vis.SOURCES, resolve_src=True)

        # Do a check of frequencies
        numpy.testing.assert_allclose(out.freq, freqs)

        # Do a check to make sure that the polarizations
        self.assertEqual(out.npol, 4)
        self.assertTrue('XX' in out.pols)
        self.assertTrue('XY' in out.pols)
        self.assertTrue('YX' in out.pols)
        self.assertTrue('YY' in out.pols)
Exemple #5
0
    def test_plotting(self):
        """Test drawing an image."""

        # Setup
        antennas = lwa1.antennas[0:20]
        freqs = numpy.arange(30e6, 50e6, 1e6)
        aa = vis.build_sim_array(lwa1, antennas, freqs)

        # Build the data dictionary
        out = vis.build_sim_data(aa, vis.SOURCES, jd=2458962.16965)

        # Build an image
        img = utils.build_gridded_image(out)

        # Plot
        fig = plt.figure()
        ax = fig.gca()
        utils.plot_gridded_image(ax, img)
Exemple #6
0
    def test_clean_leastsq(self):
        """Test CLEANing using least squares in the image plane"""

        # Setup
        antennas = lwa1.antennas[0:20]
        freqs = numpy.arange(30e6, 50e6, 1e6)
        aa = vis.build_sim_array(lwa1, antennas, freqs)

        # Build the data dictionary
        out = vis.build_sim_data(aa, vis.SOURCES, jd=2458962.16965)

        with lsl.testing.SilentVerbose():
            # Build an image
            img = utils.build_gridded_image(out)

            # CLEAN
            deconv.lsq(aa,
                       out,
                       img,
                       max_iter=2,
                       verbose=False,
                       plot=run_plotting_tests)
Exemple #7
0
    def test_plotting_graticules(self):
        """Test adding a graticule to an image."""

        # Setup
        antennas = lwa1.antennas[0:20]
        freqs = numpy.arange(30e6, 50e6, 1e6)
        aa = vis.build_sim_array(lwa1, antennas, freqs)

        # Build the data dictionary
        out = vis.build_sim_data(aa, vis.SOURCES, jd=2458962.16965)

        # Build an image
        img = utils.build_gridded_image(out)

        # Plot
        fig = plt.figure()
        ax = fig.gca()
        utils.plot_gridded_image(ax, img)
        with self.subTest(type='RA/Dec.'):
            overlay.graticule_radec(ax, aa)
        with self.subTest(type='az/alt'):
            overlay.graticule_azalt(ax, aa)
        del fig
Exemple #8
0
def main(args):
    filename = args.filename

    idi = utils.CorrelatedData(filename)
    aa = idi.get_antennaarray()
    lo = idi.get_observer()
    lo.date = idi.date_obs.strftime("%Y/%m/%d %H:%M:%S")
    jd = lo.date + astro.DJD_OFFSET
    lst = str(lo.sidereal_time())

    nStand = len(idi.stands)
    nchan = len(idi.freq)
    freq = idi.freq

    print("Raw Stand Count: %i" % nStand)
    print("Final Baseline Count: %i" % (nStand * (nStand - 1) // 2, ))
    print(
        "Spectra Coverage: %.3f to %.3f MHz in %i channels (%.2f kHz/channel)"
        % (freq[0] / 1e6, freq[-1] / 1e6, nchan,
           (freq[-1] - freq[0]) / 1e3 / nchan))
    print("Polarization Products: %i starting with %i" %
          (len(idi.pols), idi.pols[0]))
    print("JD: %.3f" % jd)

    # Pull out something reasonable
    toWork = numpy.where((freq >= args.lower) & (freq <= args.upper))[0]

    print("Reading in FITS IDI data")
    nSets = idi.total_baseline_count // (nStand * (nStand + 1) // 2)
    for set in range(1, nSets + 1):
        print("Set #%i of %i" % (set, nSets))
        fullDict = idi.get_data_set(set)
        dataDict = fullDict.get_uv_range(min_uv=14.0)
        dataDict.sort()

        # Gather up the polarizations and baselines
        pols = dataDict['jd'].keys()
        bls = dataDict['bls'][pols[0]]
        print("The reduced list has %i baselines and %i channels" %
              (len(bls), len(toWork)))

        # Build a list of unique JDs for the data
        jdList = []
        for jd in dataDict['jd'][pols[0]]:
            if jd not in jdList:
                jdList.append(jd)

        # Build the simulated visibilities
        print("Building Model")
        simDict = simVis.build_sim_data(aa,
                                        simVis.SOURCES,
                                        jd=[
                                            jdList[0],
                                        ],
                                        pols=pols,
                                        baselines=bls)

        print("Running self cal.")
        simDict = simDict.sort()
        dataDict = dataDict.sort()
        fixedDataXX, delaysXX = selfcal.delay_only(aa,
                                                   dataDict,
                                                   simDict,
                                                   toWork,
                                                   'xx',
                                                   ref_ant=args.reference,
                                                   max_iter=60)
        fixedDataYY, delaysYY = selfcal.delay_only(aa,
                                                   dataDict,
                                                   simDict,
                                                   toWork,
                                                   'yy',
                                                   ref_ant=args.reference,
                                                   max_iter=60)
        fixedFullXX = simVis.scale_data(fullDict, delaysXX * 0 + 1, delaysXX)
        fixedFullYY = simVis.scale_data(fullDict, delaysYY * 0 + 1, delaysYY)

        print("    Saving results")
        outname = os.path.split(filename)[1]
        outname = os.path.splitext(outname)[0]
        outname = "%s.sc" % outname
        fh = open(outname, 'w')
        fh.write("################################\n")
        fh.write("#                              #\n")
        fh.write("# Columns:                     #\n")
        fh.write("# 1) Stand number              #\n")
        fh.write("# 2) X pol. amplitude          #\n")
        fh.write("# 3) X pol. delay (ns)         #\n")
        fh.write("# 4) Y pol. amplitude          #\n")
        fh.write("# 5) Y pol. delay (ns)         #\n")
        fh.write("#                              #\n")
        fh.write("################################\n")
        for i in xrange(delaysXX.size):
            fh.write("%3i  %.6g  %.6g  %.6g  %.6g\n" %
                     (idi.stands[i], 1.0, delaysXX[i], 1.0, delaysYY[i]))
        fh.close()

        # Build up the images for each polarization
        if args.plot:
            print("    Gridding")
            toWork = numpy.where((freq >= 80e6) & (freq <= 82e6))[0]
            try:
                imgXX = utils.build_gridded_image(fullDict,
                                                  size=80,
                                                  res=0.5,
                                                  pol='xx',
                                                  chan=toWork)
            except:
                imgXX = None

            try:
                imgYY = utils.build_gridded_image(fullDict,
                                                  size=80,
                                                  res=0.5,
                                                  pol='yy',
                                                  chan=toWork)
            except:
                imgYY = None

            try:
                simgXX = utils.build_gridded_image(simDict,
                                                   size=80,
                                                   res=0.5,
                                                   pol='xx',
                                                   chan=toWork)
            except:
                simgXX = None
            try:
                simgYY = utils.build_gridded_image(simDict,
                                                   size=80,
                                                   res=0.5,
                                                   pol='yy',
                                                   chan=toWork)
            except:
                simgYY = None

            try:
                fimgXX = utils.build_gridded_image(fixedFullXX,
                                                   size=80,
                                                   res=0.5,
                                                   pol='xx',
                                                   chan=toWork)
            except:
                fimgXX = None
            try:
                fimgYY = utils.build_gridded_image(fixedFullYY,
                                                   size=80,
                                                   res=0.5,
                                                   pol='yy',
                                                   chan=toWork)
            except:
                fimgYY = None

            # Plots
            print("    Plotting")
            fig = plt.figure()
            ax1 = fig.add_subplot(3, 2, 1)
            ax2 = fig.add_subplot(3, 2, 2)
            ax3 = fig.add_subplot(3, 2, 3)
            ax4 = fig.add_subplot(3, 2, 4)
            ax5 = fig.add_subplot(3, 2, 5)
            ax6 = fig.add_subplot(3, 2, 6)
            for ax, img, pol in zip(
                [ax1, ax2, ax3, ax4, ax5, ax6],
                [imgXX, imgYY, simgXX, simgYY, fimgXX, fimgYY],
                ['XX', 'YY', 'simXX', 'simYY', 'scalXX', 'scalYY']):
                # Skip missing images
                if img is None:
                    ax.text(0.5,
                            0.5,
                            'Not found in file',
                            color='black',
                            size=12,
                            horizontalalignment='center')

                    ax.xaxis.set_major_formatter(NullFormatter())
                    ax.yaxis.set_major_formatter(NullFormatter())

                    ax.set_title("%s @ %s LST" % (pol, lst))
                    continue

                # Display the image and label with the polarization/LST
                out = img.image(center=(80, 80))
                print(pol, out.min(), out.max())
                #if pol == 'scalXX':
                #out = numpy.rot90(out)
                #out = numpy.rot90(out)
                cb = ax.imshow(out,
                               extent=(1, -1, -1, 1),
                               origin='lower',
                               vmin=img.image().min(),
                               vmax=img.image().max())
                fig.colorbar(cb, ax=ax)
                ax.set_title("%s @ %s LST" % (pol, lst))

                # Turn off tick marks
                ax.xaxis.set_major_formatter(NullFormatter())
                ax.yaxis.set_major_formatter(NullFormatter())

                # Compute the positions of major sources and label the images
                compSrc = {}
                ax.plot(0, 0, marker='+', markersize=10, markeredgecolor='w')
                for name, src in simVis.SOURCES.items():
                    src.compute(aa)
                    top = src.get_crds(crdsys='top', ncrd=3)
                    az, alt = aipy.coord.top2azalt(top)
                    compSrc[name] = [az, alt]
                    if alt <= 0:
                        continue
                    ax.plot(top[0],
                            top[1],
                            marker='x',
                            markerfacecolor='None',
                            markeredgecolor='w',
                            linewidth=10.0,
                            markersize=10)
                    ax.text(top[0], top[1], name, color='white', size=12)

                # Add lines of constant RA and dec.
                graticle(ax, lo.sidereal_time(), lo.lat)

            plt.show()

    print("...Done")
Exemple #9
0
    def test_scale_data(self):
        """Test that we can scale a data dictionary without error"""

        # Setup
        lwa1 = lwa_common.lwa1
        antennas = lwa1.antennas[0:20]
        freqs = numpy.arange(30e6, 50e6, 1e6)
        aa = vis.build_sim_array(lwa1, antennas, freqs)

        # Build the data dictionary
        out = vis.build_sim_data(aa, vis.SOURCES)

        # Scale
        amp = vis.scale_data(out,
                             numpy.ones(len(antennas)) * 2,
                             numpy.zeros(len(antennas)))
        # Delay
        phs = vis.scale_data(out, numpy.ones(len(antennas)),
                             numpy.ones(len(antennas)))

        #
        # Single-channel test
        #

        # Setup
        lwa1 = lwa_common.lwa1
        antennas = lwa1.antennas[0:20]
        freqs = numpy.array([
            30e6,
        ])
        aa = vis.build_sim_array(lwa1, antennas, freqs)

        # Build the data dictionary
        out = vis.build_sim_data(aa, vis.SOURCES)

        # Scale
        amp = vis.scale_data(out,
                             numpy.ones(len(antennas)) * 2,
                             numpy.zeros(len(antennas)))
        # Delay
        phs = vis.scale_data(out, numpy.ones(len(antennas)),
                             numpy.ones(len(antennas)))

        #
        # VisibilityData test
        #

        # Setup
        lwa1 = lwa_common.lwa1
        antennas = lwa1.antennas[0:20]
        freqs = numpy.arange(30e6, 50e6, 1e6)
        aa = vis.build_sim_array(lwa1, antennas, freqs)

        # Build the data dictionary
        out = vis.build_sim_data(aa, vis.SOURCES)
        out2 = VisibilityData(out)

        # Scale
        amp2 = vis.scale_data(out2,
                              numpy.ones(len(antennas)) * 2,
                              numpy.zeros(len(antennas)))
        # Delay
        phs2 = vis.scale_data(out2, numpy.ones(len(antennas)),
                              numpy.ones(len(antennas)))
Exemple #10
0
def lsq(aa,
        dataDict,
        aipyImg,
        input_image=None,
        size=80,
        res=0.50,
        wres=0.10,
        pol='XX',
        chan=None,
        gain=0.05,
        max_iter=150,
        rtol=1e-9,
        verbose=True,
        plot=False):
    """
    Given a AIPY antenna array instance, a data dictionary, and an AIPY ImgW 
    instance filled with data, return a deconvolved image.  This function 
    implements a least squares deconvolution.
    
    Least squares tuning parameters:
      * gain - least squares loop gain (default 0.05)
      * max_iter - Maximum number of iteration (default 150)
      * rtol - Minimum change in the residual RMS between iterations
               (default 1e-9)
    """

    # 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(aa.sidereal_time(), aa.lat, center=(size, size))
    ra, dec = eq2radec(xyz)

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

    # 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, :, :]

    # 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

    # Build the initial model
    mdl = img * 0 + img.max()
    mdl[numpy.where(mdl < 0)] = 0
    mdl[numpy.where(ra.mask == 1)] = 0

    # Determine the overall image->model scale factor
    bSrcs = {}
    rChan = [chan[0], chan[-1]]
    bSrcs['zenith'] = RadioFixedBody(aa.sidereal_time(),
                                     aa.lat,
                                     name='zenith',
                                     jys=1,
                                     index=0)
    simDict = build_sim_data(aa,
                             bSrcs,
                             jd=aa.get_jultime(),
                             pols=[
                                 pol,
                             ],
                             chan=rChan,
                             baselines=baselines,
                             flat_response=True)
    simImg = utils.build_gridded_image(simDict,
                                       size=size,
                                       res=res,
                                       wres=wres,
                                       chan=rChan,
                                       pol=pol,
                                       verbose=verbose)
    simImg = simImg.image(center=(size, size))

    simToModel = 1.0 / simImg.max()
    modelToSim = simImg.max() / 1.0

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

        pylab.ion()

    rChan = [chan[0], chan[-1]]
    mdl *= modelToSim
    diff = img - mdl
    diffScaled = 0.0 * diff / gain
    oldModel = mdl
    oldRMS = diff.std() * 1e6
    oldDiff = diff * 0.0
    rHist = []
    exitStatus = 'iteration'
    for k in range(max_iter):
        ## Update the model image but don't allow negative flux
        mdl += diffScaled * gain
        mdl[numpy.where(mdl <= 0)] = 0.0

        ## Convert the model image to an ensemble of point sources for forward
        ## modeling
        bSrcs = {}
        for i in range(mdl.shape[0]):
            for j in range(mdl.shape[1]):
                if dec.mask[i, j]:
                    continue
                if mdl[i, j] <= 0:
                    continue

                nm = '%i-%i' % (i, j)
                bSrcs[nm] = RadioFixedBody(ra[i, j],
                                           dec[i, j],
                                           name=nm,
                                           jys=mdl[i, j],
                                           index=0,
                                           epoch=aa.date)

        ## Model the visibilities
        simDict = build_sim_data(aa,
                                 bSrcs,
                                 jd=aa.get_jultime(),
                                 pols=[
                                     pol,
                                 ],
                                 chan=rChan,
                                 baselines=baselines,
                                 flat_response=True)

        ## Form the simulated image
        simImg = utils.build_gridded_image(simDict,
                                           size=size,
                                           res=res,
                                           wres=wres,
                                           chan=rChan,
                                           pol=pol,
                                           verbose=verbose)
        simImg = simImg.image(center=(size, size))

        ## Difference the image and the simulated image and scale it to the
        ## model's peak flux
        diff = img - simImg
        diff2 = _minor_cycle(diff, beamClean, gain=0.1, max_iter=2000)

        ## Compute the RMS and create an appropriately scaled version of the model
        RMS = diff.std()
        mdl2 = mdl * modelToSim

        ## Status report
        if verbose:
            print("Iteration %i:  %i sources used, RMS is %.4e" %
                  (k + 1, len(bSrcs.keys()), RMS))
            print("               -> maximum residual: %.4e (%.3f%% of peak)" %
                  (diff.max(), 100.0 * diff.max() / img.max()))
            print("               -> minimum residual: %.4e (%.3f%% of peak)" %
                  (diff.min(), 100.0 * diff.min() / img.max()))
            print("               -> delta RMS: %.4e (%.3f%%)" %
                  (RMS - oldRMS, 100.0 * (RMS - oldRMS) / RMS))

        ## Make the cleaned residuals map ready for updating the model
        diff = diff2
        diffScaled = diff * simToModel

        ## Has the RMS gone up?  If so, it is time to exit.  But first, restore
        ## the previous iteration
        if RMS - oldRMS > 0:
            mdl = oldModel
            diff = oldDiff

            exitStatus = 'residuals'

            break

        ## Is the RMS still changing in an acceptable manner?
        if abs(RMS - oldRMS) < rtol:
            # No need to go back a step
            #mdl = oldModel
            #diff = oldDiff

            exitStatus = 'tolerance'

            break

        ## Save the current iteration as the previous state
        rHist.append(RMS)
        oldRMS = RMS
        oldModel = mdl
        oldDiff = diff

        if plot:
            pylab.subplot(3, 2, 1)
            pylab.imshow(img,
                         origin='lower',
                         interpolation='nearest',
                         vmin=img.min(),
                         vmax=img.max())
            pylab.subplot(3, 2, 2)
            pylab.imshow(simImg,
                         origin='lower',
                         interpolation='nearest',
                         vmin=img.min(),
                         vmax=img.max())
            pylab.subplot(3, 2, 3)
            pylab.imshow(diff, origin='lower', interpolation='nearest')
            pylab.subplot(3, 2, 4)
            pylab.imshow(mdl, origin='lower', interpolation='nearest')
            pylab.subplot(3, 1, 3)
            pylab.cla()
            pylab.plot(rHist)
            pylab.draw()

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

    # Restore
    conv = convolve(mdl2, beamClean, mode='same')
    conv = numpy.ma.array(conv, mask=convMask)

    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(simImg,
                       extent=(1, -1, -1, 1),
                       origin='lower',
                       interpolation='nearest')
        fig.colorbar(d, ax=ax2)
        ax2.set_title('Realized Model')

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

        f = ax4.imshow(conv + diff,
                       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 conv + diff
Exemple #11
0
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
ovro = ovro2

# Simulation setup
nant = len(ovro.antennas) // 2
nbl = nant*(nant+1)//2
chan0 = 1234
nchan = 192
CHAN_BW = 196e6 / 8192
jd = AstroTime.now().jd

# Simulation array
freqs = (chan0 + numpy.arange(nchan)) * CHAN_BW + CHAN_BW/2
aa = simVis.build_sim_array(ovro, ovro.antennas[0::2], freqs/1e9, jd=jd)

# Simulate with bright sources only
dataSet = simVis.build_sim_data(aa, simVis.SOURCES, pols=['xx','yy'], jd=jd)

# Make the final data set that can be used with dr_visibilities.py
# NOTE:  XY and XY are 1% of XX and have sign flip between XY and YX
vis = numpy.zeros((nbl,nchan,4), dtype=numpy.complex64)
k = 0
l = 0
for i in range(nant):
    for j in range(i, nant):
        vis[l,:,0] = dataSet.XX.data[k,:]
        vis[l,:,1] = dataSet.XX.data[k,:]* 0.01
        vis[l,:,2] = dataSet.XX.data[k,:]*-0.01
        vis[l,:,3] = dataSet.YY.data[k,:]
        if i == j:
            vis[l,:,:] = vis[l,:,:].real
        else:
def main(args):
    filenames = args.filenames
    filenames.sort()

    times = []
    for filename in filenames:
        dataDict = numpy.load(filename)

        tStart = datetime.utcfromtimestamp(dataDict['tStart'])
        tInt = dataDict['tInt']
        try:
            srate = dataDict['srate']
        except KeyError:
            srate = 19.6e6
        
        freq1 = dataDict['freq1']
        freq2 = dataDict['freq2']
        stand1, stand2 = dataDict['stands']

        times.append( tStart)

    print("Got %i files from %s to %s (%s)" % (len(filenames), times[0].strftime("%Y/%m/%d %H:%M:%S"), times[-1].strftime("%Y/%m/%d %H:%M:%S"), (times[-1]-times[0])))

    iTimes = []
    for i in xrange(1, len(times)):
        dt = times[i] - times[i-1]
        iTimes.append(dt.days*24*3600 + dt.seconds + dt.microseconds/1e6)
    iTimes = numpy.array(iTimes)
    print(" -> Interval: %.3f +/- %.3f seconds (%.3f to %.3f seconds)" % (iTimes.mean(), iTimes.std(), iTimes.min(), iTimes.max()))
    
    print("Number of frequency channels: %i (~%.1f Hz/channel)" % (len(freq1)+1, freq1[1]-freq1[0]))

    # Build up the station
    if args.lwasv:
        site = stations.lwasv
    else:
        site = stations.lwa1
    
    rawAntennas = site.antennas
    
    antennas = []
    for ant in rawAntennas:
        if ant.stand.id == stand1 and ant.pol == 0:
            antennas.append(ant)
    for ant in rawAntennas:
        if ant.stand.id == stand2 and ant.pol == 0:
            antennas.append(ant)
    if len(antennas) != 2:
        raise RuntimeError("Can only find stand %i, %i and %i found in the NPZ files" % (antennas[0].stand.id, stand1, stand2))

    # Create the simulated array
    refJD = unix_to_utcjd(timegm(times[0].timetuple()))
    aa1 = simVis.build_sim_array(site, antennas, freq1/1e9, jd=refJD)
    aa2 = simVis.build_sim_array(site, antennas, freq2/1e9, jd=refJD)

    # Build the model times and range.
    jdList = []
    dTimes = []
    for i in xrange(len(times)):
        tNow = timegm(times[i].timetuple())
        jdNow = unix_to_utcjd(tNow)

        jdList.append(jdNow)
        dTimes.append( (times[i]-times[0]).seconds )
        
    # Actually run the simulations
    simDict1 = simVis.build_sim_data(aa1, simVis.SOURCES, jd=jdList, pols=['xx',], verbose=False)
    simDict2 = simVis.build_sim_data(aa2, simVis.SOURCES, jd=jdList, pols=['xx',], verbose=False)

    # Plot
    fig = plt.figure()
    ax1 = fig.add_subplot(2, 1, 1)
    ax2 = fig.add_subplot(2, 1, 2)

    vis1 = []
    for vis in simDict1['vis']['xx']:
        vis1.append( vis )
    vis2 = []
    for vis in simDict2['vis']['xx']:
        vis2.append( vis )
    
    vis1 = numpy.array(vis1)
    vis1 = numpy.ma.array(vis1, mask=~numpy.isfinite(vis1))
    vis2 = numpy.array(vis2)
    vis2 = numpy.ma.array(vis2, mask=~numpy.isfinite(vis2))

    data = numpy.abs(vis1)
    data = data.ravel()
    data.sort()
    vmin1 = data[int(round(0.15*len(data)))]
    vmax1 = data[int(round(0.85*len(data)))]
    print('Plot range for tuning 1:', vmin1, vmax1)
    
    data = numpy.abs(vis2)
    data = data.ravel()
    data.sort()
    vmin2 = data[int(round(0.15*len(data)))]
    vmax2 = data[int(round(0.85*len(data)))]
    print('Plot range for tuning 2:', vmin2, vmax2)

    ax1.imshow(numpy.abs(vis1), extent=(freq1[0], freq1[-1], dTimes[0], dTimes[-1]), origin='lower', 
            vmin=vmin1, vmax=vmax1)
    ax2.imshow(numpy.abs(vis2), extent=(freq1[0], freq1[-1], dTimes[0], dTimes[-1]), origin='lower', 
            vmin=vmin2, vmax=vmax2)

    ax1.axis('auto')
    ax2.axis('auto')

    fig.suptitle("%s to %s UTC" % (times[0].strftime("%Y/%m/%d %H:%M"), times[-1].strftime("%Y/%m/%d %H:%M")))
    ax1.set_xlabel('Frequency [MHz]')
    ax2.set_xlabel('Frequency [MHz]')
    ax1.set_ylabel('Elapsed Time [s]')
    ax2.set_ylabel('Elapsed Time [s]')

    plt.show()