def test_predict_vis(corr_shape, idm, einsum_sig1, einsum_sig2, a1j, blj, a2j,
g1j, bvis, g2j, chunks):
from africanus.rime.predict import predict_vis
s = sum(chunks['source'])
t = sum(chunks['time'])
a = sum(chunks['antenna'])
c = sum(chunks['channels'])
r = sum(chunks['rows'])
a1_jones = rc((s, t, a, c) + corr_shape)
bl_jones = rc((s, r, c) + corr_shape)
a2_jones = rc((s, t, a, c) + corr_shape)
g1_jones = rc((t, a, c) + corr_shape)
base_vis = rc((r, c) + corr_shape)
g2_jones = rc((t, a, c) + corr_shape)
# Row indices into the above time/ant indexed arrays
time_idx = np.asarray([0, 0, 1, 1, 2, 2, 2, 2, 3, 3])
ant1 = np.asarray([0, 0, 0, 0, 1, 1, 1, 2, 2, 3])
ant2 = np.asarray([0, 1, 2, 3, 1, 2, 3, 2, 3, 3])
assert ant1.size == r
model_vis = predict_vis(time_idx, ant1, ant2, a1_jones if a1j else None,
bl_jones if blj else None,
a2_jones if a2j else None,
g1_jones if g1j else None,
base_vis if bvis else None,
g2_jones if g2j else None)
assert model_vis.shape == (r, c) + corr_shape
def _id(array):
return np.broadcast_to(idm, array.shape)
# For einsum, convert (time, ant) dimensions to row
# or ID matrices if the input is present
a1_jones = a1_jones[:, time_idx, ant1] if a1j else _id(bl_jones)
bl_jones = bl_jones if blj else _id(bl_jones)
a2_jones = a2_jones[:, time_idx, ant2].conj() if a2j else _id(bl_jones)
v = np.einsum(einsum_sig1, a1_jones, bl_jones, a2_jones)
if bvis:
v += base_vis
# Convert (time, ant) dimensions to row or
# or ID matrices if input is not present
g1_jones = g1_jones[time_idx, ant1] if g1j else _id(v)
g2_jones = g2_jones[time_idx, ant2].conj() if g2j else _id(v)
v = np.einsum(einsum_sig2, g1_jones, v, g2_jones)
assert_array_almost_equal(v, model_vis)
def test_corrupt_vis(data_factory, corr_shape, jones_shape):
"""
Tests corrupt vis against predict_vis. They should do
the same thing but corrupt_vis adheres to the structure
in the africanus.calibration module.
"""
# simulate noise free data with random DDE's
n_dir = 3
n_time = 32
n_chan = 16
n_ant = 7
sigma_n = 0.0
sigma_f = 0.05
data_dict = data_factory(sigma_n, sigma_f, n_time, n_chan, n_ant, n_dir,
corr_shape, jones_shape)
# make_data uses corrupt_vis to produce the data so we only need to test
# that predict vis gives the same thing on the reshaped arrays
ant1 = data_dict['ANTENNA1']
ant2 = data_dict['ANTENNA2']
vis = data_dict['DATA']
model = data_dict['MODEL_DATA']
jones = data_dict['JONES']
time = data_dict['TIME']
# predict_vis expects (source, time, ant, chan, corr1, corr2) so
# we need to transpose the axes while preserving corr_shape and jones_shape
if jones_shape != corr_shape:
# This only happens in DIAG mode and we need to broadcast jones_shape
# to match corr_shape
tmp = np.zeros((n_time, n_ant, n_chan, n_dir) + corr_shape,
dtype=np.complex128)
tmp[:, :, :, :, 0, 0] = jones[:, :, :, :, 0]
tmp[:, :, :, :, 1, 1] = jones[:, :, :, :, 1]
jones = tmp
if len(corr_shape) == 2:
jones = np.transpose(jones, [3, 0, 1, 2, 4, 5])
model = np.transpose(model, [2, 0, 1, 3, 4])
elif len(corr_shape) == 1:
jones = np.transpose(jones, [3, 0, 1, 2, 4])
model = np.transpose(model, [2, 0, 1, 3])
else:
raise ValueError("Unsupported correlation shapes")
# get vis
time_index = np.unique(time, return_inverse=True)[1]
test_vis = predict_vis(time_index,
ant1,
ant2,
source_coh=model,
dde1_jones=jones,
dde2_jones=jones)
assert_array_almost_equal(test_vis, vis, decimal=10)
def test_residual_vis():
"""
Tests subtraction of model by subtracting all but one
direction from noise free simulated data and comparing
the output to the unsubtracted direction.
"""
np.random.seed(42)
# simulate noise free data with random DDE's
n_dir = 3
sigma_n = 0.0
sigma_f = 0.05
data_dict = make_dual_pol_data(sigma_n, n_dir, sigma_f)
time = data_dict['TIME']
_, time_bin_indices, _, time_bin_counts = unique_time(time)
ant1 = data_dict['ANTENNA1']
ant2 = data_dict['ANTENNA2']
vis = data_dict['DATA']
model = data_dict['MODEL_DATA']
jones = data_dict['JONES']
flag = data_dict['FLAG']
# split the model and jones terms
model_unsubtracted = model[:, :, 0:1]
model_subtract = model[:, :, 1::]
jones_unsubtracted = jones[:, :, :, 0:1]
jones_subtract = jones[:, :, :, 1::]
# subtract all but one direction
residual = residual_vis(time_bin_indices, time_bin_counts, ant1, ant2,
jones_subtract, vis, flag, model_subtract)
# apply gains to the unsubtracted direction
jones_tmp = np.transpose(jones_unsubtracted, [3, 0, 1, 2, 4])
model_tmp = np.transpose(model_unsubtracted, [2, 0, 1, 3])
time_index = np.unique(time, return_inverse=True)[1]
vis_unsubtracted = predict_vis(time_index,
ant1,
ant2,
dde1_jones=jones_tmp,
source_coh=model_tmp,
dde2_jones=jones_tmp)
# residual should now be equal to unsubtracted vis
assert_array_almost_equal(residual, vis_unsubtracted, decimal=5)
def test_dask_predict_vis(corr_shape, idm, einsum_sig1, einsum_sig2, a1j, blj,
a2j, g1j, bvis, g2j, chunks):
da = pytest.importorskip('dask.array')
import numpy as np
import dask
from africanus.rime.predict import predict_vis as np_predict_vis
from africanus.rime.dask import predict_vis
# chunk sizes
sc = chunks['source']
tc = chunks['time']
rrc = chunks['rows']
ac = chunks['antenna']
cc = chunks['channels']
# dimension sizes
s = sum(sc) # sources
t = sum(tc) # times
a = sum(ac) # antennas
c = sum(cc) # channels
r = sum(rrc) # rows
a1_jones = rc((s, t, a, c) + corr_shape)
a2_jones = rc((s, t, a, c) + corr_shape)
bl_jones = rc((s, r, c) + corr_shape)
g1_jones = rc((t, a, c) + corr_shape)
base_vis = rc((r, c) + corr_shape)
g2_jones = rc((t, a, c) + corr_shape)
# Row indices into the above time/ant indexed arrays
time_idx = np.asarray([0, 0, 1, 1, 2, 2, 2, 2, 3, 3])
ant1 = np.asarray([0, 0, 0, 0, 1, 1, 1, 2, 2, 3])
ant2 = np.asarray([0, 1, 2, 3, 1, 2, 3, 2, 3, 3])
assert ant1.size == r
np_model_vis = np_predict_vis(
time_idx, ant1, ant2, a1_jones if a1j else None,
bl_jones if blj else None, a2_jones if a2j else None,
g1_jones if g1j else None, base_vis if bvis else None,
g2_jones if g2j else None)
da_time_idx = da.from_array(time_idx, chunks=rrc)
da_ant1 = da.from_array(ant1, chunks=rrc)
da_ant2 = da.from_array(ant2, chunks=rrc)
da_a1_jones = da.from_array(a1_jones, chunks=(sc, tc, ac, cc) + corr_shape)
da_bl_jones = da.from_array(bl_jones, chunks=(sc, rrc, cc) + corr_shape)
da_a2_jones = da.from_array(a2_jones, chunks=(sc, tc, ac, cc) + corr_shape)
da_g1_jones = da.from_array(g1_jones, chunks=(tc, ac, cc) + corr_shape)
da_base_vis = da.from_array(base_vis, chunks=(rrc, cc) + corr_shape)
da_g2_jones = da.from_array(g2_jones, chunks=(tc, ac, cc) + corr_shape)
args = (da_time_idx, da_ant1, da_ant2, da_a1_jones if a1j else None,
da_bl_jones if blj else None, da_a2_jones if a2j else None,
da_g1_jones if g1j else None, da_base_vis if bvis else None,
da_g2_jones if g2j else None)
stream_model_vis = predict_vis(*args, streams=True)
fan_model_vis = predict_vis(*args, streams=False)
stream_model_vis, fan_model_vis = dask.compute(stream_model_vis,
fan_model_vis)
assert_array_almost_equal(fan_model_vis, np_model_vis)
assert_array_almost_equal(stream_model_vis, fan_model_vis)
def make_dual_pol_data(sigma_n, n_dir, sigma_f):
# make aux data
antenna1 = np.zeros(n_row, dtype=np.int16)
antenna2 = np.zeros(n_row, dtype=np.int16)
time = np.zeros(n_row, dtype=np.float32)
uvw = np.zeros((n_row, 3), dtype=np.float32)
unique_time = np.linspace(0, 1, n_time)
freq = np.linspace(1e9, 2e9, n_chan)
for i in range(n_time):
row = 0
for p in range(n_ant):
for q in range(p):
time[i * n_bl + row] = unique_time[i]
antenna1[i * n_bl + row] = p
antenna2[i * n_bl + row] = q
uvw[i * n_bl + row] = np.random.randn(3)
row += 1
assert time.size == n_row
# simulate visibilities
model_data = np.zeros((n_row, n_chan, n_dir, n_cor), dtype=np.complex64)
# make up some sources
lm = give_lm(n_dir)
flux = give_flux(n_dir)
# simulate model data (pure Stokes I)
for dir in range(n_dir):
this_lm = lm[dir].reshape(1, 2)
this_flux = np.tile(flux[dir], (n_chan, n_cor))[None, :, :]
model_tmp = im_to_vis(this_flux, uvw, this_lm, freq)
model_data[:, :, dir, :] = model_tmp
assert not np.isnan(model_data).any()
# simulate gains (just radnomly scattered around 1 for now)
jones = np.ones((n_time, n_ant, n_chan, n_dir, n_cor), dtype=np.complex64)
if sigma_f:
jones += sigma_f * (
np.random.randn(n_time, n_ant, n_chan, n_dir, n_cor) +
1.0j * np.random.randn(n_time, n_ant, n_chan, n_dir, n_cor))
assert (np.abs(jones) > 1e-5).all()
assert not np.isnan(jones).any()
# get vis
time_index = np.unique(time, return_inverse=True)[1]
jones_tmp = np.transpose(jones, [3, 0, 1, 2, 4])
model_tmp = np.transpose(model_data, [2, 0, 1, 3])
vis = predict_vis(time_index,
antenna1,
antenna2,
source_coh=model_tmp,
dde1_jones=jones_tmp,
dde2_jones=jones_tmp)
assert not np.isnan(vis).any()
# add noise
if sigma_n:
vis += sigma_n * (np.random.randn(n_row, n_chan, n_cor) +
1.0j * np.random.randn(n_row, n_chan, n_cor))
weights = np.ones((n_row, n_chan, n_cor), dtype=np.float32)
if sigma_n:
weights /= sigma_n**2
flag = np.zeros((n_row, n_chan, n_cor), dtype=np.bool)
data_dict = {}
data_dict["DATA"] = vis
data_dict["MODEL_DATA"] = model_data
data_dict["WEIGHT_SPECTRUM"] = weights
data_dict["TIME"] = time
data_dict["ANTENNA1"] = antenna1
data_dict["ANTENNA2"] = antenna2
data_dict["FLAG"] = flag
data_dict['JONES'] = jones
return data_dict
def test_dask_predict_vis(corr_shape, idm, einsum_sig1, einsum_sig2, a1j, blj,
a2j, g1j, bvis, g2j):
da = pytest.importorskip('dask.array')
import numpy as np
from africanus.rime.predict import predict_vis as np_predict_vis
from africanus.rime.dask import predict_vis
# chunk sizes
sc = (2, 3, 4) # sources
tc = (2, 1, 1) # times
rrc = (4, 4, 2) # rows
ac = (4, ) # antennas
cc = (3, 2) # channels
# dimension sizes
s = sum(sc) # sources
t = sum(tc) # times
a = sum(ac) # antennas
c = sum(cc) # channels
r = sum(rrc) # rows
a1_jones = rc((s, t, a, c) + corr_shape)
a2_jones = rc((s, t, a, c) + corr_shape)
bl_jones = rc((s, r, c) + corr_shape)
g1_jones = rc((t, a, c) + corr_shape)
base_vis = rc((r, c) + corr_shape)
g2_jones = rc((t, a, c) + corr_shape)
# Row indices into the above time/ant indexed arrays
time_idx = np.asarray([0, 0, 1, 1, 2, 2, 2, 2, 3, 3])
ant1 = np.asarray([0, 0, 0, 0, 1, 1, 1, 2, 2, 3])
ant2 = np.asarray([0, 1, 2, 3, 1, 2, 3, 2, 3, 3])
assert ant1.size == r
np_model_vis = np_predict_vis(
time_idx, ant1, ant2, a1_jones if a1j else None,
bl_jones if blj else None, a2_jones if a2j else None,
g1_jones if g1j else None, base_vis if bvis else None,
g2_jones if g2j else None)
da_time_idx = da.from_array(time_idx, chunks=rrc)
da_ant1 = da.from_array(ant1, chunks=rrc)
da_ant2 = da.from_array(ant2, chunks=rrc)
da_a1_jones = da.from_array(a1_jones, chunks=(sc, tc, ac, cc) + corr_shape)
da_bl_jones = da.from_array(bl_jones, chunks=(sc, rrc, cc) + corr_shape)
da_a2_jones = da.from_array(a2_jones, chunks=(sc, tc, ac, cc) + corr_shape)
da_g1_jones = da.from_array(g1_jones, chunks=(tc, ac, cc) + corr_shape)
da_base_vis = da.from_array(base_vis, chunks=(rrc, cc) + corr_shape)
da_g2_jones = da.from_array(g2_jones, chunks=(tc, ac, cc) + corr_shape)
model_vis = predict_vis(
da_time_idx, da_ant1, da_ant2, da_a1_jones if a1j else None,
da_bl_jones if blj else None, da_a2_jones if a2j else None,
da_g1_jones if g1j else None, da_base_vis if bvis else None,
da_g2_jones if g2j else None)
model_vis = model_vis.compute()
if not np.allclose(model_vis, np_model_vis):
diff = model_vis - np_model_vis
diff[np.abs(diff) < 1e-10] = 0.0
problems = np.array(np.nonzero(diff)).T
for p in (tuple(p.tolist()) for p in problems):
print(p, model_vis[p], np_model_vis[p])
assert np.allclose(model_vis, np_model_vis)