def test_average_chunks2_2d(self):
     arr = numpy.linspace(0.0, 120.0, 121).reshape(11, 11)
     wts = numpy.ones_like(arr)
     carr, cwts = average_chunks2(arr, wts, (5, 2))
     assert len(carr) == len(cwts)
     answerarr = numpy.array([32., 87., 120.])
     answerwts = numpy.array([5., 5., 1.])
     numpy.testing.assert_array_equal(carr[:, 5], answerarr)
     numpy.testing.assert_array_equal(cwts[:, 5], answerwts)
 def test_average_chunks2_1d_trans(self):
     arr = numpy.linspace(0.0, 100.0, 11).reshape([11, 1])
     wts = numpy.ones_like(arr)
     carr, cwts = average_chunks2(arr, wts, (2, 1))
     assert len(carr) == len(cwts)
     answerarr = numpy.array([[5.], [25.], [45.], [65.0], [85.0], [100.0]])
     answerwts = numpy.array([[2.0], [2.0], [2.0], [2.0], [2.0], [1.0]])
     numpy.testing.assert_array_equal(carr, answerarr)
     numpy.testing.assert_array_equal(cwts, answerwts)
 def test_average_chunks2_1d(self):
     arr = numpy.linspace(0.0, 100.0, 11).reshape([1, 11])
     wts = numpy.ones_like(arr)
     carr, cwts = average_chunks2(arr, wts, (1, 2))
     assert len(carr) == len(cwts)
     answerarr = numpy.array([[5., 25., 45., 65.0, 85.0, 100.0]])
     answerwts = numpy.array([[2.0, 2.0, 2.0, 2.0, 2.0, 1.0]])
     numpy.testing.assert_array_equal(carr, answerarr)
     numpy.testing.assert_array_equal(cwts, answerwts)
Exemple #4
0
 def sum_from_grid(arr):
     result = average_chunks2(
         arr, allpwtsgrid[:, a2, a1, :],
         (time_average[a2, a1], frequency_average[a2, a1]))
     return result[0] * result[0].size
Exemple #5
0
 def average_from_grid(arr):
     return average_chunks2(
         arr, allpwtsgrid[:, a2, a1, :],
         (time_average[a2, a1], frequency_average[a2, a1]))[0]
Exemple #6
0
def average_in_blocks(vis,
                      uvw,
                      wts,
                      times,
                      integration_time,
                      frequency,
                      channel_bandwidth,
                      time_coal=1.0,
                      max_time_coal=100,
                      frequency_coal=1.0,
                      max_frequency_coal=100):
    # Calculate the averaging factors for time and frequency making them the same for all times
    # for this baseline
    # Find the maximum possible baseline and then scale to this.

    # The input visibility is a block of shape [ntimes, nant, nant, nchan, npol]. We will map this
    # into rows like vis[npol] and with additional columns antenna1, antenna2, frequency

    ntimes, nant, _, nchan, npol = vis.shape

    # Pol independent weighting
    allpwtsgrid = numpy.sum(wts, axis=4)
    # Pol and frequency independent weighting
    allcpwtsgrid = numpy.sum(allpwtsgrid, axis=3)
    # Pol and time independent weighting
    alltpwtsgrid = numpy.sum(allpwtsgrid, axis=0)

    # Now calculate on a baseline basis the time and frequency averaging. We do this by looking at
    # the maximum uv distance for all data and for a given baseline. The integration time and
    # channel bandwidth are scale appropriately.
    uvmax = numpy.sqrt(numpy.max(uvw[:, 0]**2 + uvw[:, 1]**2 + uvw[:, 2]**2))
    time_average = numpy.ones([nant, nant], dtype='int')
    frequency_average = numpy.ones([nant, nant], dtype='int')
    ua = numpy.arange(nant)
    for a2 in ua:
        for a1 in ua:
            if allpwtsgrid[:, a2, a1, :].any() > 0.0:
                uvdist = numpy.max(numpy.sqrt(uvw[:, a2, a1, 0]**2 +
                                              uvw[:, a2, a1, 1]**2),
                                   axis=0)
                if uvdist > 0.0:
                    time_average[a2, a1] = min(
                        max_time_coal,
                        max(1, int(round((time_coal * uvmax / uvdist)))))
                    frequency_average[a2, a1] = min(
                        max_frequency_coal,
                        max(1, int(round(frequency_coal * uvmax / uvdist))))
                else:
                    time_average[a2, a1] = max_time_coal
                    frequency_average[a2, a1] = max_frequency_coal

    # See how many time chunks and frequency we need for each baseline. To do this we use the same averaging that
    # we will use later for the actual data. This tells us the number of chunks required for each baseline.
    frequency_grid, time_grid = numpy.meshgrid(frequency, times)
    channel_bandwidth_grid, integration_time_grid = numpy.meshgrid(
        channel_bandwidth, integration_time)
    cnvis = 0
    time_chunk_len = numpy.ones([nant, nant], dtype='int')
    frequency_chunk_len = numpy.ones([nant, nant], dtype='int')
    for a2 in ua:
        for a1 in ua:
            if (time_average[a2, a1] >
                    0) & (frequency_average[a2, a1] > 0 &
                          (allpwtsgrid[:, a2, a1, ...].any() > 0.0)):
                time_chunks, _ = average_chunks(times, allcpwtsgrid[:, a2, a1],
                                                time_average[a2, a1])
                time_chunk_len[a2, a1] = time_chunks.shape[0]
                frequency_chunks, _ = average_chunks(frequency,
                                                     alltpwtsgrid[a2, a1, :],
                                                     frequency_average[a2, a1])
                frequency_chunk_len[a2, a1] = frequency_chunks.shape[0]
                nrows = time_chunk_len[a2, a1] * frequency_chunk_len[a2, a1]
                cnvis += nrows

    # Now we know enough to define the output coalesced arrays. The shape will be
    # succesive a1, a2: [len_time_chunks[a2,a1], a2, a1, len_frequency_chunks[a2,a1]]
    ctime = numpy.zeros([cnvis])
    cfrequency = numpy.zeros([cnvis])
    cchannel_bandwidth = numpy.zeros([cnvis])
    cvis = numpy.zeros([cnvis, npol], dtype='complex')
    cwts = numpy.zeros([cnvis, npol])
    cuvw = numpy.zeros([cnvis, 3])
    ca1 = numpy.zeros([cnvis], dtype='int')
    ca2 = numpy.zeros([cnvis], dtype='int')
    cintegration_time = numpy.zeros([cnvis])

    # For decoalescence we keep an index to map back to the original BlockVisibility
    rowgrid = numpy.zeros([ntimes, nant, nant, nchan], dtype='int')
    rowgrid.flat = range(rowgrid.size)

    cindex = numpy.zeros([rowgrid.size], dtype='int')

    # Now go through, chunking up the various arrays. Everything is converted into an array with
    # axes [time, channel] and then it is averaged over time and frequency chunks for
    # this baseline.
    # To aid decoalescence we will need an index of which output elements a given input element
    # contributes to. This is a many to one. The decoalescence will then just consist of using
    # this index to extract the coalesced value that a given input element contributes towards.

    visstart = 0
    for a2 in ua:
        for a1 in ua:
            if (time_chunk_len[a2, a1] > 0) & (frequency_chunk_len[a2, a1] > 0) & \
                    (allpwtsgrid[:, a2, a1, :].any() > 0.0):

                nrows = time_chunk_len[a2, a1] * frequency_chunk_len[a2, a1]
                rows = slice(visstart, visstart + nrows)

                cindex.flat[rowgrid[:, a2, a1, :]] = numpy.array(
                    range(visstart, visstart + nrows))

                ca1[rows] = a1
                ca2[rows] = a2

                # Average over time and frequency for case where polarisation isn't an issue
                def average_from_grid(arr):
                    return average_chunks2(
                        arr, allpwtsgrid[:, a2, a1, :],
                        (time_average[a2, a1], frequency_average[a2, a1]))[0]

                ctime[rows] = average_from_grid(time_grid).flatten()
                cfrequency[rows] = average_from_grid(frequency_grid).flatten()

                for axis in range(3):
                    uvwgrid = numpy.outer(uvw[:, a2, a1, axis],
                                          frequency / constants.c.value)
                    cuvw[rows, axis] = average_from_grid(uvwgrid).flatten()

                # For some variables, we need the sum not the average
                def sum_from_grid(arr):
                    result = average_chunks2(
                        arr, allpwtsgrid[:, a2, a1, :],
                        (time_average[a2, a1], frequency_average[a2, a1]))
                    return result[0] * result[0].size

                cintegration_time[rows] = sum_from_grid(
                    integration_time_grid).flatten()
                cchannel_bandwidth[rows] = sum_from_grid(
                    channel_bandwidth_grid).flatten()

                # For the polarisations we have to perform the time-frequency average separately for each polarisation
                for pol in range(npol):
                    result = average_chunks2(
                        vis[:, a2, a1, :, pol], wts[:, a2, a1, :, pol],
                        (time_average[a2, a1], frequency_average[a2, a1]))
                    cvis[rows, pol], cwts[
                        rows, pol] = result[0].flatten(), result[1].flatten()

                visstart += nrows

    assert cnvis == visstart, "Mismatch between number of rows in coalesced visibility and index"

    return cvis, cuvw, cwts, ctime, cfrequency, cchannel_bandwidth, ca1, ca2, cintegration_time, cindex