def test_average_chunks_jit(self):
arr = numpy.linspace(0.0, 100.0, 11)
wts = numpy.ones_like(arr)
carr, cwts = average_chunks(arr, wts, 2)
carr_jit, cwts_jit = average_chunks_jit(arr, wts, 2)
numpy.testing.assert_array_equal(carr, carr_jit)
numpy.testing.assert_array_equal(cwts, cwts_jit)
def test_average_chunks_zero(self):
arr = numpy.linspace(0.0, 90.0, 10)
wts = numpy.ones_like(arr)
carr, cwts = average_chunks(arr, wts, 0)
assert len(carr) == len(cwts)
numpy.testing.assert_array_equal(carr, arr)
numpy.testing.assert_array_equal(cwts, wts)
def test_average_chunks_single(self):
arr = numpy.linspace(0.0, 100.0, 11)
wts = numpy.ones_like(arr)
carr, cwts = average_chunks(arr, wts, 12)
assert len(carr) == len(cwts)
answerarr = numpy.array([50.0])
answerwts = numpy.array([11.0])
numpy.testing.assert_array_equal(carr, answerarr)
numpy.testing.assert_array_equal(cwts, answerwts)
def test_average_chunks_exact(self):
arr = numpy.linspace(0.0, 90.0, 10)
wts = numpy.ones_like(arr)
carr, cwts = average_chunks(arr, wts, 2)
assert len(carr) == len(cwts)
answerarr = numpy.array([5., 25., 45., 65.0, 85.0])
answerwts = numpy.array([2.0, 2.0, 2.0, 2.0, 2.0])
numpy.testing.assert_array_equal(carr, answerarr)
numpy.testing.assert_array_equal(cwts, answerwts)
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_models. 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
def simulate_rfi_image(config, times, frequency, channel_bandwidth,
phasecentre, polarisation_frame, time_average,
channel_average, attenuation, noise, emitter_location,
emitter_power, use_pole, waterfall, write_ms):
averaged_frequency = numpy.array(
average_chunks(frequency, numpy.ones_like(frequency),
channel_average))[0]
averaged_channel_bandwidth, wts = numpy.array(
average_chunks(channel_bandwidth, numpy.ones_like(frequency),
channel_average))
averaged_channel_bandwidth *= wts
averaged_times = numpy.array(
average_chunks(times, numpy.ones_like(times), time_average))[0]
s2r = numpy.pi / 43200.0
bvis = create_blockvisibility(config,
s2r * times,
frequency,
channel_bandwidth=channel_bandwidth,
phasecentre=phasecentre,
polarisation_frame=polarisation_frame,
zerow=False)
bvis = simulate_rfi_block(bvis,
emitter_location=emitter_location,
emitter_power=emitter_power,
attenuation=attenuation,
use_pole=use_pole)
if noise:
bvis = add_noise(bvis)
if waterfall:
plot_waterfall(bvis)
if write_ms:
msname = "simulate_rfi_%.1f.ms" % (times[0])
export_blockvisibility_to_ms(msname, [bvis], "RFI")
averaged_bvis = create_blockvisibility(
config,
s2r * averaged_times,
averaged_frequency,
channel_bandwidth=averaged_channel_bandwidth,
phasecentre=phasecentre,
polarisation_frame=polarisation_frame,
zerow=False)
npol = 1
for itime, _ in enumerate(averaged_times):
atime = itime * time_average
for ant2 in range(nants):
for ant1 in range(ant2, nants):
for ichan, _ in enumerate(averaged_frequency):
achan = ichan * channel_average
for pol in range(npol):
averaged_bvis.data['vis'][itime, ant2, ant1, ichan, pol] = \
calculate_averaged_correlation(
bvis.data['vis'][atime:(atime+time_average), ant2, ant1, achan:(achan+channel_average), pol],
time_average, channel_average)[0,0]
averaged_bvis.data['vis'][itime, ant1, ant2, ichan, pol] = \
numpy.conjugate(averaged_bvis.data['vis'][itime, ant2, ant1, ichan, pol])
achan += 1
atime += 1
del bvis
if noise:
averaged_bvis = add_noise(averaged_bvis)
return averaged_bvis
nintegrations_per_chunk * integration_time)
print("Start times", start_times)
results = list()
pole_results = list()
chunk_start_times = [
start_times[i:i + args.ngroup_visibility]
for i in range(0, len(start_times), args.ngroup_visibility)
]
print("Chunk start times", [c[0] for c in chunk_start_times])
dopsf = args.do_psf == "True"
# Find the average frequencies
averaged_frequency = numpy.array(
average_chunks(frequency, numpy.ones_like(frequency),
channel_average))[0]
if len(averaged_frequency) > 1:
step = abs(averaged_frequency[-1] -
averaged_frequency[0]) / (len(averaged_frequency) - 1)
else:
step = channel_average * channel_bandwidth
averaged_channel_bandwidth = step * numpy.ones_like(averaged_frequency)
print("Each averaged chunk has %d integrations of duration %.2f (s)" %
(nintegrations_per_chunk // time_average,
time_average * integration_time))
print("Each averaged chunk has %d channels of width %.3f (MHz)" %
(len(averaged_frequency), 1e-6 * averaged_channel_bandwidth[0]))
print("Processing %d time chunks in groups of %d" %
(len(start_times), args.ngroup_visibility))