def set_sparse_weights(data_shape, psf, **kwargs):
"""Set the sparsity weights
This method defines the weights for thresholding in the sparse domain and
add them to the keyword arguments. It additionally defines the shape of the
dual variable.
Parameters
----------
data_shape : tuple
Shape of the input data array
psf : np.ndarray
PSF data (2D or 3D array)
Returns
-------
dict Updated keyword arguments
"""
# Convolve the PSF with the wavelet filters
if kwargs['psf_type'] == 'fixed':
filter_conv = (filter_convolve(np.rot90(psf, 2),
kwargs['wavelet_filters']))
filter_norm = np.array([
norm(a) * b * np.ones(data_shape[1:])
for a, b in zip(filter_conv, kwargs['wave_thresh_factor'])
])
filter_norm = np.array([filter_norm for i in xrange(data_shape[0])])
else:
filter_conv = (filter_convolve_stack(np.rot90(psf, 2),
kwargs['wavelet_filters']))
filter_norm = np.array([[
norm(b) * c * np.ones(data_shape[1:])
for b, c in zip(a, kwargs['wave_thresh_factor'])
] for a in filter_conv])
# Define a reweighting instance
kwargs['reweight'] = cwbReweight(kwargs['noise_est'] * filter_norm)
# Set the shape of the dual variable
dual_shape = ([kwargs['wavelet_filters'].shape[0]] + list(data_shape))
dual_shape[0], dual_shape[1] = dual_shape[1], dual_shape[0]
kwargs['dual_shape'] = dual_shape
return kwargs
def reconstruct_map(data, noise_est, layout, psf=None, psf_pcs=None,
psf_coef=None, wavelet_levels=1, wavelet_opt=None,
wave_thresh_factor=1, lowr_thresh_factor=1,
n_reweights=0, mode='all'):
######
# SET THE GRADIENT OPERATOR
if not isinstance(psf, type(None)):
grad_op = FixedPSF(data, psf)
else:
grad_op = PixelVariantPSF(data, psf_pcs, psf_coef)
print ' - Spetral Radius:', grad_op.spec_rad
######
# SET THE LINEAR OPERATOR
# linear_op = LinearCombo([Identity()])
# linear_op = LinearCombo([Wavelet(data.shape, wavelet_levels, wavelet_opt)])
linear_op = LinearCombo([Wavelet(data.shape, wavelet_levels, wavelet_opt),
Identity])
# print linear_op.operators[0].filters.shape
# exit()
######
# ESTIMATE THE NOISE
if not isinstance(psf, type(None)):
noise_est *= norm(convolve_mr_filters(np.rot90(psf, 2), linear_op.operators[0].filters))
# noise_est *= norm(linear_op.operators[0].op(np.rot90(psf, 2)))
noise_est *= np.ones(linear_op.operators[0].filters.shape)
else:
noise_est = np.sqrt(convolve_mr_filters(noise_est,
linear_op.operators[0].filters ** 2))
print noise_est.shape
######
# SET THE WEIGHTS
rw = cwbReweight(wave_thresh_factor * noise_est)
######
# SET RHO, SIGMA AND TAU
l1norm_filters = linear_op.operators[0].l1norm
tau = 1.0 / (grad_op.spec_rad + l1norm_filters)
sigma = tau
rho = 0.5
print ' - 1/tau - sigma||L||^2 >= beta/2:', (1 / tau - sigma *
l1norm_filters ** 2 >=
grad_op.spec_rad / 2)
######
# SET THE SHAPE OF THE DUAL
dual_shape = [linear_op.operators[0].filters.shape[0]] + list(data.shape)
######
# INITALISE THE PRIMAL AND DUAL VALUES
primal = np.ones(data.shape)
# dual = np.array([np.ones(dual_shape)])
dual = np.array([np.ones(dual_shape), np.ones(data.shape)])
grad_op.get_grad(primal)
######
# SET THE PROXIMITY OPERATORS
# prox_op = Identity()
prox_op = Positive()
# prox_dual_op = ProximityCombo([LowRankMatrix(layout, thresh_factor=3,
# grad_class=grad_op,
# grad_factor=1/sigma)])
# prox_dual_op = ProximityCombo([Threshold(rw.weights / sigma)])
prox_dual_op = ProximityCombo([Threshold(rw.weights / sigma),
LowRankMatrix(layout,
thresh_factor=lowr_thresh_factor,
grad_class=grad_op,
grad_factor=1/sigma)])
######
# SET THE COST FUNCTION.
cost_op = costTest(data, grad_op.MX)
# cost_op = posThresh(data, sigma, (grad_op, linear_op), rw.weights)
######
# PERFORM THE OPTIMISATION
opt = Condat(primal, dual, grad_op, prox_op, prox_dual_op, linear_op,
cost_op, sigma=sigma, tau=tau, print_cost=True,
auto_iterate=False)
opt.iterate(max_iter=150)
######
# REWEIGHTING
# for i in range(n_reweights):
#
# rw.reweight(linear_op.op(opt.x_new))
# prox_dual_op.update_weights(rw.weights / sigma)
# cost_op.update_weights(rw.weights)
# opt.iterate()
return opt.x_final
def reconstruct(data, noise_est, layout, psf=None, psf_pcs=None,
psf_coef=None, wavelet_levels=1, wavelet_opt=None,
wave_thresh_factor=1, lowr_thresh_factor=1,
n_reweights=0, mode='all', data_format='cube'):
######
# SET THE GRADIENT OPERATOR
if not isinstance(psf, type(None)):
grad_op = FixedPSF(data, psf, data_format=data_format)
else:
grad_op = PixelVariantPSF(data, psf_pcs, psf_coef,
data_format=data_format)
print ' - Spetral Radius:', grad_op.spec_rad
######
# SET THE LINEAR OPERATOR
if mode == 'all':
linear_op = LinearCombo([Wavelet(data.shape, wavelet_levels,
wavelet_opt, data_format=data_format),
Identity()])
elif mode == 'lowr':
linear_op = LinearCombo([Identity()])
elif mode == 'wave':
linear_op = LinearCombo([Wavelet(data.shape, wavelet_levels,
wavelet_opt, data_format=data_format)])
######
# ESTIMATE THE NOISE
if mode in ('all', 'wave'):
if not isinstance(psf, type(None)):
noise_est *= norm(convolve_mr_filters(np.rot90(psf, 2),
linear_op.operators[0].filters))
# noise_est *= norm(linear_op.operators[0].op(np.rot90(psf, 2)))
if data_format == 'map':
noise_est *= np.ones(linear_op.operators[0].filters.shape)
else:
noise_est *= np.ones([data.shape[0]] +
list(linear_op.operators[0].filters.shape))
else:
noise_est = np.sqrt(convolve_cube(noise_est,
linear_op.operators[0].filters ** 2))
print noise_est.shape
######
# SET THE WEIGHTS
rw = cwbReweight(wave_thresh_factor * noise_est)
######
# SET THE SHAPE OF THE DUAL
dual_shape = [linear_op.operators[0].filters.shape[0]] + list(data.shape)
if data_format == 'cube':
dual_shape[0], dual_shape[1] = dual_shape[1], dual_shape[0]
######
# SET RHO, SIGMA AND TAU
l1norm_filters = linear_op.operators[0].l1norm
tau = 1.0 / (grad_op.spec_rad + l1norm_filters)
sigma = tau
rho = 0.5
print ' - 1/tau - sigma||L||^2 >= beta/2:', (1 / tau - sigma *
l1norm_filters ** 2 >=
grad_op.spec_rad / 2)
######
# INITALISE THE PRIMAL AND DUAL VALUES
# 1 Primal Operator (Positivity or Identity)
primal = np.ones(data.shape)
if mode == 'all':
# 2 Dual Operators (Wavelet + Threshold and Identity + LowRankMatrix)
dual = np.empty(2, dtype=np.ndarray)
dual[0] = np.ones(dual_shape)
dual[1] = np.ones(data.shape)
elif mode == 'lowr':
# 1 Dual Operator (Identity + LowRankMatrix)
dual = np.empty(1, dtype=np.ndarray)
dual[0] = np.ones(data.shape)
elif mode == 'wave':
# 1 Dual Operator (Wavelet + Threshold or Identity + LowRankMatrix)
dual = np.empty(1, dtype=np.ndarray)
dual[0] = np.ones(dual_shape)
# Get the initial gradient value
grad_op.get_grad(primal)
print ' - Primal Variable Shape:', primal.shape
print ' - Dual Variable Shape:', dual.shape
######
# SET THE PROXIMITY OPERATORS
# prox_op = Identity()
prox_op = Positive()
if mode == 'all':
prox_dual_op = ProximityCombo([Threshold(rw.weights / sigma),
LowRankMatrix(layout,
thresh_factor=lowr_thresh_factor,
grad_class=grad_op,
grad_factor=1/sigma,
data_format=data_format)])
elif mode == 'lowr':
prox_dual_op = ProximityCombo([LowRankMatrix(layout,
thresh_factor=lowr_thresh_factor,
grad_class=grad_op,
grad_factor=1/sigma,
data_format=data_format)])
elif mode == 'wave':
prox_dual_op = ProximityCombo([Threshold(rw.weights / sigma),])
######
# SET THE COST FUNCTION.
cost_op = costTest(data, grad_op.MX)
# cost_op = posThresh(data, sigma, (grad_op, linear_op), rw.weights)
######
# PERFORM THE OPTIMISATION
opt = Condat(primal, dual, grad_op, prox_op, prox_dual_op, linear_op,
cost_op, sigma=sigma, tau=tau, print_cost=True,
auto_iterate=False)
opt.iterate(max_iter=150)
######
# REWEIGHTING
# for i in range(n_reweights):
#
# rw.reweight(linear_op.op(opt.x_new))
# prox_dual_op.update_weights(rw.weights / sigma)
# cost_op.update_weights(rw.weights)
# opt.iterate()
return opt.x_final, opt.y_final
