def phase_only_gauss_newton(time_bin_indices, time_bin_counts, antenna1,
antenna2, jones, vis, flag, model,
weight, tol=1e-4, maxiter=100):
# whiten data
sqrtweights = np.sqrt(weight)
vis *= sqrtweights
model *= sqrtweights[:, :, None]
eps = 1.0
k = 0
while eps > tol and k < maxiter:
# keep track of old phases
phases = np.angle(jones)
# get residual
residual = residual_vis(time_bin_indices, time_bin_counts, antenna1,
antenna2, jones, vis, flag, model)
# get diag(jhj) and jhr
jhj, jhr = jhj_and_jhr(time_bin_indices, time_bin_counts, antenna1,
antenna2, jones, residual, model, flag)
# implement update
phases_new = phases + (jhr/jhj).real
jones = np.exp(1.0j * phases_new)
# check convergence/iteration control
eps = np.abs(phases_new - phases).max()
k += 1
return jones, jhj, jhr, k
def test_residual_vis_dask(data_factory, corr_shape, jones_shape):
da = pytest.importorskip("dask.array")
# simulate noise free data with only DIE's
n_dir = 3
n_time = 32
n_chan = 16
n_ant = 4
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)
vis = data_dict['DATA']
ant1 = data_dict['ANTENNA1']
ant2 = data_dict['ANTENNA2']
model = data_dict['MODEL_DATA'] # what we need to compare to
jones = data_dict['JONES']
time = data_dict['TIME']
flag = data_dict['FLAG']
# get chunking scheme
ncpu = 8
utimes_per_chunk = n_time // ncpu
row_chunks, time_bin_idx, time_bin_counts = chunkify_rows(
time, utimes_per_chunk)
# set up dask arrays
da_time_bin_idx = da.from_array(time_bin_idx, chunks=(utimes_per_chunk))
da_time_bin_counts = da.from_array(time_bin_counts,
chunks=(utimes_per_chunk))
da_ant1 = da.from_array(ant1, chunks=row_chunks)
da_ant2 = da.from_array(ant2, chunks=row_chunks)
da_vis = da.from_array(vis, chunks=(row_chunks, (n_chan, )) + (corr_shape))
da_model = da.from_array(model,
chunks=(row_chunks, (n_chan, ),
(n_dir, )) + (corr_shape))
da_jones = da.from_array(jones,
chunks=(utimes_per_chunk, n_ant, n_chan, n_dir) +
jones_shape)
da_flag = da.from_array(flag,
chunks=(row_chunks, (n_chan, )) + (corr_shape))
from africanus.calibration.utils import residual_vis as residual_vis_np
residual = residual_vis_np(time_bin_idx, time_bin_counts, ant1, ant2,
jones, vis, flag, model)
from africanus.calibration.utils.dask import residual_vis
da_residual = residual_vis(da_time_bin_idx, da_time_bin_counts, da_ant1,
da_ant2, da_jones, da_vis, da_flag, da_model)
residual2 = da_residual.compute()
assert_array_almost_equal(residual, residual2, decimal=10)
def gauss_newton(time_bin_indices, time_bin_counts, antenna1,
antenna2, jones, vis, flag, model,
weight, tol=1e-4, maxiter=100):
# whiten data
sqrtweights = np.sqrt(weight)
vis *= sqrtweights
model *= sqrtweights[:, :, None]
mode = check_type(jones, vis)
# can avoid recomputing JHJ in DIAG_DIAG mode
if mode == DIAG_DIAG:
jhj = compute_jhj(time_bin_indices, time_bin_counts,
antenna1, antenna2, jones, model, flag)
else:
raise NotImplementedError("Only DIAG_DIAG mode implemented")
eps = 1.0
k = 0
while eps > tol and k < maxiter:
# keep track of old phases
phases = np.angle(jones)
# get residual TODO - we can avoid this in DIE case
residual = residual_vis(time_bin_indices, time_bin_counts, antenna1,
antenna2, jones, vis, flag, model)
jhr = compute_jhr(time_bin_indices, time_bin_counts,
antenna1, antenna2,
jones, residual, model, flag)
# implement update
phases_new = phases + 0.5 * (jhr/jhj).real
jones = np.exp(1.0j * phases_new)
# check convergence/iteration control
eps = np.abs(phases_new - phases).max()
k += 1
return jones, jhj, jhr, k
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_residual_vis(data_factory, corr_shape, jones_shape):
"""
Tests subtraction of model by subtracting all but one
direction from noise free simulated data and comparing
the output to the unsubtracted direction.
"""
from africanus.calibration.utils import residual_vis, corrupt_vis
# 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)
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
vis_unsubtracted = corrupt_vis(time_bin_indices, time_bin_counts, ant1,
ant2, jones_unsubtracted,
model_unsubtracted)
# residual should now be equal to unsubtracted vis
assert_array_almost_equal(residual, vis_unsubtracted, decimal=10)