def restrict_solution(self):
PerIntervalGains.restrict_solution(self)
if self.ref_ant is not None:
phase = np.angle(self.gains[...,self.ref_ant,0,0])
self.gains[:,:,:,:,(0,1),(0,1)] *= np.exp(-1j*phase)[:,:,:,np.newaxis,np.newaxis]
def __init__(self, label, data_arr, ndir, nmod, chunk_ts, chunk_fs,
chunk_label, options):
"""
Initialises a 2x2 complex 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.
"""
# note that this sets up self.kernel
PerIntervalGains.__init__(self, label, data_arr, ndir, nmod, chunk_ts,
chunk_fs, chunk_label, options)
# try guesstimating the PZD
self._estimate_pzd = options["estimate-pzd"]
self._offdiag_only = options["offdiag-only"]
def __init__(self, label, data_arr, ndir, nmod, chunk_ts, chunk_fs,
chunk_label, options):
"""
Initialises a 2x2 complex 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.
"""
PerIntervalGains.__init__(self, label, data_arr, ndir, nmod, chunk_ts,
chunk_fs, chunk_label, options,
self.get_kernel(options))
def precompute_attributes(self, data_arr, model_arr, flags_arr, noise):
"""
"""
PerIntervalGains.precompute_attributes(self, data_arr, model_arr,
flags_arr, noise)
if self._estimate_pzd:
marr = model_arr[..., (0, 1), (1, 0)][:, 0].sum(0)
darr = data_arr[..., (0, 1), (1, 0)][0]
mask = (flags_arr[..., np.newaxis] != 0) | (marr == 0)
dm = darr * (np.conj(marr) / abs(marr))
dabs = np.abs(darr)
dm[mask] = 0
dabs[mask] = 0
# collapse time/freq axis into intervals and sum antenna axes
dm_sum = self.interval_sum(dm).sum(axis=(2, 3))
dabs_sum = self.interval_sum(dabs).sum(axis=(2, 3))
# sum off-diagonal terms
dm_sum = dm_sum[..., 0] + np.conj(dm_sum[..., 1])
dabs_sum = dabs_sum[..., 0] + np.conj(dabs_sum[..., 1])
pzd = np.angle(dm_sum / dabs_sum)
pzd[dabs_sum == 0] = 0
print("{}: PZD estimate {}".format(self.chunk_label, pzd),
file=log(2))
self.gains[:, :, :, :, 1, 1] = np.exp(-1j * pzd)[np.newaxis, :, :,
np.newaxis]
def __init__(self, label, data_arr, ndir, nmod, chunk_ts, chunk_fs, chunk_label, options):
"""
Initialises a diagonal phase-only 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.
"""
PerIntervalGains.__init__(self, label, data_arr, ndir, nmod,
chunk_ts, chunk_fs, chunk_label, options)
# 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.phases = self.kernel.allocate_gain_array(self.gain_shape, dtype=self.ftype, zeros=True)
self.gains = np.empty_like(self.phases, dtype=self.dtype)
self.gains[:] = np.eye(self.n_cor)
self.old_gains = self.gains.copy()
def precompute_attributes(self, model_arr, flags_arr, noise):
"""
Set the initial weights to 1 and set the weights of the flags data points to 0
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.
flags_arr (np.ndarray):
Shape (n_tim, n_fre, n_ant, n_ant) array containing flags
"""
PerIntervalGains.precompute_attributes(self, model_arr, flags_arr,
noise)
self.weights_shape = [
self.n_mod, self.n_tim, self.n_fre, self.n_ant, self.n_ant, 1
]
self.weights = np.ones(self.weights_shape, dtype=self.dtype)
self.weights[:, flags_arr != 0] = 0
self._init_flags = flags_arr
unflagged = flags_arr == 0
self.num_init_unflaged_eqs = np.sum(unflagged)
self.v = 2. #t-distribution number of degrees of freedom
def __init__(self, label, data_arr, ndir, nmod, chunk_ts, chunk_fs,
chunk_label, options):
"""
Initialises a diagonal phase-only 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.
"""
PerIntervalGains.__init__(self, label, data_arr, ndir, nmod, chunk_ts,
chunk_fs, chunk_label, options,
self.get_kernel(options))
self.phases = self.cykernel.allocate_gain_array(self.gain_shape,
dtype=self.ftype,
zeros=True)
self.gains = np.empty_like(self.phases, dtype=self.dtype)
self.gains[:] = np.eye(self.n_cor)
self.old_gains = self.gains.copy()
def restrict_solution(self):
"""
Restricts the solution by invoking the inherited restrict_soultion method and applying
any machine specific restrictions.
"""
PerIntervalGains.restrict_solution(self)
if self.ref_ant is not None:
phase = np.angle(self.gains[...,self.ref_ant,(0,1),(0,1)])
self.gains *= np.exp(-1j*phase)[:,:,:,np.newaxis,:,np.newaxis]
def restrict_solution(self):
"""
Restricts the solution by invoking the inherited restrict_soultion method and applying
any machine specific restrictions.
"""
PerIntervalGains.restrict_solution(self)
if self.ref_ant is not None:
self.phases -= self.phases[:,:,:,self.ref_ant,:,:][:,:,:,np.newaxis,0,0]
for idir in self.fix_directions:
self.phases[idir, ...] = 0
def compute_update(self, model_arr, obser_arr):
"""
This function computes the update step of the weighted GN/LM method. This is
equivalent to the complete (((J^H)WJ)^-1)(J^H)WR.
Args:
obser_arr (np.array): Array containing the observed visibilities.
model_arr (np.array): Array containing the model visibilities.
gains (np.array): Array containing the current gain estimates.
jhwjinv (np.array): Array containing (J^H)WJ)^-1. (Invariant)
Returns:
update (np.array): Array containing the result of computing
(((J^H)WJ)^-1)(J^H)WR
"""
flag_count = PerIntervalGains.compute_update(model_arr, obser_arr)
#Computing the weights
resid_arr = np.empty_like(obser_arr)
residuals = self.compute_residual(obser_arr, model_arr, resid_arr)
covinv = self.compute_covinv(residuals)
self.weights, self.v = self.update_weights(residuals, covinv,
self.weights, self.v)
self.restrict_solution()
return flag_count
def __init__(self, label, data_arr, ndir, nmod, chunk_ts, chunk_fs,
chunk_label, options):
"""
Initialises a weighted complex 2x2 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.
"""
# clumsy but necessary: can't import at top level (OMP must not be touched before worker processes
# are forked off), so we import it only in here
PerIntervalGains.__init__(self, label, data_arr, ndir, nmod, chunk_ts,
chunk_fs, chunk_label, options,
self.get_kernel(options))
self.residuals = np.empty_like(data_arr)
self.save_weights = options.get("robust-save-weights", False)
self.label = label
self.cov_type = options.get(
"robust-cov", "hybrid"
) #adding an option to compute residuals covariance or just assume 1 as in Robust-t paper
self.npol = options.get(
"robust-npol", 2
) #testing if the number of polarizations really have huge effects
self.v_int = options.get("robust-int", 5)
def precompute_attributes(self, data_arr, model_arr, flags_arr, noise):
"""
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.
"""
PerIntervalGains.precompute_attributes(self, data_arr, model_arr, flags_arr, noise)
self.jhjinv = np.zeros_like(self.gains)
self.kernel_solve.compute_jhj(model_arr, self.jhjinv, self.t_int, self.f_int)
self.kernel_solve.compute_jhjinv(self.jhjinv, self.jhjinv, self.gflags, self.eps, FL.ILLCOND)
self.jhjinv = self.jhjinv.real
def importable_solutions(self, grid0):
""" Returns a dictionary of importable solutions for this machine type. """
# can import 2x2 complex gain, with a corr1/corr2 axis
solutions = PerIntervalGains.importable_solutions(self, grid0)
# but also phase (with a corr axis)
solutions['phase'] = dict(dir=grid0['dir'] if self.dd_term else [0], ant=grid0['ant'], corr=grid0['corr'], **self.interval_grid)
return solutions
def exportable_solutions():
""" Returns a dictionary of exportable solutions for this machine type. """
exportables = PerIntervalGains.exportable_solutions()
exportables.update({
"phase": (0., ("dir", "time", "freq", "ant", "corr")),
"phase.err": (0., ("dir", "time", "freq", "ant", "corr")),
})
return exportables
def __init__(self, label, data_arr, ndir, nmod,
chunk_ts, chunk_fs, chunk_label, options):
"""
Initialises a weighted complex 2x2 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.
"""
PerIntervalGains.__init__(self, label, data_arr, ndir, nmod, chunk_ts, chunk_fs, chunk_label, options)
self.kernel_robust = cubical.kernels.import_kernel("full_W_complex")
self.residuals = np.empty_like(data_arr)
self.save_weights = options.get("robust-save-weights", False)
self.label = label
self.cov_type = options.get("robust-cov", "compute") #adding an option to compute residuals covariance or just assume 1 as in Robust-t paper
self.npol = options.get("robust-npol", 2.) #testing if the number of polarizations really have huge effects
self.v_int = options.get("robust-int", 1)
self.cov_scale = options.get("robust-scale", True) # scale the covariance by n_corr*2
def __init__(self, label, model_arr, ndir, nmod, chunk_ts, chunk_fs,
chunk_label, options):
# clumsy but necessary: can't import at top level (OMP must not be touched before worker processes
# are forked off), so we import it only in here
import cubical.kernels.cyfull_W_complex
global cyfull
cyfull = cubical.kernels.cyfull_W_complex
PerIntervalGains.__init__(self, label, model_arr, ndir, nmod, chunk_ts,
chunk_fs, chunk_label, options)
self.gains = np.empty(self.gain_shape, dtype=self.dtype)
self.gains[:] = np.eye(self.n_cor)
self.old_gains = self.gains.copy()
self.weights_shape = [
self.n_mod, self.n_tim, self.n_fre, self.n_ant, self.n_ant, 1
]
self.weights = np.ones(self.weights_shape, dtype=self.dtype)
self.v = 2.
def __init__(self, label, data_arr, ndir, nmod,
chunk_ts, chunk_fs, chunk_label, options):
"""
Initialises a weighted complex 2x2 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.
"""
PerIntervalGains.__init__(self, label, data_arr, ndir, nmod, chunk_ts, chunk_fs, chunk_label, options)
if self.is_diagonal:
self.kernel_robust = cubical.kernels.import_kernel("diag_robust")
else:
self.kernel_robust = cubical.kernels.import_kernel("full_W_complex")
self.residuals = np.empty_like(data_arr)
self.save_weights = options.get("robust-save-weights", False)
self.label = label
self.cov_type = options.get("robust-cov", "compute") # adding an option to compute residuals covariance or just assume 1 as in Robust-t paper
self.npol = options.get("robust-npol", 2.) # testing if the number of polarizations really have huge effects
self.v_int = options.get("robust-int", 5)
self.cov_scale = options.get("robust-scale", 0) # scale down the covariance by this factor
self.cov_thresh = options.get("robust-cov-thresh", 1)
self.sigma_thresh = options.get("robust-sigma-thresh", 3)
self.robust_flag_weights = options.get("robust-flag-weights", True)
self.fixed_v = False
self.not_all_flagged = True
self._flag = True
self.any_new = 0
self.is_robust = True # to identify this machine as the robust solver
self._estimate_pzd = options["estimate-pzd"]
# this will be set to PZD and exp(-i*PZD) once the PZD estimate is done
self._pzd = 0
self._exp_pzd = 1
def exportable_solutions():
""" Returns a dictionary of exportable solutions for this machine type. """
sols = PerIntervalGains.exportable_solutions()
sols["pzd"] = (0.0, ("time", "freq"))
return sols