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