def _predict_coh_wrapper(time_index, antenna1, antenna2,
dde1_jones, source_coh, dde2_jones,
base_vis,
reduce_single_source=False):
if reduce_single_source:
# All these arrays contract over a single 'source' chunk
dde1_jones = dde1_jones[0] if dde1_jones else None
source_coh = source_coh[0] if source_coh else None
dde2_jones = dde2_jones[0] if dde2_jones else None
vis = np_predict_vis(time_index, antenna1, antenna2,
# dde1_jones contracts over a single 'ant' chunk
dde1_jones[0] if dde1_jones else None,
source_coh,
# dde2_jones contracts over a single 'ant' chunk
dde2_jones[0] if dde2_jones else None,
None,
base_vis,
None)
if reduce_single_source:
return vis
return vis[None, ...]
def test_cuda_predict_vis(corr_shape, idm, einsum_sig1, einsum_sig2, a1j, blj,
a2j, g1j, bvis, g2j, chunks):
np.random.seed(40)
cp = pytest.importorskip('cupy')
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)
# Add 10 to the index to test time index normalisation
time_idx = np.concatenate([
np.full(rows, i + 10, dtype=np.int32)
for i, rows in enumerate(chunks['rows'])
])
ant1 = np.concatenate([
np.random.randint(0, a, rows, dtype=np.int32)
for rows in chunks['rows']
])
ant2 = np.concatenate([
np.random.randint(0, a, rows, dtype=np.int32)
for rows in chunks['rows']
])
assert ant1.size == r
model_vis = predict_vis(cp.asarray(time_idx), cp.asarray(ant1),
cp.asarray(ant2),
cp.asarray(a1_jones) if a1j else None,
cp.asarray(bl_jones) if blj else None,
cp.asarray(a2_jones) if a2j else None,
cp.asarray(g1_jones) if g1j else None,
cp.asarray(base_vis) if bvis else None,
cp.asarray(g2_jones) if g2j else None)
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)
np.testing.assert_array_almost_equal(cp.asnumpy(model_vis), np_model_vis)
def _predict_dies_wrapper(time_index, antenna1, antenna2,
die1_jones, base_vis, die2_jones):
return np_predict_vis(time_index, antenna1, antenna2,
None,
None,
None,
# die1_jones loses the 'ant' dim
die1_jones[0] if die1_jones else None,
base_vis,
# die2_jones loses the 'ant' dim
die2_jones[0] if die2_jones else None)
def _predict_dies_wrapper(time_index, antenna1, antenna2,
dde1_jones, source_coh, dde2_jones,
die1_jones, base_vis, die2_jones):
return np_predict_vis(time_index, antenna1, antenna2,
# dde1_jones loses the 'source' and 'ant' dims
dde1_jones[0][0] if dde1_jones else None,
# source_coh loses the 'source' dim
source_coh[0] if source_coh else None,
# dde2_jones loses the 'source' and 'ant' dims
dde2_jones[0][0] if dde2_jones else None,
# die1_jones loses the 'ant' dim
die1_jones[0] if die1_jones else None,
base_vis,
# die2_jones loses the 'ant' dim
die2_jones[0] if die2_jones else None)
def _predict_coh_wrapper(time_index, antenna1, antenna2,
dde1_jones, source_coh, dde2_jones,
die1_jones, base_vis, die2_jones):
return (np_predict_vis(time_index, antenna1, antenna2,
# dde1_jones loses the 'ant' dim
dde1_jones[0] if dde1_jones else None,
# source_coh loses the 'source' dim
source_coh,
# dde2_jones loses the 'source' and 'ant' dims
dde2_jones[0] if dde2_jones else None,
# die1_jones loses the 'ant' dim
die1_jones[0] if die1_jones else None,
base_vis,
# die2_jones loses the 'ant' dim
die2_jones[0] if die2_jones else None)
# Introduce an extra dimension (source dim reduced to 1)
[None, ...])
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 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)