def __init__(self, label, data_arr, ndir, nmod, chunk_ts, chunk_fs,
chunk_label, options):
"""
Initialises a diagonal phase-slope gain machine.
Args:
label (str):
Label identifying the Jones term.
data_arr (np.ndarray):
Shape (n_mod, n_tim, n_fre, n_ant, n_ant, n_cor, n_cor) array containing observed
visibilities.
ndir (int):
Number of directions.
nmod (nmod):
Number of models.
chunk_ts (np.ndarray):
Times for the data being processed.
chunk_fs (np.ndarray):
Frequencies for the data being processsed.
options (dict):
Dictionary of options.
"""
ParameterisedGains.__init__(self, label, data_arr, ndir, nmod,
chunk_ts, chunk_fs, chunk_label, options)
self.slope_type = options["type"]
self.n_param = 3 if self.slope_type == "tf-plane" else 2
self.param_shape = [
self.n_dir, self.n_timint, self.n_freint, self.n_ant, self.n_param,
self.n_cor, self.n_cor
]
self.slope_params = np.zeros(self.param_shape, dtype=self.ftype)
self.posterior_slope_error = None
self.chunk_ts = _normalize(chunk_ts, self.ftype)
self.chunk_fs = _normalize(chunk_fs, self.ftype)
if self.slope_type == "tf-plane":
self.slope = cubical.kernels.import_kernel("tf_plane")
self._labels = dict(phase=0, delay=1, rate=2)
elif self.slope_type == "f-slope":
self.slope = cubical.kernels.import_kernel("f_slope")
self._labels = dict(phase=0, delay=1)
elif self.slope_type == "t-slope":
self.slope = cubical.kernels.import_kernel("t_slope")
self._labels = dict(phase=0, rate=1)
else:
raise RuntimeError("unknown machine type '{}'".format(
self.slope_type))
# kernel used in solver is diag-diag in diag mode, else uses full kernel version
if options.get('diag-data') or options.get('diag-only'):
self.kernel_solve = cubical.kernels.import_kernel(
'diag_phase_only')
else:
self.kernel_solve = cubical.kernels.import_kernel('phase_only')
self._jhr0 = self._gerr = None
def restrict_solution(self):
"""
Restricts the solution by invoking the inherited restrict_soultion method and applying
any machine specific restrictions.
"""
ParameterisedGains.restrict_solution(self)
if self.ref_ant is not None:
self.slope_params -= self.slope_params[:, :, :, self.
ref_ant, :, :, :][:, :, :,
np.
newaxis, :, :, :]
for idir in self.fix_directions:
self.slope_params[idir, ...] = 0
def precompute_attributes(self, data_arr, model_arr, flags_arr,
inv_var_chan):
"""
Precompute (J\ :sup:`H`\J)\ :sup:`-1`, which does not vary with iteration.
Args:
model_arr (np.ndarray):
Shape (n_dir, n_mod, n_tim, n_fre, n_ant, n_ant, n_cor, n_cor) array containing
model visibilities.
"""
ParameterisedGains.precompute_attributes(self, data_arr, model_arr,
flags_arr, inv_var_chan)
jhj1_shape = [self.n_dir, self.n_tim, self.n_fre, self.n_ant, 2, 2]
jhj1 = self.slope.allocate_gain_array(jhj1_shape,
dtype=self.dtype,
zeros=True)
# use appropriate phase-only kernel (with 1,1 intervals) to compute inner JHJ
self.kernel_solve.compute_jhj(model_arr, jhj1, 1, 1)
blocks_per_inverse = 6 if self.slope_type == "tf-plane" else 3
jhj_shape = [
self.n_dir, self.n_timint, self.n_freint, self.n_ant,
blocks_per_inverse, 2, 2
]
jhj = self.slope.allocate_param_array(jhj_shape,
dtype=self.ftype,
zeros=True)
self.slope.compute_jhj(jhj1.real, jhj, self.chunk_ts, self.chunk_fs,
self.t_int, self.f_int)
self.jhjinv = np.zeros_like(jhj)
self.slope.compute_jhjinv(jhj, self.jhjinv, self.eps)
def exportable_solutions():
""" Returns a dictionary of exportable solutions for this machine type. """
exportables = ParameterisedGains.exportable_solutions()
exportables.update({
"phase": (0., ("dir", "time", "freq", "ant", "corr")),
"delay": (0., ("dir", "time", "freq", "ant", "corr")),
"rate": (0., ("dir", "time", "freq", "ant", "corr")),
"phase.err": (0., ("dir", "time", "freq", "ant", "corr")),
"delay.err": (0., ("dir", "time", "freq", "ant", "corr")),
"rate.err": (0., ("dir", "time", "freq", "ant", "corr")),
})
return exportables
def export_solutions(self):
""" Saves the solutions to a dict of {label: solutions,grids} items. """
solutions = ParameterisedGains.export_solutions(self)
for label, num in self._labels.iteritems():
solutions[label] = masked_array(
self.slope_params[..., num, (0, 1),
(0, 1)]), self.interval_grid
if self.posterior_slope_error is not None:
solutions[label + ".err"] = masked_array(
self.posterior_slope_error[...,
num, :]), self.interval_grid
return solutions
def export_solutions(self):
""" Saves the solutions to a dict of {label: solutions,grids} items. """
solutions = ParameterisedGains.export_solutions(self)
for label, num in self._labels.items():
solutions[label] = masked_array(
self.slope_params[..., num, (0, 1),
(0, 1)]), self.interval_grid
if self.posterior_slope_error is not None:
solutions[label + ".err"] = masked_array(
self.posterior_slope_error[...,
num, :]), self.interval_grid
with np.printoptions(precision=3, linewidth=100000, threshold=10000):
log(1).print("{}: slope solutions are {}, {}".format(
self.chunk_label, self.slope_params[..., :, :, 1, 0, 0],
self.slope_params[..., :, :, 1, 1, 1]))
return solutions
def __init__(self, label, data_arr, ndir, nmod, chunk_ts, chunk_fs,
chunk_label, options):
"""
Initialises a diagonal phase-slope gain machine.
Args:
label (str):
Label identifying the Jones term.
data_arr (np.ndarray):
Shape (n_mod, n_tim, n_fre, n_ant, n_ant, n_cor, n_cor) array containing observed
visibilities.
ndir (int):
Number of directions.
nmod (nmod):
Number of models.
chunk_ts (np.ndarray):
Times for the data being processed.
chunk_fs (np.ndarray):
Frequencies for the data being processsed.
options (dict):
Dictionary of options.
"""
ParameterisedGains.__init__(self, label, data_arr, ndir, nmod,
chunk_ts, chunk_fs, chunk_label, options)
self.slope_type = options["type"]
self.n_param = 3 if self.slope_type == "tf-plane" else 2
self._estimate_delays = options["estimate-delays"]
self._pin_slope_iters = options["pin-slope-iters"]
if self._pin_slope_iters > self.maxiter:
raise ValueError(
"Slope pinning iterations is greater than the maximum "
"number of iterations. Please check term-iters, max-iters "
"and pin-slope-iters.")
if (self.slope_type != "f-slope") and (self._pin_slope_iters != 0):
raise ValueError(
"Slope pinning is not supported for modes other than "
"f-slope. Please ensure that pin-slope-iters is zero.")
self.param_shape = [
self.n_dir, self.n_timint, self.n_freint, self.n_ant, self.n_param,
self.n_cor, self.n_cor
]
# This needs to change - we want to initialise these slope params
# from the fourier transform along the relevant axis.
self.slope_params = np.zeros(self.param_shape, dtype=self.ftype)
self.posterior_slope_error = None
self.chunk_ts = _normalize(chunk_ts, self.ftype)
self.chunk_fs = _normalize(chunk_fs, self.ftype)
if self.slope_type == "tf-plane":
self.slope = cubical.kernels.import_kernel("tf_plane")
self._labels = dict(phase=0, delay=1, rate=2)
elif self.slope_type == "f-slope":
self.slope = cubical.kernels.import_kernel("f_slope")
self._labels = dict(phase=0, delay=1)
if self._estimate_delays:
# Select all baselines containing the reference antenna.
ref_ant_data = data_arr[:, :, :, self.ref_ant]
# Average over time solution intervals. This should improve SNR,
# but is not always necesssary.
interval_data = np.add.reduceat(ref_ant_data, self.t_bins, 1)
interval_smpl = np.add.reduceat(ref_ant_data != 0, self.t_bins,
1)
selection = np.where(interval_smpl)
interval_data[selection] /= interval_smpl[selection]
bad_bins = np.add.reduceat(interval_data, self.f_bins, 2)
# FFT the data along the frequency axis. TODO: Need to consider
# what happens if we solve for few delays across the band. As there
# is no guarantee that the band will be perfectly split, this may
# need to be a loop over frequency solution intervals.
pad_factor = options["delay-estimate-pad-factor"]
for i in range(self.n_freint):
edges = self.f_bins + [None]
slice_fs = self.chunk_fs[edges[i]:edges[i + 1]]
slice_data = interval_data[:, :, edges[i]:edges[i + 1]]
slice_nchan = slice_data.shape[2]
fft_data = np.abs(
np.fft.fft(slice_data,
n=slice_nchan * pad_factor,
axis=2))
fft_data = np.fft.fftshift(fft_data, axes=2)
# Convert the normalised frequency values into delay values.
delta_freq = slice_fs[1] - slice_fs[0]
fft_freq = np.fft.fftfreq(slice_nchan * pad_factor,
delta_freq)
fft_freq = np.fft.fftshift(fft_freq)
# Find the delay value at which the FFT of the data is
# maximised. As we do not pad the values, this only a rough
# approximation of the delay. We also reintroduce the
# frequency axis for consistency.
delay_est_ind = np.argmax(fft_data, axis=2)
delay_est_ind = np.expand_dims(delay_est_ind, axis=2)
delay_est = fft_freq[delay_est_ind]
delay_est[..., (0, 1), (1, 0)] = 0
# Check for bad data points (bls missing across all channels)
# and set their estimates to zero.
selection = np.where(bad_bins[:, :, i:i + 1] == 0)
delay_est[selection] = 0
# Zero the off diagonals and take the negative delay values -
# this is necessary as we technically get the delay
# corresponding to the conjugate term.
self.slope_params[..., i:i + 1, :,
1, :, :] = -1 * delay_est
self.slope_params[..., (1, 0), (0, 1)] = 0
log(1).print("{}: slope estimates are {}, {}".format(
chunk_label, " ".join(
map(str, self.slope_params[..., i:i + 1, :, 1, 0,
0])),
" ".join(
map(str, self.slope_params[..., i:i + 1, :, 1, 1,
1]))))
elif self.slope_type == "t-slope":
self.slope = cubical.kernels.import_kernel("t_slope")
self._labels = dict(phase=0, rate=1)
else:
raise RuntimeError("unknown machine type '{}'".format(
self.slope_type))
self.slope.construct_gains(self.slope_params, self.gains,
self.chunk_ts, self.chunk_fs, self.t_int,
self.f_int)
# kernel used in solver is diag-diag in diag mode, else uses full kernel version
if options.get('diag-data') or options.get('diag-only'):
self.kernel_solve = cubical.kernels.import_kernel(
'diag_phase_only')
else:
self.kernel_solve = cubical.kernels.import_kernel('phase_only')
self._jhr0 = self._gerr = None