def restore(self):
"""
This method constructs the restoring beam and then adds the convolution to the residual.
"""
clean_beam, beam_params = beam_fit(self.psf_data,
self.psf_hdu_list[0].header)
if np.all(
np.array(self.psf_data_shape) == 2 *
np.array(self.dirty_data_shape)):
self.restored = np.fft.fftshift(
np.fft.irfft2(
np.fft.rfft2(conv.pad_array(self.model)) *
np.fft.rfft2(clean_beam)))
self.restored = self.restored[self.dirty_data_shape[0] /
2:-self.dirty_data_shape[0] / 2,
self.dirty_data_shape[1] /
2:-self.dirty_data_shape[1] / 2]
else:
self.restored = np.fft.fftshift(
np.fft.irfft2(
np.fft.rfft2(self.model) * np.fft.rfft2(clean_beam)))
self.restored += self.residual
self.restored = self.restored.astype(np.float32)
self.img_hdu_list[0].header.update('BMAJ', beam_params[0])
self.img_hdu_list[0].header.update('BMIN', beam_params[1])
self.img_hdu_list[0].header.update('BPA', beam_params[2])
def restore(self):
"""
This method constructs the restoring beam and then adds the convolution to the residual.
"""
clean_beam, beam_params = beam_fit(self.psf_data, self.cdelt1,
self.cdelt2)
if np.all(
np.array(self.psf_data_shape) == 2 *
np.array(self.dirty_data_shape)):
self.restored = np.fft.fftshift(
np.fft.irfft2(
np.fft.rfft2(conv.pad_array(self.model)) *
np.fft.rfft2(clean_beam)))
self.restored = self.restored[self.dirty_data_shape[0] /
2:-self.dirty_data_shape[0] / 2,
self.dirty_data_shape[1] /
2:-self.dirty_data_shape[1] / 2]
else:
self.restored = np.fft.fftshift(
np.fft.irfft2(
np.fft.rfft2(self.model) * np.fft.rfft2(clean_beam)))
self.restored += self.residual
self.restored = self.restored.astype(np.float32)
return beam_params
def restore(self):
"""
This method constructs the restoring beam and then adds the convolution to the residual.
"""
clean_beam, beam_params = beam_fit(self.psf_data, self.cdelt1, self.cdelt2)
if np.all(np.array(self.psf_data_shape)==2*np.array(self.dirty_data_shape)):
self.restored = np.fft.fftshift(np.fft.irfft2(np.fft.rfft2(conv.pad_array(self.model))*np.fft.rfft2(clean_beam)))
self.restored = self.restored[self.dirty_data_shape[0]/2:-self.dirty_data_shape[0]/2,
self.dirty_data_shape[1]/2:-self.dirty_data_shape[1]/2]
else:
self.restored = np.fft.fftshift(np.fft.irfft2(np.fft.rfft2(self.model)*np.fft.rfft2(clean_beam)))
self.restored += self.residual
self.restored = self.restored.astype(np.float32)
return beam_params
def restore(self):
"""
This method constructs the restoring beam and then adds the convolution to the residual.
"""
clean_beam, beam_params = beam_fit(self.psf_data, self.psf_hdu_list[0].header)
if np.all(np.array(self.psf_data_shape)==2*np.array(self.dirty_data_shape)):
self.restored = np.fft.fftshift(np.fft.irfft2(np.fft.rfft2(conv.pad_array(self.model))*np.fft.rfft2(clean_beam)))
self.restored = self.restored[self.dirty_data_shape[0]/2:-self.dirty_data_shape[0]/2,
self.dirty_data_shape[1]/2:-self.dirty_data_shape[1]/2]
else:
self.restored = np.fft.fftshift(np.fft.irfft2(np.fft.rfft2(self.model)*np.fft.rfft2(clean_beam)))
self.restored += self.residual
self.restored = self.restored.astype(np.float32)
self.img_hdu_list[0].header.update('BMAJ',beam_params[0])
self.img_hdu_list[0].header.update('BMIN',beam_params[1])
self.img_hdu_list[0].header.update('BPA',beam_params[2])
def moresane(self, subregion=None, scale_count=None, sigma_level=4, loop_gain=0.1, tolerance=0.75, accuracy=1e-6,
major_loop_miter=100, minor_loop_miter=30, all_on_gpu=False, decom_mode="ser", core_count=1,
conv_device='cpu', conv_mode='linear', extraction_mode='cpu', enforce_positivity=False,
edge_suppression=False, edge_offset=0, flux_threshold=0,
neg_comp=False, edge_excl=0, int_excl=0):
"""
Primary method for wavelet analysis and subsequent deconvolution.
INPUTS:
subregion (default=None): Size, in pixels, of the central region to be analyzed and deconvolved.
scale_count (default=None): Maximum scale to be considered - maximum scale considered during
initialisation.
sigma_level (default=4) Number of sigma at which thresholding is to be performed.
loop_gain (default=0.1): Loop gain for the deconvolution.
tolerance (default=0.75): Tolerance level for object extraction. Significant objects contain
wavelet coefficients greater than the tolerance multiplied by the
maximum wavelet coefficient in the scale under consideration.
accuracy (default=1e-6): Threshold on the standard deviation of the residual noise. Exit main
loop when this threshold is reached.
major_loop_miter (default=100): Maximum number of iterations allowed in the major loop. Exit
condition.
minor_loop_miter (default=30): Maximum number of iterations allowed in the minor loop. Serves as an
exit condition when the SNR is does not reach a maximum.
all_on_gpu (default=False): Boolean specifier to toggle all gpu modes on.
decom_mode (default='ser'): Specifier for decomposition mode - serial, multiprocessing, or gpu.
core_count (default=1): For multiprocessing, specifies the number of cores.
conv_device (default='cpu'): Specifier for device to be used - cpu or gpu.
conv_mode (default='linear'): Specifier for convolution mode - linear or circular.
extraction_mode (default='cpu'): Specifier for mode to be used - cpu or gpu.
enforce_positivity (default=False): Boolean specifier for whether or not a model must be strictly positive.
edge_suppression (default=False): Boolean specifier for whether or not the edges are to be suprressed.
edge_offset (default=0): Numeric value for an additional user-specified number of edge pixels
to be ignored. This is added to the minimum suppression.
flux_threshold (default=0): Float value, assumed to be in Jy, which specifies an approximate
convolution depth.
OUTPUTS:
self.model (no default): Model extracted by the algorithm.
self.residual (no default): Residual signal after deconvolution.
"""
# If neither subregion nor scale_count is specified, the following handles the assignment of default values.
# The default value for subregion is the whole image. The default value for scale_count is the log to the
# base two of the image dimensions minus one.
logger.info("Starting...")
if (self.dirty_data_shape[0]%2)==1:
logger.error("Image size is uneven. Please use even dimensions.")
raise ValueError("Image size is uneven. Please use even dimensions.")
if (subregion is None)|(subregion>self.dirty_data_shape[0]):
subregion = self.dirty_data_shape[0]
logger.info("Assuming subregion is {}px.".format(self.dirty_data_shape[0]))
if (scale_count is None) or (scale_count>(np.log2(self.dirty_data_shape[0])-1)):
scale_count = int(np.log2(self.dirty_data_shape[0])-1)
logger.info("Assuming maximum scale is {}.".format(scale_count))
if all_on_gpu:
decom_mode = 'gpu'
conv_device = 'gpu'
extraction_mode = 'gpu'
# The following creates arrays with dimensions equal to subregion and containing the values of the dirty
# image and psf in their central subregions.
subregion_slice = tuple([slice(self.dirty_data_shape[0]/2-subregion/2, self.dirty_data_shape[0]/2+subregion/2),
slice(self.dirty_data_shape[1]/2-subregion/2, self.dirty_data_shape[1]/2+subregion/2)])
dirty_subregion = self.dirty_data[subregion_slice]
if np.all(np.array(self.psf_data_shape)==2*np.array(self.dirty_data_shape)):
psf_subregion = self.psf_data[self.psf_data_shape[0]/2-subregion/2:self.psf_data_shape[0]/2+subregion/2,
self.psf_data_shape[1]/2-subregion/2:self.psf_data_shape[1]/2+subregion/2]
else:
psf_subregion = self.psf_data[subregion_slice]
# The following pre-loads the gpu with the fft of both the full PSF and the subregion of interest. If usegpu
# is false, this simply precomputes the fft of the PSF.
if conv_device=="gpu":
if conv_mode=="circular":
if np.all(np.array(self.psf_data_shape)==2*np.array(self.dirty_data_shape)):
psf_subregion_fft = conv.gpu_r2c_fft(psf_subregion, is_gpuarray=False, store_on_gpu=True)
psf_slice = tuple([slice(self.psf_data_shape[0]/2-self.dirty_data_shape[0]/2, self.psf_data_shape[0]/2+self.dirty_data_shape[0]/2),
slice(self.psf_data_shape[1]/2-self.dirty_data_shape[1]/2, self.psf_data_shape[1]/2+self.dirty_data_shape[1]/2)])
psf_data_fft = self.psf_data[psf_slice]
psf_data_fft = conv.gpu_r2c_fft(psf_data_fft, is_gpuarray=False, store_on_gpu=True)
else:
psf_subregion_fft = conv.gpu_r2c_fft(psf_subregion, is_gpuarray=False, store_on_gpu=True)
if psf_subregion.shape==self.psf_data_shape:
psf_data_fft = psf_subregion_fft
else:
psf_data_fft = conv.gpu_r2c_fft(self.psf_data, is_gpuarray=False, store_on_gpu=True)
if conv_mode=="linear":
if np.all(np.array(self.psf_data_shape)==2*np.array(self.dirty_data_shape)):
if np.all(np.array(self.dirty_data_shape)==subregion):
psf_subregion_fft = conv.gpu_r2c_fft(self.psf_data, is_gpuarray=False, store_on_gpu=True)
psf_data_fft = psf_subregion_fft
logger.info("Using double size PSF.")
else:
psf_slice = tuple([slice(self.psf_data_shape[0]/2-subregion, self.psf_data_shape[0]/2+subregion),
slice(self.psf_data_shape[1]/2-subregion, self.psf_data_shape[1]/2+subregion)])
psf_subregion_fft = self.psf_data[psf_slice]
psf_subregion_fft = conv.gpu_r2c_fft(psf_subregion_fft, is_gpuarray=False, store_on_gpu=True)
psf_data_fft = conv.gpu_r2c_fft(self.psf_data, is_gpuarray=False, store_on_gpu=True)
else:
if np.all(np.array(self.dirty_data_shape)==subregion):
psf_subregion_fft = conv.pad_array(self.psf_data)
psf_subregion_fft = conv.gpu_r2c_fft(psf_subregion_fft, is_gpuarray=False, store_on_gpu=True)
psf_data_fft = psf_subregion_fft
else:
psf_slice = tuple([slice(self.psf_data_shape[0]/2-subregion, self.psf_data_shape[0]/2+subregion),
slice(self.psf_data_shape[1]/2-subregion, self.psf_data_shape[1]/2+subregion)])
psf_subregion_fft = self.psf_data[psf_slice]
psf_subregion_fft = conv.gpu_r2c_fft(psf_subregion_fft, is_gpuarray=False, store_on_gpu=True)
psf_data_fft = conv.pad_array(self.psf_data)
psf_data_fft = conv.gpu_r2c_fft(psf_data_fft, is_gpuarray=False, store_on_gpu=True)
elif conv_device=="cpu":
if conv_mode=="circular":
if np.all(np.array(self.psf_data_shape)==2*np.array(self.dirty_data_shape)):
psf_subregion_fft = np.fft.rfft2(psf_subregion)
psf_slice = tuple([slice(self.psf_data_shape[0]/2-self.dirty_data_shape[0]/2, self.psf_data_shape[0]/2+self.dirty_data_shape[0]/2),
slice(self.psf_data_shape[1]/2-self.dirty_data_shape[1]/2, self.psf_data_shape[1]/2+self.dirty_data_shape[1]/2)])
psf_data_fft = self.psf_data[psf_slice]
psf_data_fft = np.fft.rfft2(psf_data_fft)
else:
psf_subregion_fft = np.fft.rfft2(psf_subregion)
if psf_subregion.shape==self.psf_data_shape:
psf_data_fft = psf_subregion_fft
else:
psf_data_fft = np.fft.rfft2(self.psf_data)
if conv_mode=="linear":
if np.all(np.array(self.psf_data_shape)==2*np.array(self.dirty_data_shape)):
if np.all(np.array(self.dirty_data_shape)==subregion):
psf_subregion_fft = np.fft.rfft2(self.psf_data)
psf_data_fft = psf_subregion_fft
logger.info("Using double size PSF.")
else:
psf_slice = tuple([slice(self.psf_data_shape[0]/2-subregion, self.psf_data_shape[0]/2+subregion),
slice(self.psf_data_shape[1]/2-subregion, self.psf_data_shape[1]/2+subregion)])
psf_subregion_fft = self.psf_data[psf_slice]
psf_subregion_fft = np.fft.rfft2(psf_subregion_fft)
psf_data_fft = np.fft.rfft2(self.psf_data)
else:
if np.all(np.array(self.dirty_data_shape)==subregion):
psf_subregion_fft = conv.pad_array(self.psf_data)
psf_subregion_fft = np.fft.rfft2(psf_subregion_fft)
psf_data_fft = psf_subregion_fft
else:
psf_slice = tuple([slice(self.psf_data_shape[0]/2-subregion, self.psf_data_shape[0]/2+subregion),
slice(self.psf_data_shape[1]/2-subregion, self.psf_data_shape[1]/2+subregion)])
psf_subregion_fft = self.psf_data[psf_slice]
psf_subregion_fft = np.fft.rfft2(psf_subregion_fft)
psf_data_fft = conv.pad_array(self.psf_data)
psf_data_fft = np.fft.rfft2(psf_data_fft)
# The following is a call to the first of the IUWT (Isotropic Undecimated Wavelet Transform) functions. This
# returns the decomposition of the PSF. The norm of each scale is found - these correspond to the energies or
# weighting factors which must be applied when locating maxima.
### REPLACE SCALECOUNT WITH: int(np.log2(self.dirty_data_shape[0])-1)
psf_decomposition = iuwt.iuwt_decomposition(psf_subregion, scale_count, mode=decom_mode, core_count=core_count)
psf_energies = np.empty([psf_decomposition.shape[0],1,1], dtype=np.float32)
for i in range(psf_energies.shape[0]):
psf_energies[i] = np.sqrt(np.sum(np.square(psf_decomposition[i,:,:])))
# INCORPORATE IF NECESSARY. POSSIBLY AT OUTER LEVEL
# psf_decomposition = psf_decomposition/psf_energies
# print(np.unravel_index(np.argmax(psf_decomposition), psf_decomposition.shape)[0])
######################################################MAJOR LOOP######################################################
major_loop_niter = 0
max_coeff = 1
model = np.zeros_like(self.dirty_data)
std_current = 1000
std_last = 1
std_ratio = 1
min_scale = 0 # The current minimum scale of interest. If this ever equals or exceeds the scale_count
# value, it will also break the following loop.
# In the case that edge_supression is desired, the following sets up a masking array.
if edge_suppression:
edge_corruption = 0
suppression_array = np.zeros([scale_count,subregion,subregion],np.float32)
for i in range(scale_count):
edge_corruption += 2*2**i
if edge_offset>edge_corruption:
suppression_array[i,edge_offset:-edge_offset, edge_offset:-edge_offset] = 1
else:
suppression_array[i,edge_corruption:-edge_corruption, edge_corruption:-edge_corruption] = 1
elif edge_offset>0:
suppression_array = np.zeros([scale_count,subregion,subregion],np.float32)
suppression_array[:,edge_offset:-edge_offset, edge_offset:-edge_offset] = 1
# The following is the major loop. Its exit conditions are reached if if the number of major loop iterations
# exceeds a user defined value, the maximum wavelet coefficient is zero or the standard deviation of the
# residual drops below a user specified accuracy threshold.
while (((major_loop_niter<major_loop_miter) & (max_coeff>0)) & ((std_ratio>accuracy)
& (np.max(dirty_subregion)>flux_threshold))):
# The first interior loop allows for the model to be re-estimated at a higher scale in the case of a poor
# SNR. If, however, a better job cannot be done, the loop will terminate.
while (min_scale<scale_count):
# This is the IUWT decomposition of the dirty image subregion up to scale_count, followed by a
# thresholding of the resulting wavelet coefficients based on the MAD estimator. This is a denoising
# operation.
if min_scale==0:
dirty_decomposition = iuwt.iuwt_decomposition(dirty_subregion, scale_count, 0, decom_mode, core_count)
thresholds = tools.estimate_threshold(dirty_decomposition, edge_excl, int_excl)
if self.mask_name is not None:
dirty_decomposition = iuwt.iuwt_decomposition(dirty_subregion*self.mask[subregion_slice], scale_count, 0,
decom_mode, core_count)
dirty_decomposition_thresh = tools.apply_threshold(dirty_decomposition, thresholds,
sigma_level=sigma_level)
# If edge_supression is desired, the following simply masks out the offending wavelet coefficients.
if edge_suppression|(edge_offset>0):
dirty_decomposition_thresh *= suppression_array
# The following calculates and stores the normalised maximum at each scale.
normalised_scale_maxima = np.empty_like(psf_energies)
for i in range(dirty_decomposition_thresh.shape[0]):
normalised_scale_maxima[i] = np.max(dirty_decomposition_thresh[i,:,:])/psf_energies[i]
# The following stores the index, scale and value of the global maximum coefficient.
max_index = np.argmax(normalised_scale_maxima[min_scale:,:,:]) + min_scale
max_scale = max_index + 1
max_coeff = normalised_scale_maxima[max_index,0,0]
# This is an escape condition for the loop. If the maximum coefficient is zero, then there is no
# useful information left in the wavelets and MORESANE is complete.
if max_coeff == 0:
logger.info("No significant wavelet coefficients detected.")
break
logger.info("Minimum scale = {}".format(min_scale))
logger.info("Maximum scale = {}".format(max_scale))
# The following constitutes a major change to the original implementation - the aim is to establish
# as soon as possible which scales are to be omitted on the current iteration. This attempts to find
# a local maxima or empty scales below the maximum scale. If either is found, that scale all those
# below it are ignored.
scale_adjust = 0
for i in range(max_index-1,-1,-1):
# if max_index > 1:
# if (normalised_scale_maxima[i,0,0] > normalised_scale_maxima[i+1,0,0]):
# scale_adjust = i + 1
# logger.info("Scale {} contains a local maxima. Ignoring scales <= {}"
# .format(scale_adjust, scale_adjust))
# break
if (normalised_scale_maxima[i,0,0] == 0):
scale_adjust = i + 1
logger.info("Scale {} is empty. Ignoring scales <= {}".format(scale_adjust, scale_adjust))
break
# We choose to only consider scales up to the scale containing the maximum wavelet coefficient,
# and ignore scales at or below the scale adjustment.
thresh_slice = dirty_decomposition_thresh[scale_adjust:max_scale,:,:]
# The following is a call to the externally defined source extraction function. It returns an array
# populated with the wavelet coefficients of structures of interest in the image. This basically refers
# to objects containing a maximum wavelet coefficient within some user-specified tolerance of the
# maximum at that scale.
extracted_sources, extracted_sources_mask = \
tools.source_extraction(thresh_slice, tolerance,
mode=extraction_mode, store_on_gpu=all_on_gpu,
neg_comp=neg_comp)
# for blah in range(extracted_sources.shape[0]):
#
# plt.imshow(extracted_sources[blah,:,:],
# interpolation="none")
# plt.show()
# The wavelet coefficients of the extracted sources are recomposed into a single image,
# which should contain only the structures of interest.
recomposed_sources = iuwt.iuwt_recomposition(extracted_sources, scale_adjust, decom_mode, core_count)
######################################################MINOR LOOP######################################################
x = np.zeros_like(recomposed_sources)
r = recomposed_sources.copy()
p = recomposed_sources.copy()
minor_loop_niter = 0
snr_last = 0
snr_current = 0
# The following is the minor loop of the algorithm. In particular, we make use of the conjugate
# gradient descent method to optimise our model. The variables have been named in order to appear
# consistent with the algorithm.
while (minor_loop_niter<minor_loop_miter):
Ap = conv.fft_convolve(p, psf_subregion_fft, conv_device, conv_mode, store_on_gpu=all_on_gpu)
Ap = iuwt.iuwt_decomposition(Ap, max_scale, scale_adjust, decom_mode, core_count,
store_on_gpu=all_on_gpu)
Ap = extracted_sources_mask*Ap
Ap = iuwt.iuwt_recomposition(Ap, scale_adjust, decom_mode, core_count)
alpha_denominator = np.dot(p.reshape(1,-1),Ap.reshape(-1,1))[0,0]
alpha_numerator = np.dot(r.reshape(1,-1),r.reshape(-1,1))[0,0]
alpha = alpha_numerator/alpha_denominator
xn = x + alpha*p
# The following enforces the positivity constraint which necessitates some recalculation.
if (np.min(xn)<0) & (enforce_positivity):
xn[xn<0] = 0
p = (xn-x)/alpha
Ap = conv.fft_convolve(p, psf_subregion_fft, conv_device, conv_mode, store_on_gpu=all_on_gpu)
Ap = iuwt.iuwt_decomposition(Ap, max_scale, scale_adjust, decom_mode, core_count,
store_on_gpu=all_on_gpu)
Ap = extracted_sources_mask*Ap
Ap = iuwt.iuwt_recomposition(Ap, scale_adjust, decom_mode, core_count)
rn = r - alpha*Ap
beta_numerator = np.dot(rn.reshape(1,-1), rn.reshape(-1,1))[0,0]
beta_denominator = np.dot(r.reshape(1,-1), r.reshape(-1,1))[0,0]
beta = beta_numerator/beta_denominator
p = rn + beta*p
model_sources = conv.fft_convolve(xn, psf_subregion_fft, conv_device, conv_mode, store_on_gpu=all_on_gpu)
model_sources = iuwt.iuwt_decomposition(model_sources, max_scale, scale_adjust, decom_mode,
core_count, store_on_gpu=all_on_gpu)
model_sources = extracted_sources_mask*model_sources
if all_on_gpu:
model_sources = model_sources.get()
# We compare our model to the sources extracted from the data.
snr_last = snr_current
snr_current = tools.snr_ratio(extracted_sources, model_sources)
minor_loop_niter += 1
logger.debug("SNR at iteration {0} = {1}".format(minor_loop_niter, snr_current))
# The following flow control determines whether or not the model is adequate and if a
# recalculation is required.
if (minor_loop_niter==1)&(snr_current>40):
logger.info("SNR too large on first iteration - false detection. "
"Incrementing the minimum scale.")
min_scale += 1
break
if snr_current>40:
logger.info("Model has reached <1% error - exiting minor loop.")
x = xn
min_scale = 0
break
if (minor_loop_niter>2)&(snr_current<=snr_last):
if (snr_current>10.5):
logger.info("SNR has decreased - Model has reached ~{}% error - exiting minor loop." \
.format(int(100/np.power(10,snr_current/20))))
min_scale = 0
break
else:
logger.info("SNR has decreased - SNR too small. Incrementing the minimum scale.")
min_scale += 1
break
r = rn
x = xn
logger.info("{} minor loop iterations performed.".format(minor_loop_niter))
if ((minor_loop_niter==minor_loop_miter)&(snr_current>10.5)):
logger.info("Maximum number of minor loop iterations exceeded. Model reached ~{}% error." \
.format(int(100/np.power(10,snr_current/20))))
min_scale = 0
break
if (min_scale==0):
break
###################################################END OF MINOR LOOP###################################################
if min_scale==scale_count:
logger.info("All scales are performing poorly - stopping.")
break
# The following handles the deconvolution step. The model convolved with the psf is subtracted from the
# dirty image to give the residual.
if max_coeff>0:
# x[abs(x)<0.8*np.max(np.abs(x))] = 0
model[subregion_slice] += loop_gain*x
residual = self.dirty_data - conv.fft_convolve(model, psf_data_fft, conv_device, conv_mode)
# The following assesses whether or not the residual has improved.
std_last = std_current
std_current = np.std(residual[subregion_slice])
std_ratio = (std_last-std_current)/std_last
# If the most recent deconvolution step is poor, the following reverts the changes so that the
# previous model and residual are preserved.
if std_ratio<0:
logger.info("Residual has worsened - reverting changes.")
model[subregion_slice] -= loop_gain*x
residual = self.dirty_data - conv.fft_convolve(model, psf_data_fft, conv_device, conv_mode)
# The current residual becomes the dirty image for the subsequent iteration.
dirty_subregion = residual[subregion_slice]
major_loop_niter += 1
logger.info("{} major loop iterations performed.".format(major_loop_niter))
# The following condition will only trigger if MORESANE did no work - this is an exit condition for the
# by-scale approach.
if (major_loop_niter==0):
logger.info("Current MORESANE iteration did no work - finished.")
self.complete = True
break
# If MORESANE did work at the current iteration, the following simply updates the values in the class
# variables self.model and self.residual.
if major_loop_niter>0:
self.model += model
self.residual = residual
def moresane(self,
subregion=None,
scale_count=None,
sigma_level=4,
loop_gain=0.1,
tolerance=0.75,
accuracy=1e-6,
major_loop_miter=100,
minor_loop_miter=30,
all_on_gpu=False,
decom_mode="ser",
core_count=1,
conv_device='cpu',
conv_mode='linear',
extraction_mode='cpu',
enforce_positivity=False,
edge_suppression=False,
edge_offset=0,
flux_threshold=0,
neg_comp=False,
edge_excl=0,
int_excl=0):
"""
Primary method for wavelet analysis and subsequent deconvolution.
INPUTS:
subregion (default=None): Size, in pixels, of the central region to be analyzed and deconvolved.
scale_count (default=None): Maximum scale to be considered - maximum scale considered during
initialisation.
sigma_level (default=4) Number of sigma at which thresholding is to be performed.
loop_gain (default=0.1): Loop gain for the deconvolution.
tolerance (default=0.75): Tolerance level for object extraction. Significant objects contain
wavelet coefficients greater than the tolerance multiplied by the
maximum wavelet coefficient in the scale under consideration.
accuracy (default=1e-6): Threshold on the standard deviation of the residual noise. Exit main
loop when this threshold is reached.
major_loop_miter (default=100): Maximum number of iterations allowed in the major loop. Exit
condition.
minor_loop_miter (default=30): Maximum number of iterations allowed in the minor loop. Serves as an
exit condition when the SNR is does not reach a maximum.
all_on_gpu (default=False): Boolean specifier to toggle all gpu modes on.
decom_mode (default='ser'): Specifier for decomposition mode - serial, multiprocessing, or gpu.
core_count (default=1): For multiprocessing, specifies the number of cores.
conv_device (default='cpu'): Specifier for device to be used - cpu or gpu.
conv_mode (default='linear'): Specifier for convolution mode - linear or circular.
extraction_mode (default='cpu'): Specifier for mode to be used - cpu or gpu.
enforce_positivity (default=False): Boolean specifier for whether or not a model must be strictly positive.
edge_suppression (default=False): Boolean specifier for whether or not the edges are to be suprressed.
edge_offset (default=0): Numeric value for an additional user-specified number of edge pixels
to be ignored. This is added to the minimum suppression.
flux_threshold (default=0): Float value, assumed to be in Jy, which specifies an approximate
convolution depth.
OUTPUTS:
self.model (no default): Model extracted by the algorithm.
self.residual (no default): Residual signal after deconvolution.
"""
# If neither subregion nor scale_count is specified, the following handles the assignment of default values.
# The default value for subregion is the whole image. The default value for scale_count is the log to the
# base two of the image dimensions minus one.
logger.info("Starting...")
if (self.dirty_data_shape[0] % 2) == 1:
logger.error("Image size is uneven. Please use even dimensions.")
raise ValueError(
"Image size is uneven. Please use even dimensions.")
if (subregion is None) | (subregion > self.dirty_data_shape[0]):
subregion = self.dirty_data_shape[0]
logger.info("Assuming subregion is {}px.".format(
self.dirty_data_shape[0]))
if (scale_count is None) or (scale_count >
(np.log2(self.dirty_data_shape[0]) - 1)):
scale_count = int(np.log2(self.dirty_data_shape[0]) - 1)
logger.info("Assuming maximum scale is {}.".format(scale_count))
if all_on_gpu:
decom_mode = 'gpu'
conv_device = 'gpu'
extraction_mode = 'gpu'
# The following creates arrays with dimensions equal to subregion and containing the values of the dirty
# image and psf in their central subregions.
subregion_slice = tuple([
slice(self.dirty_data_shape[0] / 2 - subregion / 2,
self.dirty_data_shape[0] / 2 + subregion / 2),
slice(self.dirty_data_shape[1] / 2 - subregion / 2,
self.dirty_data_shape[1] / 2 + subregion / 2)
])
dirty_subregion = self.dirty_data[subregion_slice]
if np.all(
np.array(self.psf_data_shape) == 2 *
np.array(self.dirty_data_shape)):
psf_subregion = self.psf_data[
self.psf_data_shape[0] / 2 -
subregion / 2:self.psf_data_shape[0] / 2 + subregion / 2,
self.psf_data_shape[1] / 2 -
subregion / 2:self.psf_data_shape[1] / 2 + subregion / 2]
else:
psf_subregion = self.psf_data[subregion_slice]
# The following pre-loads the gpu with the fft of both the full PSF and the subregion of interest. If usegpu
# is false, this simply precomputes the fft of the PSF.
if conv_device == "gpu":
if conv_mode == "circular":
if np.all(
np.array(self.psf_data_shape) == 2 *
np.array(self.dirty_data_shape)):
psf_subregion_fft = conv.gpu_r2c_fft(psf_subregion,
is_gpuarray=False,
store_on_gpu=True)
psf_slice = tuple([
slice(
self.psf_data_shape[0] / 2 -
self.dirty_data_shape[0] / 2,
self.psf_data_shape[0] / 2 +
self.dirty_data_shape[0] / 2),
slice(
self.psf_data_shape[1] / 2 -
self.dirty_data_shape[1] / 2,
self.psf_data_shape[1] / 2 +
self.dirty_data_shape[1] / 2)
])
psf_data_fft = self.psf_data[psf_slice]
psf_data_fft = conv.gpu_r2c_fft(psf_data_fft,
is_gpuarray=False,
store_on_gpu=True)
else:
psf_subregion_fft = conv.gpu_r2c_fft(psf_subregion,
is_gpuarray=False,
store_on_gpu=True)
if psf_subregion.shape == self.psf_data_shape:
psf_data_fft = psf_subregion_fft
else:
psf_data_fft = conv.gpu_r2c_fft(self.psf_data,
is_gpuarray=False,
store_on_gpu=True)
if conv_mode == "linear":
if np.all(
np.array(self.psf_data_shape) == 2 *
np.array(self.dirty_data_shape)):
if np.all(np.array(self.dirty_data_shape) == subregion):
psf_subregion_fft = conv.gpu_r2c_fft(self.psf_data,
is_gpuarray=False,
store_on_gpu=True)
psf_data_fft = psf_subregion_fft
logger.info("Using double size PSF.")
else:
psf_slice = tuple([
slice(self.psf_data_shape[0] / 2 - subregion,
self.psf_data_shape[0] / 2 + subregion),
slice(self.psf_data_shape[1] / 2 - subregion,
self.psf_data_shape[1] / 2 + subregion)
])
psf_subregion_fft = self.psf_data[psf_slice]
psf_subregion_fft = conv.gpu_r2c_fft(psf_subregion_fft,
is_gpuarray=False,
store_on_gpu=True)
psf_data_fft = conv.gpu_r2c_fft(self.psf_data,
is_gpuarray=False,
store_on_gpu=True)
else:
if np.all(np.array(self.dirty_data_shape) == subregion):
psf_subregion_fft = conv.pad_array(self.psf_data)
psf_subregion_fft = conv.gpu_r2c_fft(psf_subregion_fft,
is_gpuarray=False,
store_on_gpu=True)
psf_data_fft = psf_subregion_fft
else:
psf_slice = tuple([
slice(self.psf_data_shape[0] / 2 - subregion,
self.psf_data_shape[0] / 2 + subregion),
slice(self.psf_data_shape[1] / 2 - subregion,
self.psf_data_shape[1] / 2 + subregion)
])
psf_subregion_fft = self.psf_data[psf_slice]
psf_subregion_fft = conv.gpu_r2c_fft(psf_subregion_fft,
is_gpuarray=False,
store_on_gpu=True)
psf_data_fft = conv.pad_array(self.psf_data)
psf_data_fft = conv.gpu_r2c_fft(psf_data_fft,
is_gpuarray=False,
store_on_gpu=True)
elif conv_device == "cpu":
if conv_mode == "circular":
if np.all(
np.array(self.psf_data_shape) == 2 *
np.array(self.dirty_data_shape)):
psf_subregion_fft = np.fft.rfft2(psf_subregion)
psf_slice = tuple([
slice(
self.psf_data_shape[0] / 2 -
self.dirty_data_shape[0] / 2,
self.psf_data_shape[0] / 2 +
self.dirty_data_shape[0] / 2),
slice(
self.psf_data_shape[1] / 2 -
self.dirty_data_shape[1] / 2,
self.psf_data_shape[1] / 2 +
self.dirty_data_shape[1] / 2)
])
psf_data_fft = self.psf_data[psf_slice]
psf_data_fft = np.fft.rfft2(psf_data_fft)
else:
psf_subregion_fft = np.fft.rfft2(psf_subregion)
if psf_subregion.shape == self.psf_data_shape:
psf_data_fft = psf_subregion_fft
else:
psf_data_fft = np.fft.rfft2(self.psf_data)
if conv_mode == "linear":
if np.all(
np.array(self.psf_data_shape) == 2 *
np.array(self.dirty_data_shape)):
if np.all(np.array(self.dirty_data_shape) == subregion):
psf_subregion_fft = np.fft.rfft2(self.psf_data)
psf_data_fft = psf_subregion_fft
logger.info("Using double size PSF.")
else:
psf_slice = tuple([
slice(self.psf_data_shape[0] / 2 - subregion,
self.psf_data_shape[0] / 2 + subregion),
slice(self.psf_data_shape[1] / 2 - subregion,
self.psf_data_shape[1] / 2 + subregion)
])
psf_subregion_fft = self.psf_data[psf_slice]
psf_subregion_fft = np.fft.rfft2(psf_subregion_fft)
psf_data_fft = np.fft.rfft2(self.psf_data)
else:
if np.all(np.array(self.dirty_data_shape) == subregion):
psf_subregion_fft = conv.pad_array(self.psf_data)
psf_subregion_fft = np.fft.rfft2(psf_subregion_fft)
psf_data_fft = psf_subregion_fft
else:
psf_slice = tuple([
slice(self.psf_data_shape[0] / 2 - subregion,
self.psf_data_shape[0] / 2 + subregion),
slice(self.psf_data_shape[1] / 2 - subregion,
self.psf_data_shape[1] / 2 + subregion)
])
psf_subregion_fft = self.psf_data[psf_slice]
psf_subregion_fft = np.fft.rfft2(psf_subregion_fft)
psf_data_fft = conv.pad_array(self.psf_data)
psf_data_fft = np.fft.rfft2(psf_data_fft)
# The following is a call to the first of the IUWT (Isotropic Undecimated Wavelet Transform) functions. This
# returns the decomposition of the PSF. The norm of each scale is found - these correspond to the energies or
# weighting factors which must be applied when locating maxima.
### REPLACE SCALECOUNT WITH: int(np.log2(self.dirty_data_shape[0])-1)
psf_decomposition = iuwt.iuwt_decomposition(psf_subregion,
scale_count,
mode=decom_mode,
core_count=core_count)
psf_energies = np.empty([psf_decomposition.shape[0], 1, 1],
dtype=np.float32)
for i in range(psf_energies.shape[0]):
psf_energies[i] = np.sqrt(
np.sum(np.square(psf_decomposition[i, :, :])))
# INCORPORATE IF NECESSARY. POSSIBLY AT OUTER LEVEL
# psf_decomposition = psf_decomposition/psf_energies
# print(np.unravel_index(np.argmax(psf_decomposition), psf_decomposition.shape)[0])
######################################################MAJOR LOOP######################################################
major_loop_niter = 0
max_coeff = 1
model = np.zeros_like(self.dirty_data)
std_current = 1000
std_last = 1
std_ratio = 1
min_scale = 0 # The current minimum scale of interest. If this ever equals or exceeds the scale_count
# value, it will also break the following loop.
# In the case that edge_supression is desired, the following sets up a masking array.
if edge_suppression:
edge_corruption = 0
suppression_array = np.zeros([scale_count, subregion, subregion],
np.float32)
for i in range(scale_count):
edge_corruption += 2 * 2**i
if edge_offset > edge_corruption:
suppression_array[i, edge_offset:-edge_offset,
edge_offset:-edge_offset] = 1
else:
suppression_array[i, edge_corruption:-edge_corruption,
edge_corruption:-edge_corruption] = 1
elif edge_offset > 0:
suppression_array = np.zeros([scale_count, subregion, subregion],
np.float32)
suppression_array[:, edge_offset:-edge_offset,
edge_offset:-edge_offset] = 1
# The following is the major loop. Its exit conditions are reached if if the number of major loop iterations
# exceeds a user defined value, the maximum wavelet coefficient is zero or the standard deviation of the
# residual drops below a user specified accuracy threshold.
while (((major_loop_niter < major_loop_miter) & (max_coeff > 0)) &
((std_ratio > accuracy)
& (np.max(dirty_subregion) > flux_threshold))):
# The first interior loop allows for the model to be re-estimated at a higher scale in the case of a poor
# SNR. If, however, a better job cannot be done, the loop will terminate.
while (min_scale < scale_count):
# This is the IUWT decomposition of the dirty image subregion up to scale_count, followed by a
# thresholding of the resulting wavelet coefficients based on the MAD estimator. This is a denoising
# operation.
if min_scale == 0:
dirty_decomposition = iuwt.iuwt_decomposition(
dirty_subregion, scale_count, 0, decom_mode,
core_count)
thresholds = tools.estimate_threshold(
dirty_decomposition, edge_excl, int_excl)
if self.mask_name is not None:
dirty_decomposition = iuwt.iuwt_decomposition(
dirty_subregion * self.mask[subregion_slice],
scale_count, 0, decom_mode, core_count)
dirty_decomposition_thresh = tools.apply_threshold(
dirty_decomposition,
thresholds,
sigma_level=sigma_level)
# If edge_supression is desired, the following simply masks out the offending wavelet coefficients.
if edge_suppression | (edge_offset > 0):
dirty_decomposition_thresh *= suppression_array
# The following calculates and stores the normalised maximum at each scale.
normalised_scale_maxima = np.empty_like(psf_energies)
for i in range(dirty_decomposition_thresh.shape[0]):
normalised_scale_maxima[i] = np.max(
dirty_decomposition_thresh[
i, :, :]) / psf_energies[i]
# The following stores the index, scale and value of the global maximum coefficient.
max_index = np.argmax(
normalised_scale_maxima[min_scale:, :, :]) + min_scale
max_scale = max_index + 1
max_coeff = normalised_scale_maxima[max_index, 0, 0]
# This is an escape condition for the loop. If the maximum coefficient is zero, then there is no
# useful information left in the wavelets and MORESANE is complete.
if max_coeff == 0:
logger.info(
"No significant wavelet coefficients detected.")
break
logger.info("Minimum scale = {}".format(min_scale))
logger.info("Maximum scale = {}".format(max_scale))
# The following constitutes a major change to the original implementation - the aim is to establish
# as soon as possible which scales are to be omitted on the current iteration. This attempts to find
# a local maxima or empty scales below the maximum scale. If either is found, that scale all those
# below it are ignored.
scale_adjust = 0
for i in range(max_index - 1, -1, -1):
# if max_index > 1:
# if (normalised_scale_maxima[i,0,0] > normalised_scale_maxima[i+1,0,0]):
# scale_adjust = i + 1
# logger.info("Scale {} contains a local maxima. Ignoring scales <= {}"
# .format(scale_adjust, scale_adjust))
# break
if (normalised_scale_maxima[i, 0, 0] == 0):
scale_adjust = i + 1
logger.info(
"Scale {} is empty. Ignoring scales <= {}".format(
scale_adjust, scale_adjust))
break
# We choose to only consider scales up to the scale containing the maximum wavelet coefficient,
# and ignore scales at or below the scale adjustment.
thresh_slice = dirty_decomposition_thresh[
scale_adjust:max_scale, :, :]
# The following is a call to the externally defined source extraction function. It returns an array
# populated with the wavelet coefficients of structures of interest in the image. This basically refers
# to objects containing a maximum wavelet coefficient within some user-specified tolerance of the
# maximum at that scale.
extracted_sources, extracted_sources_mask = \
tools.source_extraction(thresh_slice, tolerance,
mode=extraction_mode, store_on_gpu=all_on_gpu,
neg_comp=neg_comp)
# for blah in range(extracted_sources.shape[0]):
#
# plt.imshow(extracted_sources[blah,:,:],
# interpolation="none")
# plt.show()
# The wavelet coefficients of the extracted sources are recomposed into a single image,
# which should contain only the structures of interest.
recomposed_sources = iuwt.iuwt_recomposition(
extracted_sources, scale_adjust, decom_mode, core_count)
######################################################MINOR LOOP######################################################
x = np.zeros_like(recomposed_sources)
r = recomposed_sources.copy()
p = recomposed_sources.copy()
minor_loop_niter = 0
snr_last = 0
snr_current = 0
# The following is the minor loop of the algorithm. In particular, we make use of the conjugate
# gradient descent method to optimise our model. The variables have been named in order to appear
# consistent with the algorithm.
while (minor_loop_niter < minor_loop_miter):
Ap = conv.fft_convolve(p,
psf_subregion_fft,
conv_device,
conv_mode,
store_on_gpu=all_on_gpu)
Ap = iuwt.iuwt_decomposition(Ap,
max_scale,
scale_adjust,
decom_mode,
core_count,
store_on_gpu=all_on_gpu)
Ap = extracted_sources_mask * Ap
Ap = iuwt.iuwt_recomposition(Ap, scale_adjust, decom_mode,
core_count)
alpha_denominator = np.dot(p.reshape(1, -1),
Ap.reshape(-1, 1))[0, 0]
alpha_numerator = np.dot(r.reshape(1, -1),
r.reshape(-1, 1))[0, 0]
alpha = alpha_numerator / alpha_denominator
xn = x + alpha * p
# The following enforces the positivity constraint which necessitates some recalculation.
if (np.min(xn) < 0) & (enforce_positivity):
xn[xn < 0] = 0
p = (xn - x) / alpha
Ap = conv.fft_convolve(p,
psf_subregion_fft,
conv_device,
conv_mode,
store_on_gpu=all_on_gpu)
Ap = iuwt.iuwt_decomposition(Ap,
max_scale,
scale_adjust,
decom_mode,
core_count,
store_on_gpu=all_on_gpu)
Ap = extracted_sources_mask * Ap
Ap = iuwt.iuwt_recomposition(Ap, scale_adjust,
decom_mode, core_count)
rn = r - alpha * Ap
beta_numerator = np.dot(rn.reshape(1, -1),
rn.reshape(-1, 1))[0, 0]
beta_denominator = np.dot(r.reshape(1, -1),
r.reshape(-1, 1))[0, 0]
beta = beta_numerator / beta_denominator
p = rn + beta * p
model_sources = conv.fft_convolve(xn,
psf_subregion_fft,
conv_device,
conv_mode,
store_on_gpu=all_on_gpu)
model_sources = iuwt.iuwt_decomposition(
model_sources,
max_scale,
scale_adjust,
decom_mode,
core_count,
store_on_gpu=all_on_gpu)
model_sources = extracted_sources_mask * model_sources
if all_on_gpu:
model_sources = model_sources.get()
# We compare our model to the sources extracted from the data.
snr_last = snr_current
snr_current = tools.snr_ratio(extracted_sources,
model_sources)
minor_loop_niter += 1
logger.debug("SNR at iteration {0} = {1}".format(
minor_loop_niter, snr_current))
# The following flow control determines whether or not the model is adequate and if a
# recalculation is required.
if (minor_loop_niter == 1) & (snr_current > 40):
logger.info(
"SNR too large on first iteration - false detection. "
"Incrementing the minimum scale.")
min_scale += 1
break
if snr_current > 40:
logger.info(
"Model has reached <1% error - exiting minor loop."
)
x = xn
min_scale = 0
break
if (minor_loop_niter > 2) & (snr_current <= snr_last):
if (snr_current > 10.5):
logger.info("SNR has decreased - Model has reached ~{}% error - exiting minor loop." \
.format(int(100/np.power(10,snr_current/20))))
min_scale = 0
break
else:
logger.info(
"SNR has decreased - SNR too small. Incrementing the minimum scale."
)
min_scale += 1
break
r = rn
x = xn
logger.info("{} minor loop iterations performed.".format(
minor_loop_niter))
if ((minor_loop_niter == minor_loop_miter) &
(snr_current > 10.5)):
logger.info("Maximum number of minor loop iterations exceeded. Model reached ~{}% error." \
.format(int(100/np.power(10,snr_current/20))))
min_scale = 0
break
if (min_scale == 0):
break
###################################################END OF MINOR LOOP###################################################
if min_scale == scale_count:
logger.info("All scales are performing poorly - stopping.")
break
# The following handles the deconvolution step. The model convolved with the psf is subtracted from the
# dirty image to give the residual.
if max_coeff > 0:
# x[abs(x)<0.8*np.max(np.abs(x))] = 0
model[subregion_slice] += loop_gain * x
residual = self.dirty_data - conv.fft_convolve(
model, psf_data_fft, conv_device, conv_mode)
# The following assesses whether or not the residual has improved.
std_last = std_current
std_current = np.std(residual[subregion_slice])
std_ratio = (std_last - std_current) / std_last
# If the most recent deconvolution step is poor, the following reverts the changes so that the
# previous model and residual are preserved.
if std_ratio < 0:
logger.info("Residual has worsened - reverting changes.")
model[subregion_slice] -= loop_gain * x
residual = self.dirty_data - conv.fft_convolve(
model, psf_data_fft, conv_device, conv_mode)
# The current residual becomes the dirty image for the subsequent iteration.
dirty_subregion = residual[subregion_slice]
major_loop_niter += 1
logger.info("{} major loop iterations performed.".format(
major_loop_niter))
# The following condition will only trigger if MORESANE did no work - this is an exit condition for the
# by-scale approach.
if (major_loop_niter == 0):
logger.info(
"Current MORESANE iteration did no work - finished.")
self.complete = True
break
# If MORESANE did work at the current iteration, the following simply updates the values in the class
# variables self.model and self.residual.
if major_loop_niter > 0:
self.model += model
self.residual = residual