def impl(sigma_n,
sigma_f,
n_time,
n_chan,
n_ant,
n_dir,
corr_shape,
jones_shape,
phase_only_gains=False):
rs = np.random.RandomState(42)
n_bl = n_ant * (n_ant - 1) // 2
n_row = n_bl * n_time
# 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.float64)
uvw = np.zeros((n_row, 3), dtype=np.float64)
time_values = 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] = time_values[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) + corr_shape,
dtype=np.complex128)
# make up some sources
lm = lm_factory(n_dir, rs)
alpha = -0.7
freq0 = freq[n_chan // 2]
flux = flux_factory(n_dir, n_chan, corr_shape, alpha, freq, freq0, rs)
# simulate model data
for dir in range(n_dir):
dir_lm = lm[dir].reshape(1, 2)
# Get flux for source (keep source axis, flatten cor axis)
dir_flux = flux[dir].reshape(1, n_chan, np.prod(corr_shape))
tmp = im_to_vis(dir_flux, uvw, dir_lm, freq)
model_data[:, :, dir] = tmp.reshape((n_row, n_chan) + corr_shape)
assert not np.isnan(model_data).any()
# simulate gains (just randomly scattered around 1 for now)
jones = np.ones((n_time, n_ant, n_chan, n_dir) + jones_shape,
dtype=np.complex128)
if sigma_f:
if phase_only_gains:
jones = np.exp(
1.0j * rs.normal(loc=0.0, scale=sigma_f, size=jones.shape))
else:
jones += (
rs.normal(loc=0.0, scale=sigma_f, size=jones.shape) +
1.0j * rs.normal(loc=0.0, scale=sigma_f, size=jones.shape))
assert (np.abs(jones) > 1e-5).all()
assert not np.isnan(jones).any()
# get vis
_, time_bin_indices, time_bin_counts = chunkify_rows(time, n_time)
vis = corrupt_vis(time_bin_indices, time_bin_counts, antenna1,
antenna2, jones, model_data)
assert not np.isnan(vis).any()
# add noise
if sigma_n:
vis += (rs.normal(loc=0.0, scale=sigma_n, size=vis.shape) +
1.0j * rs.normal(loc=0.0, scale=sigma_n, size=vis.shape))
weights = np.ones(vis.shape, dtype=np.float64)
if sigma_n:
weights /= sigma_n**2
flag = np.zeros(vis.shape, 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 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 calibrate(args, jones, alphas):
# simple calibration to test if simulation went as expected.
# Note do not run on large data set
# load data
ms = table(args.ms)
time = ms.getcol('TIME')
_, tbin_idx, tbin_counts = chunkify_rows(time, args.utimes_per_chunk)
n_time = tbin_idx.size
ant1 = ms.getcol('ANTENNA1')
ant2 = ms.getcol('ANTENNA2')
n_ant = np.maximum(ant1.max(), ant2.max()) + 1
uvw = ms.getcol('UVW').astype(np.float64)
data = ms.getcol(args.out_col) # this is where we put the data
# we know it is pure Stokes I so we can solve using diagonals only
data = data[:, :, (0, 3)].astype(np.complex128)
n_row, n_freq, n_corr = data.shape
flag = ms.getcol('FLAG')
flag = flag[:, :, (0, 3)]
# get phase dir
radec0 = table(args.ms + '::FIELD').getcol('PHASE_DIR').squeeze().astype(
np.float64)
# get freqs
freq = table(args.ms + '::SPECTRAL_WINDOW').getcol('CHAN_FREQ')[0].astype(
np.float64)
assert freq.size == n_freq
# now get the model
# get source coordinates from lsm
lsm = Tigger.load(args.sky_model)
radec = []
stokes = []
spi = []
ref_freqs = []
for source in lsm.sources:
radec.append([source.pos.ra, source.pos.dec])
stokes.append([source.flux.I])
tmp_spec = source.spectrum
spi.append([tmp_spec.spi if tmp_spec is not None else 0.0])
ref_freqs.append([tmp_spec.freq0 if tmp_spec is not None else 1.0])
n_dir = len(stokes)
radec = np.asarray(radec)
lm = radec_to_lm(radec, radec0)
# get model visibilities
model = np.zeros((n_row, n_freq, n_dir, 2), dtype=np.complex)
stokes = np.asarray(stokes)
ref_freqs = np.asarray(ref_freqs)
spi = np.asarray(spi)
for d in range(n_dir):
Stokes_I = stokes[d] * (freq / ref_freqs[d])**spi[d]
model[:, :, d, 0:1] = im_to_vis(Stokes_I[None, :, None], uvw,
lm[d:d + 1], freq)
model[:, :, d, 1] = model[:, :, d, 0]
# set weights to unity
weight = np.ones_like(data, dtype=np.float64)
# initialise gains
jones0 = np.ones((n_time, n_ant, n_freq, n_dir, n_corr),
dtype=np.complex128)
# calibrate
ti = timeit()
jones_hat, jhj, jhr, k = gauss_newton(tbin_idx,
tbin_counts,
ant1,
ant2,
jones0,
data,
flag,
model,
weight,
tol=1e-5,
maxiter=100)
print("%i iterations took %fs" % (k, timeit() - ti))
# verify result
for p in range(2):
for q in range(p):
diff_true = np.angle(jones[:, p] * jones[:, q].conj())
diff_hat = np.angle(jones_hat[:, p] * jones_hat[:, q].conj())
try:
assert_array_almost_equal(diff_true, diff_hat, decimal=2)
except Exception as e:
print(e)