From 50af0befbcd0f6adc440ab42343c00ea03baf641 Mon Sep 17 00:00:00 2001 From: Chris Clark Date: Fri, 19 Mar 2021 01:26:27 -0400 Subject: [PATCH] Initial commit of FourierFuncs This is the version of these functions used to perform the feathered presented in Clark et al., soon to be submitted. --- FourierFuncs.py | 270 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 270 insertions(+) create mode 100644 FourierFuncs.py diff --git a/FourierFuncs.py b/FourierFuncs.py new file mode 100644 index 0000000..53e0c39 --- /dev/null +++ b/FourierFuncs.py @@ -0,0 +1,270 @@ +# Define function for clever fourier combination of images, +# Inputs: HDU containing low-res data; HDU containing high-res data; array of low-res beam gridded to the pixel scale of high-res image; array of high-res beam gridded to pixel scale of high-res image(); (optional boolean/float for giving angular scale in degrees at which to apply a tapering transition; boolean of whether to employ subpixel low-pass filter to low-res image to remove pixel edge artefacts; boolean/string for saving combined image to file instead of returning; boolean for also returning additional data; boolean describing whether to use beam-mediated feather and only use window for cross-calibrating +# Outputs: The combined image +def FourierCombine(lores_hdu, hires_hdu, lores_beam_img, hires_beam_img, + taper_cutoffs_deg=False, apodise=False, to_file=False, return_all=False, beam_cross_corr=False): + + # If input images are being provided as paths to files, discern this and read them in + lores_hdu_path = False + if isinstance(lores_hdu, str): + lores_hdu_path = lores_hdu + lores_hdu = astropy.io.fits.PrimaryHDU(data=astropy.io.fits.getdata(lores_hdu), + header=astropy.io.fits.getheader(lores_hdu)) + hires_hdu_path = False + if isinstance(hires_hdu, str): + hires_hdu_path = hires_hdu + hires_hdu = astropy.io.fits.PrimaryHDU(data=astropy.io.fits.getdata(hires_hdu), + header=astropy.io.fits.getheader(hires_hdu)) + if isinstance(lores_beam_img, str): + lores_beam_img = astropy.io.fits.getdata(lores_beam_img) + if isinstance(hires_beam_img, str): + hires_beam_img = astropy.io.fits.getdata(hires_beam_img) + + # Make clean copies of input arrays, then put into float32, and delete HDUs, to save memory + lores_img = lores_hdu.data.copy().astype(np.float32) + lores_hdr = lores_hdu.header.copy() + hires_img = hires_hdu.data.copy().astype(np.float32) + hires_img_orig = hires_hdu.data.copy().astype(np.float32) + hires_hdr = hires_hdu.header.copy() + lores_beam_img = lores_beam_img.copy().astype(np.float32) + hires_beam_img = hires_beam_img.copy().astype(np.float32) + del(lores_hdu, hires_hdu) + + # Calculate low -resolution pixel size + lores_wcs = astropy.wcs.WCS(lores_hdr) + lores_pix_width_arcsec = 3600.0 * np.abs(np.max(lores_wcs.pixel_scale_matrix)) + + # Calculate high-reoslution pixel size + hires_wcs = astropy.wcs.WCS(hires_hdr) + hires_pix_width_arcsec = 3600.0 * np.abs(np.max(hires_wcs.pixel_scale_matrix)) + hires_pix_width_deg = hires_pix_width_arcsec / 3600.0 + + # Impute (temporarily) any NaNs surrounding the coverage region with the clipped average of the data (so that the fourier transformers play nice) + hires_img[np.where(np.isnan(hires_img))] = SigmaClip(hires_img, median=True, sigma_thresh=1.0)[1] + + # Temporarily interpolate over any NaNs in low-resolution data, reproject to high-resolution pixel scale + lores_img = ImputeImage(lores_img) + lores_img = astropy.convolution.interpolate_replace_nans(lores_img, astropy.convolution.Gaussian2DKernel(3), + astropy.convolution.convolve_fft, allow_huge=True, boundary='wrap').astype(np.float32) + lores_img = reproject.reproject_interp((lores_img, lores_hdr), hires_hdr, order='bicubic')[0].astype(np.float32) # Ie, following how SWarp supersamples images + lores_edge = np.zeros(lores_img.shape) + lores_edge[np.where(np.isnan(lores_img))] = 1.0 + lores_img[np.where(lores_edge == 1.0)] = np.nanmedian(lores_img) + + # If requested, low-pass filter low-resolution data to remove pixel-edge effects + if apodise: + lores_apodisation_kernel_sigma = 0.5 * 2.0**-0.5 * (lores_pix_width_arcsec / hires_pix_width_arcsec) # 0.5 * 2.0**-0.5 + lores_apodisation_kernel = astropy.convolution.Gaussian2DKernel(lores_apodisation_kernel_sigma).array + lores_img = astropy.convolution.convolve_fft(lores_img, lores_apodisation_kernel, + boundary='reflect', allow_huge=True, preserve_nan=False) # As NaNs already removed + + # Incorporate apodisation filter into the low-resolution beam (as we have to account for the fact that its resolution is now ever-so-slightly lower) + lores_beam_img = astropy.convolution.convolve_fft(lores_beam_img, lores_apodisation_kernel, + boundary='reflect', allow_huge=True, preserve_nan=False) + lores_beam_img -= np.min(lores_beam_img) + lores_beam_img /= np.sum(lores_beam_img) + + # Fourier transform all the things + hires_beam_fourier = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(hires_beam_img))).astype(np.complex64) + lores_beam_fourier = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(lores_beam_img))).astype(np.complex64) + hires_fourier = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(hires_img))).astype(np.complex64) + lores_fourier = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(lores_img))).astype(np.complex64) + del(hires_img, lores_img, hires_beam_img) + + # Add miniscule offset to any zero-value elements to stop inf and nan values from appearing later. + lores_beam_fourier.real[np.where(lores_beam_fourier.real == 0)] = 1E-50 + + # Divide the low-resolution data by the low-resolution beam (ie, deconvolve it), then multiply by the high-resoluiton beam, to normalise amplitudes + fourier_norm = hires_beam_fourier / lores_beam_fourier + fourier_norm[np.where(np.isinf(fourier_norm))] = 1E-50 + lores_fourier *= fourier_norm + + # If requested, start by cross-calibrating the hires and lores data within the tapering angular window + if taper_cutoffs_deg != False: + hires_fourier_corr_factor = FourierCalibrate(lores_fourier, hires_fourier, taper_cutoffs_deg, hires_pix_width_deg) + hires_fourier_corr = hires_fourier * hires_fourier_corr_factor[0] + + # Perform tapering between specificed angular scales to weight data in Fourier space, following a Hann filter profile + taper_filter = FourierTaper(taper_cutoffs_deg, hires_wcs) + hires_weight = 1.0 - taper_filter + hires_fourier_weighted = hires_fourier_corr.copy() + hires_fourier_weighted *= hires_weight + lores_weight = taper_filter + lores_fourier_weighted = lores_fourier.copy() + lores_fourier_weighted *= lores_weight + + # Otherwise, in standard operation, use low-resolution beam to weight the tapering from low-resolution to high-resolution data + elif taper_cutoffs_deg == False: + hires_fourier_corr_factor = [1.0, 0.0] + hires_weight = 1.0 - lores_beam_fourier + hires_fourier_weighted = hires_fourier * hires_weight + lores_weight = 1.0 * lores_beam_fourier + lores_fourier_weighted = lores_fourier * lores_weight + + # Remove edge effects from high-resolution map, where possible + hires_weighted_img = np.fft.fftshift(np.real(np.fft.ifft2(np.fft.ifftshift(hires_fourier_weighted)))).astype(np.float32) + hires_weighted_img[np.where(np.isnan(hires_img_orig))] = SigmaClip(hires_weighted_img, median=True, sigma_thresh=1.0)[1] + hires_fourier_weighted = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(hires_weighted_img))).astype(np.complex64) + + # Combine the images, then convert back out of Fourier space + comb_fourier = lores_fourier_weighted + hires_fourier_weighted + comb_fourier_shift = np.fft.ifftshift(comb_fourier) + comb_img = np.fft.fftshift(np.real(np.fft.ifft2(comb_fourier_shift))).astype(np.float32) + del(hires_weight,lores_weight,hires_fourier_weighted, + lores_fourier_weighted,hires_weighted_img,comb_fourier,comb_fourier_shift) + + # Estimate the size of the low-resolution beam + lores_beam_demislice = lores_beam_img[int(round(0.5*lores_beam_img.shape[1])):,int(round(0.5*lores_beam_img.shape[1]))] + lores_beam_width_pix = float(np.argmin(np.abs(lores_beam_demislice - (0.5 * lores_beam_demislice.max())))) + + # Purely for prettiness sake, identify any messy transition region between high- and low-resolution + hires_mask = hires_img_orig.copy() + hires_mask = (hires_mask * 0.0) + 1.0 + hires_mask[np.where(np.isnan(hires_mask))] = 0.0 + hires_mask_dilated_out = scipy.ndimage.binary_dilation(hires_mask, iterations=int(1.5*round(lores_beam_width_pix))).astype(int) + hires_mask = -1.0 * (hires_mask - 1.0) + hires_mask_dilated_in = scipy.ndimage.binary_dilation(hires_mask, iterations=int(1.5*round(lores_beam_width_pix))).astype(int) + hires_mask_border = (hires_mask_dilated_out + hires_mask_dilated_in) - 1 + comb_img[np.where(hires_mask_border)] = np.nan + comb_img[np.where(comb_img == 0)] = np.nan + + # Create an even larger border region for creating the interpolation to make the transition smooth + hires_mask_border_expanded = scipy.ndimage.binary_dilation(hires_mask_border, iterations=int(2.0*round(lores_beam_width_pix))).astype(int) + comb_fill = comb_img.copy() + comb_fill[np.where(hires_mask_border_expanded)] = np.nan + comb_fill = astropy.convolution.interpolate_replace_nans(comb_fill, + astropy.convolution.Gaussian2DKernel(round(2.0*lores_beam_width_pix)), + astropy.convolution.convolve_fft, + allow_huge=True, + boundary='wrap') + comb_img[np.where(hires_mask_border==1)] = comb_fill[np.where(hires_mask_border==1)] + #comb_img = astropy.convolution.interpolate_replace_nans(comb_img, astropy.convolution.Gaussian2DKernel(round(2.0*lores_beam_width_pix)), astropy.convolution.convolve_fft, allow_huge=True, boundary='wrap') + comb_img[np.where(lores_edge == 1.0)] = np.nan + + # Create mask describing tegions of the map where the only data is the good high-resolution combined data + comb_mask = hires_mask_border.copy() + comb_mask[np.where(np.isnan(hires_img_orig))] = 0 + + # Remove any temp files + if lores_hdu_path != False: + os.remove(lores_hdu_path) + if hires_hdu_path != False: + os.remove(hires_hdu_path) + + # Return combined image (or save to file if that was requested, with added possibility of returing calibration correction) + if not to_file: + if return_all: + return (comb_img, hires_fourier_corr_factor[0], hires_fourier_corr_factor[1], comb_mask) + else: + return comb_img + else: + astropy.io.fits.writeto(to_file, data=comb_img, header=hires_hdr, overwrite=True) + if return_all: + return (to_file, hires_fourier_corr_factor[0], hires_fourier_corr_factor[1], comb_mask) + else: + return to_file + + + +# Function to create a 2D tapering fourier-space filter, transitioning between two defined angular scales, according to a Hann (ie, cosine bell) filter +# Inputs: Iterable giving angular scale of high- and low resolution cutoffs (in deg); WCS of the data to be combined +# Outputs: Array containing requested filter, with low-frequency passpand and high-frequency stopband +def FourierTaper(taper_cutoffs_deg, in_wcs): + + # Use provided WCS to construct array to hold the output filter + if in_wcs.array_shape[0] != in_wcs.array_shape[1]: + raise Exception('Input data to be combined are not square; this will not end well') + out_filter = np.zeros([in_wcs.array_shape[1], in_wcs.array_shape[0]]) + out_i_centre = (0.5*out_filter.shape[0])-0.5 + out_j_centre = (0.5*out_filter.shape[1])-0.5 + + # Convert cutoff scales to units of fourier-pixels + pix_width_deg = np.abs( np.max( in_wcs.pixel_scale_matrix ) ) + cutoff_min_deg = min(taper_cutoffs_deg) + cutoff_max_deg = max(taper_cutoffs_deg) + cutoff_max_frac = (cutoff_max_deg / pix_width_deg) / out_filter.shape[0] + cutoff_min_frac = (cutoff_min_deg / pix_width_deg) / out_filter.shape[0] + cutoff_max_pix = 1.0 * (1.0 / cutoff_max_frac) + cutoff_min_pix = 1.0 * (1.0 / cutoff_min_frac) + + # Use meshgrids to find distance of each pixel from centre of filter array + i_linespace = np.linspace(0, out_filter.shape[0]-1, out_filter.shape[0]) + j_linespace = np.linspace(0, out_filter.shape[1]-1, out_filter.shape[1]) + i_grid, j_grid = np.meshgrid(i_linespace, j_linespace, indexing='ij') + i_grid -= out_i_centre + j_grid -= out_j_centre + rad_grid = np.sqrt(i_grid**2.0 + j_grid**2.0) + + # Rejigger the distance grid to give distance past inner edge of taper region + rad_grid -= cutoff_max_pix + + # Set central region (ie, low-resolution regime) to be entirely passband + out_filter[np.where(rad_grid <= 0)] = 1.0 + + # Construct well-sampled Hann filter to compute taper for transition region + hann_filter = np.hanning(2000)[1000:] + hann_pix = np.linspace(0, cutoff_min_pix-cutoff_max_pix, num=1000) + hann_interp = scipy.interpolate.interp1d(hann_pix, hann_filter, bounds_error=False, fill_value=np.nan) + + # Identify pixels where Hann filter is to be applied, and compute taper + hann_where = np.where((rad_grid > 0) & (rad_grid <= cutoff_min_pix)) + out_filter[hann_where] = hann_interp(rad_grid[hann_where]) + out_filter[np.where(np.isnan(out_filter))] = 0.0 + + # Return final filter + return out_filter + + + +# Function to cross-calibrate two data sets' power over a range of angular scales, as a precursor to fourier combination +# Input: Fourier transform of low-resolution data; fourier transform of high-resolution data; iterable giving angular scale of high- and low resolution cutoffs (in deg) +# Output: Correction factor to be applied to high-resolution data; uncertainty on the correction factor +def FourierCalibrate(lores_fourier, hires_fourier, taper_cutoffs_deg, hires_pix_width_deg): + + # Find centre of fourier arrays (ie, the zeroth order fourier frequency) + freq_i_centre = (0.5*hires_fourier.shape[0])-0.5 + freq_j_centre = (0.5*hires_fourier.shape[1])-0.5 + + # Calculate grid of radial fourier frequency distance from zeroth order + i_linespace = np.linspace(0, hires_fourier.shape[0]-1, hires_fourier.shape[0]) + j_linespace = np.linspace(0, hires_fourier.shape[1]-1, hires_fourier.shape[1]) + i_grid, j_grid = np.meshgrid(i_linespace, j_linespace, indexing='ij') + i_grid -= freq_i_centre + j_grid -= freq_j_centre + rad_grid = np.sqrt(i_grid**2.0 + j_grid**2.0) + + # Convert cutoff scales from degrees to fourier frequencies + cutoff_min_deg = min(taper_cutoffs_deg) + cutoff_max_deg = max(taper_cutoffs_deg) + cutoff_min_frac = (cutoff_min_deg / hires_pix_width_deg) / hires_fourier.shape[0] + cutoff_max_frac = (cutoff_max_deg / hires_pix_width_deg) / hires_fourier.shape[0] + cutoff_min_pix = 1.0 / cutoff_min_frac + cutoff_max_pix = 1.0 / cutoff_max_frac + + """# Slice out subsets of radial distance grid corresponding to cutoffs, and their overlap + rad_grid_cutoff_min = rad_grid[np.where(rad_gridcutoff_max_pix)] + rad_grid_overlap = rad_grid[np.where((rad_grid>cutoff_max_pix) & (rad_gridcutoff_max_pix) & (rad_gridcutoff_max_pix) & (rad_grid