def test_fft(data_dir, shape): a = create_random_data(shape, -100, 100, 'complex') start = time.time() ia = fft(a) stop = time.time() print('Original Time: {:.2f}s'.format(stop - start)) store_data(os.path.join(data_dir, 'a.dat'), a) store_data(os.path.join(data_dir, 'ia.dat'), ia)
def predict_2d_base(vis: Union[BlockVisibility, Visibility], model: Image, **kwargs) -> Union[BlockVisibility, Visibility]: """ Predict using convolutional degridding. This is at the bottom of the layering i.e. all transforms are eventually expressed in terms of this function. Any shifting needed is performed here. :param vis: Visibility to be predicted :param model: model image :return: resulting visibility (in place works) """ if isinstance(vis, BlockVisibility): log.debug("imaging.predict: coalescing prior to prediction") avis = coalesce_visibility(vis, **kwargs) else: avis = vis assert isinstance(avis, Visibility), avis _, _, ny, nx = model.data.shape padding = {} if get_parameter(kwargs, "padding", False): padding = {'padding': get_parameter(kwargs, "padding", False)} spectral_mode, vfrequencymap = get_frequency_map(avis, model) polarisation_mode, vpolarisationmap = get_polarisation_map(avis, model) uvw_mode, shape, padding, vuvwmap = get_uvw_map(avis, model, **padding) kernel_name, gcf, vkernellist = get_kernel_list(avis, model, **kwargs) uvgrid = fft((pad_mid(model.data, int(round(padding * nx))) * gcf).astype(dtype=complex)) avis.data['vis'] = convolutional_degrid(vkernellist, avis.data['vis'].shape, uvgrid, vuvwmap, vfrequencymap, vpolarisationmap) # Now we can shift the visibility from the image frame to the original visibility frame svis = shift_vis_to_image(avis, model, tangent=True, inverse=True) if isinstance(vis, BlockVisibility) and isinstance(svis, Visibility): log.debug("imaging.predict decoalescing post prediction") return decoalesce_visibility(svis) else: return svis
def fft_image(im, template_image=None): """ FFT an image, transform WCS as well Prefer to use axes 'UU---SIN' and 'VV---SIN' but astropy will not accept. :param im: :param template: :return: """ assert len(im.shape) == 4 d2r = numpy.pi / 180.0 ft_wcs = copy.deepcopy(im.wcs) ft_shape = im.shape if im.wcs.wcs.ctype[0] == 'RA---SIN' and im.wcs.wcs.ctype[1] == 'DEC--SIN': ft_wcs.wcs.axis_types[0] = 0 ft_wcs.wcs.axis_types[1] = 0 ft_wcs.wcs.crval[0] = 0.0 ft_wcs.wcs.crval[1] = 0.0 ft_wcs.wcs.crpix[0] = ft_shape[3] // 2 + 1 ft_wcs.wcs.crpix[1] = ft_shape[2] // 2 + 1 ft_wcs.wcs.ctype[0] = 'UU' ft_wcs.wcs.ctype[1] = 'VV' ft_wcs.wcs.cdelt[0] = 1.0 / (ft_shape[3] * d2r * im.wcs.wcs.cdelt[0]) ft_wcs.wcs.cdelt[1] = 1.0 / (ft_shape[2] * d2r * im.wcs.wcs.cdelt[1]) ft_data = ifft(im.data.astype('complex')) return create_image_from_array( ft_data, wcs=ft_wcs, polarisation_frame=im.polarisation_frame) elif im.wcs.wcs.ctype[0] == 'UU' and im.wcs.wcs.ctype[1] == 'VV': ft_wcs.wcs.crval[0] = template_image.wcs.wcs.crval[0] ft_wcs.wcs.crval[1] = template_image.wcs.wcs.crval[1] ft_wcs.wcs.crpix[0] = template_image.wcs.wcs.crpix[0] ft_wcs.wcs.crpix[0] = template_image.wcs.wcs.crpix[1] ft_wcs.wcs.ctype[0] = template_image.wcs.wcs.ctype[0] ft_wcs.wcs.ctype[1] = template_image.wcs.wcs.ctype[1] ft_wcs.wcs.cdelt[0] = template_image.wcs.wcs.cdelt[0] ft_wcs.wcs.cdelt[1] = template_image.wcs.wcs.cdelt[1] ft_data = fft(im.data.astype('complex')) return create_image_from_array( ft_data, wcs=ft_wcs, polarisation_frame=im.polarisation_frame) else: raise NotImplementedError("Cannot FFT specified axes")
def predict_2d_base(vis: Visibility, model: Image, **kwargs) -> Visibility: """ Predict using convolutional degridding. This is at the bottom of the layering i.e. all transforms are eventually expressed in terms of this function. Any shifting needed is performed here. :param vis: Visibility to be predicted :param model: model image :return: resulting visibility (in place works) """ if type(vis) is not Visibility: avis = coalesce_visibility(vis, **kwargs) else: avis = vis _, _, ny, nx = model.data.shape # print(model.shape) spectral_mode, vfrequencymap = get_frequency_map(avis, model) # 可以并行 polarisation_mode, vpolarisationmap = get_polarisation_map( avis, model, **kwargs) # 可以并行 uvw_mode, shape, padding, vuvwmap = get_uvw_map(avis, model, **kwargs) # 可以并行 kernel_name, gcf, vkernellist = get_kernel_list(avis, model, **kwargs) uvgrid = fft((pad_mid(model.data, int(round(padding * nx))) * gcf).astype(dtype=complex)) avis.data['vis'] = convolutional_degrid(vkernellist, avis.data['vis'].shape, uvgrid, vuvwmap, vfrequencymap, vpolarisationmap) # Now we can shift the visibility from the image frame to the original visibility frame svis = shift_vis_to_image(avis, model, tangent=True, inverse=True) if type(vis) is not Visibility: return decoalesce_visibility(svis) else: return svis
def predict_2d_base_timing(vis: Visibility, model: Image, **kwargs) -> (Visibility, tuple): """ Predict using convolutional degridding. This is at the bottom of the layering i.e. all transforms are eventually expressed in terms of this function. Any shifting needed is performed here. :param vis: Visibility to be predicted :param model: model image :return: resulting visibility (in place works) """ if not isinstance(vis, Visibility): avis = coalesce_visibility(vis, **kwargs) else: avis = vis _, _, ny, nx = model.data.shape opt = get_parameter(kwargs, 'opt', False) if not opt: log.debug('Using original algorithm') else: log.debug('Using optimized algorithm') padding = {} if get_parameter(kwargs, "padding", False): padding = {'padding': get_parameter(kwargs, "padding", False)} spectral_mode, vfrequencymap = get_frequency_map(avis, model, opt) polarisation_mode, vpolarisationmap = get_polarisation_map(avis, model) uvw_mode, shape, padding, vuvwmap = get_uvw_map(avis, model, **padding) kernel_name, gcf, vkernellist = get_kernel_list(avis, model, **kwargs) inarr = (pad_mid(model.data, int(round(padding * nx))) * gcf).astype(dtype=complex) # Use original algorithm if not opt: time_fft = -time.time() uvgrid = fft(inarr) time_fft += time.time() time_degrid = -time.time() vt = convolutional_degrid(vkernellist, avis.data['vis'].shape, uvgrid, vuvwmap, vfrequencymap, vpolarisationmap) time_degrid += time.time() # Use optimized algorithm else: time_fft = -time.time() uvgrid = numpy.zeros(inarr.shape, dtype=inarr.dtype) fft_c(uvgrid, inarr) time_fft += time.time() time_degrid = -time.time() kernel_indices, kernels = vkernellist ks0, ks1, ks2, ks3 = kernels[0].shape kernels_c = numpy.zeros((len(kernels), ks0, ks1, ks2, ks3), dtype=kernels[0].dtype) for i in range(len(kernels)): kernels_c[i, ...] = kernels[i] vfrequencymap_c = numpy.array(vfrequencymap, dtype=numpy.int32) vt = numpy.zeros(avis.data['vis'].shape, dtype=numpy.complex128) convolutional_degrid_c(vt, native_order(kernels_c), native_order(kernel_indices), native_order(uvgrid), native_order(vuvwmap), native_order(vfrequencymap_c)) time_degrid += time.time() avis.data['vis'] = vt # Now we can shift the visibility from the image frame to the original visibility frame svis = shift_vis_to_image(avis, model, tangent=True, inverse=True) if not isinstance(vis, Visibility): svis = decoalesce_visibility(svis) return svis, (time_degrid, time_fft)