def convolveWithBeam(self, ell, B_ell, threads=1): """ Return a liteMap object holding the data convolved with the beam specified as input in Fourier space Parameters ---------- ell : array_like The 1D ell values corresponding to the beam. B_ell : array_like, The 1D beam values in Fourier space. threads : int, optional Number of threads to use in pyFFTW calculations. Default is 1. """ beamFT = self.fillFourierTransform(ell, B_ell, elTrim=numpy.amax(ell), threads=threads) mapFT = fftTools.fftFromLiteMap(self, threads=threads) mapFT.kMap[:] *= beamFT.kMap[:] if have_pyFFTW: data_conv = numpy.real(ifft2(mapFT.kMap, threads=threads)) else: data_conv = numpy.real(ifft2(mapFT.kMap)) out = self.copy() out.data[:] = data_conv[:] return out
def _findCarrier(self, band0, band1, mask): band = band0 * band1 ixf = abs(fft.fftshift(fft.fft2(fft.fftshift(band)))) if self.debug: plt.figure() plt.title('Find carrier') plt.imshow((ixf - gaussian_filter(ixf, 20)) * mask) pyc0, pxc0 = self._findPeak((ixf - gaussian_filter(ixf, 20)) * mask) ixfz, Kx, Ky = self._zoomf(band, self.N, self._kx[pyc0, pxc0], self._ky[pyc0, pxc0], 50, self._dk * self.N) pyc, pxc = self._findPeak(abs(ixfz)) if self.debug: plt.figure() plt.title('Zoon Find carrier') plt.imshow(abs(ixfz)) kx = Kx[pxc] ky = Ky[pyc] otf_exclude_min_radius = 0.5 otf_exclude_max_radius = 1.5 kr = sqrt(self._kx**2 + self._ky**2) m = (kr < 2) otf = fft.fftshift(self._tfm(kr, m) + (1 - m)) otf_mask = (kr > otf_exclude_min_radius) & (kr < otf_exclude_max_radius) otf_mask_for_band_common_freq = fft.fftshift( otf_mask & scipy.ndimage.shift(otf_mask, (pyc0 - (self.N // 2 + 1), pxc0 - (self.N // 2 + 1)), order=0)) band0_common = fft.ifft2( fft.fft2(band0) / otf * otf_mask_for_band_common_freq) xx = np.arange(-self.N / 2 * self._dx, self.N / 2 * self._dx, self._dx, dtype=np.single) phase_shift_to_xpeak = exp(-1j * kx * xx * 2 * pi * self.NA / self.wavelength) phase_shift_to_ypeak = exp(-1j * ky * xx * 2 * pi * self.NA / self.wavelength) band1_common = fft.ifft2( fft.fft2(band1) / otf * otf_mask_for_band_common_freq) * np.outer( phase_shift_to_ypeak, phase_shift_to_xpeak) scaling = 1 / np.sum(band0_common * np.conjugate(band0_common)) cross_corr_result = np.sum(band0_common * band1_common) * scaling ampl = np.abs(cross_corr_result) * 2 phase = np.angle(cross_corr_result) return kx, ky, phase, ampl
def fillWithGaussianRandomField(self, ell, Cell, bufferFactor=1, threads=1): """ Generate a Gaussian random field from an input power spectrum specified as ell, Cell. Notes ----- BufferFactor = 1 means the map will have periodic boundary function, while BufferFactor > 1 means the map will be genrated on a patch bufferFactor times larger in each dimension and then cut out so as to have non-periodic boundary conditions. Fills the data field of the map with the GRF realization. """ ft = fftTools.fftFromLiteMap(self, threads=threads) Ny = self.Ny * bufferFactor Nx = self.Nx * bufferFactor bufferFactor = int(bufferFactor) realPart = numpy.zeros([Ny, Nx]) imgPart = numpy.zeros([Ny, Nx]) ly = fftfreq(Ny, d=self.pixScaleY) * (2 * numpy.pi) lx = fftfreq(Nx, d=self.pixScaleX) * (2 * numpy.pi) # print ly modLMap = numpy.zeros([Ny, Nx]) iy, ix = numpy.mgrid[0:Ny, 0:Nx] modLMap[iy, ix] = numpy.sqrt(ly[iy] ** 2 + lx[ix] ** 2) s = splrep(ell, Cell, k=3) ll = numpy.ravel(modLMap) kk = splev(ll, s) id = numpy.where(ll > ell.max()) kk[id] = 0.0 # add a cosine ^2 falloff at the very end # id2 = numpy.where( (ll> (ell.max()-500)) & (ll<ell.max())) # lEnd = ll[id2] # kk[id2] *= numpy.cos((lEnd-lEnd.min())/(lEnd.max() -lEnd.min())*numpy.pi/2) # pylab.loglog(ll,kk) area = Nx * Ny * self.pixScaleX * self.pixScaleY p = numpy.reshape(kk, [Ny, Nx]) / area * (Nx * Ny) ** 2 realPart = numpy.sqrt(p) * numpy.random.randn(Ny, Nx) imgPart = numpy.sqrt(p) * numpy.random.randn(Ny, Nx) kMap = realPart + 1j * imgPart if have_pyFFTW: data = numpy.real(ifft2(kMap, threads=threads)) else: data = numpy.real(ifft2(kMap)) b = bufferFactor self.data = data[(b - 1) / 2 * self.Ny : (b + 1) / 2 * self.Ny, (b - 1) / 2 * self.Nx : (b + 1) / 2 * self.Nx]
def _refineCarrier(self, band0, band1, kx_in, ky_in): pxc0 = np.int(np.round(kx_in / self._dk) + self.N / 2) pyc0 = np.int(np.round(ky_in / self._dk) + self.N / 2) otf_exclude_min_radius = self.eta / 2 otf_exclude_max_radius = 1.5 m = (self._kr < 2) otf = fft.fftshift(self._tfm(self._kr, m) + (1 - m) * 0.0001) otf_mask = (self._kr > otf_exclude_min_radius) & ( self._kr < otf_exclude_max_radius) otf_mask_for_band_common_freq = fft.fftshift( otf_mask & scipy.ndimage.shift(otf_mask, (pyc0 - (self.N // 2), pxc0 - (self.N // 2)), order=0)) band0_common = fft.ifft2( fft.fft2(band0) / otf * otf_mask_for_band_common_freq) band1_common = fft.ifft2( fft.fft2(band1) / otf * otf_mask_for_band_common_freq) band = band0_common * band1_common mag = 25 * self.N / 256 ixfz, Kx, Ky = self._zoomf(band, self.N, np.single(self._k[pxc0]), np.single(self._k[pyc0]), mag, self._dk * self.N) pyc, pxc = self._findPeak(abs(ixfz)) if self.debug: plt.figure() plt.title('Zoom Find carrier') plt.imshow(abs(ixfz)) kx = Kx[pxc] ky = Ky[pyc] xx = np.arange(-self.N / 2 * self._dx, self.N / 2 * self._dx, self._dx, dtype=np.double) phase_shift_to_xpeak = exp(-1j * kx * xx * 2 * pi * self.NA / self.wavelength) phase_shift_to_ypeak = exp(-1j * ky * xx * 2 * pi * self.NA / self.wavelength) scaling = 1 / np.sum(band0_common * np.conjugate(band0_common)) cross_corr_result = np.sum(band0_common * band1_common * np.outer( phase_shift_to_ypeak, phase_shift_to_xpeak)) * scaling ampl = np.abs(cross_corr_result) * 2 phase = np.angle(cross_corr_result) return kx, ky, phase, ampl
def reconstruct_fftw(self, img): imf = fft.fft2(img) * self._prefilter self._carray[:, 0:self.N // 2, 0:self.N // 2] = imf[:, 0:self.N // 2, 0:self.N // 2] self._carray[:, 0:self.N // 2, 3 * self.N // 2:2 * self.N] = imf[:, 0:self.N // 2, self.N // 2:self.N] self._carray[:, 3 * self.N // 2:2 * self.N, 0:self.N // 2] = imf[:, self.N // 2:self.N, 0:self.N // 2] self._carray[:, 3 * self.N // 2:2 * self.N, 3 * self.N // 2:2 * self.N] = imf[:, self.N // 2:self.N, self.N // 2:self.N] img2 = np.sum(np.real(fft.ifft2(self._carray)).real * self._reconfactor, 0) self._imgstore = img self._bigimgstore = fft.ifft2(fft.fft2(img2) * self._postfilter).real return self._bigimgstore
def solve_spectral(Winit, expU, expT): B = fft2(Winit, axes=(0,1)) # (x,y,px,py) -> (λx,λy,px,py) B *= expT B = ifft2(B, axes=(0,1)) # (λx,λy,px,py) -> (x,y,px,py) B = fft2(B, axes=(2,3)) # (x,y,px,py) -> (x,y,θx,θy) B *= expU B = ifft2(B, axes=(2,3)) # (x,y,θx,θy) -> (x,y,px,py) B = fft2(B, axes=(0,1)) # (x,y,px,py) -> (λx,λy,px,py) B *= expT B = ifft2(B, axes=(0,1)) # (λx,λy,px,py) -> (x,y,px,py) return real(B) # to avoid python warning
def fast_icfft2(x, axes=(-1, -2)): if len(x.shape) == 2: return fftshift(np.transpose(ifft2(np.transpose(ifftshift(x))))) elif len(x.shape) == 3: y = ifftshift(x, axes=axes) y = ifft2(y, axes=axes) y = fftshift(y, axes=axes) return y else: raise ValueError("x must be 2D or 3D")
def reconstructframe_fftw(self, img, i): diff = img - self._imgstore[i, :, :] imf = fft.fft2(diff) * self._prefilter self._carray[0, 0:self.N // 2, 0:self.N // 2] = imf[0:self.N // 2, 0:self.N // 2] self._carray[0, 0:self.N // 2, 3 * self.N // 2:2 * self.N] = imf[0:self.N // 2, self.N // 2:self.N] self._carray[0, 3 * self.N // 2:2 * self.N, 0:self.N // 2] = imf[self.N // 2:self.N, 0:self.N // 2] self._carray[0, 3 * self.N // 2:2 * self.N, 3 * self.N // 2:2 * self.N] = imf[self.N // 2:self.N, self.N // 2:self.N] img2 = fft.ifft2(self._carray[0, :, :]).real * self._reconfactor[i, :, :] self._imgstore[i, :, :] = img self._bigimgstore = self._bigimgstore + fft.ifft2(fft.fft2(img2) * self._postfilter).real return self._bigimgstore
def optimal_binary_image_subtraction(R,N,Pr,Pn,sr,sn): R_hat = fft.fft2(R) N_hat = fft.fft2(N) Pn_hat = fft.fft2(Pn) Pr_hat = fft.fft2(Pr) G_hat = (Pr_hat*N_hat - Pn_hat*R_hat) / np.sqrt((sr**2*abs(Pn_hat**2) + sn**2*abs(Pr_hat**2))) P_G_hat = (Pr_hat*Pn_hat) / np.sqrt((sr**2*abs(Pn_hat**2) + sn**2*abs(Pr_hat**2))) S_hat = G_hat*conj(P_G_hat) #S_hat = (conj(Pn_hat)*np.abs(Pr_hat)**2*N_hat - conj(Pr_hat)*np.abs(Pn_hat)**2*R_hat) / (sr**2*abs(Pn_hat**2) + sn**2*abs(Pr_hat**2)) S = fft.ifft2(S_hat) G = fft.ifft2(G_hat) P_G = real(fft.ifft2(P_G_hat)) return S/std(S[15::30,15::30]), G/std(G[15::30,15::30]), P_G / sum(P_G)
def multislice_propagate_batch_numpy(grid_delta_batch, grid_beta_batch, probe_real, probe_imag, energy_ev, psize_cm, free_prop_cm=None, obj_batch_shape=None): minibatch_size = obj_batch_shape[0] grid_shape = obj_batch_shape[1:] voxel_nm = np.array([psize_cm] * 3) * 1.e7 wavefront = np.zeros([minibatch_size, obj_batch_shape[1], obj_batch_shape[2]], dtype='complex64') wavefront += (probe_real + 1j * probe_imag) lmbda_nm = 1240. / energy_ev mean_voxel_nm = np.prod(voxel_nm) ** (1. / 3) size_nm = np.array(grid_shape) * voxel_nm n_slice = obj_batch_shape[-1] delta_nm = voxel_nm[-1] # h = get_kernel_ir(delta_nm, lmbda_nm, voxel_nm, grid_shape) h = get_kernel(delta_nm, lmbda_nm, voxel_nm, grid_shape) k = 2. * PI * delta_nm / lmbda_nm probe_array = [] for i in range(n_slice): delta_slice = grid_delta_batch[:, :, :, i] beta_slice = grid_beta_batch[:, :, :, i] c = np.exp(1j * k * delta_slice) * np.exp(-k * beta_slice) wavefront = wavefront * c if i < n_slice - 1: wavefront = ifft2(np_ifftshift(np_fftshift(fft2(wavefront), axes=[1, 2]) * h, axes=[1, 2])) probe_array.append(wavefront) if free_prop_cm is not None: #1dxchange.write_tiff(abs(wavefront), '2d_1024/monitor_output/wv', dtype='float32', overwrite=True) if free_prop_cm == 'inf': wavefront = np_fftshift(fft2(wavefront), axes=[1, 2]) else: dist_nm = free_prop_cm * 1e7 l = np.prod(size_nm)**(1. / 3) crit_samp = lmbda_nm * dist_nm / l algorithm = 'TF' if mean_voxel_nm > crit_samp else 'IR' # print(algorithm) algorithm = 'TF' if algorithm == 'TF': h = get_kernel(dist_nm, lmbda_nm, voxel_nm, grid_shape) wavefront = ifft2(np_ifftshift(np_fftshift(fft2(wavefront), axes=[1, 2]) * h, axes=[1, 2])) else: h = get_kernel_ir(dist_nm, lmbda_nm, voxel_nm, grid_shape) wavefront = ifft2(np_ifftshift(np_fftshift(fft2(wavefront), axes=[1, 2]) * h, axes=[1, 2])) # dxchange.write_tiff(abs(wavefront), '2d_512/monitor_output/wv', dtype='float32', overwrite=True) # dxchange.write_tiff(np.angle(h), '2d_512/monitor_output/h', dtype='float32', overwrite=True) return wavefront, np.array(probe_array)
def fast_icfft2(x): if len(x.shape) == 2: return fftshift(np.transpose(ifft2(np.transpose(ifftshift(x))))) elif len(x.shape) == 3: y = ifftshift(x, (1, 2)) y = np.transpose(y, (0, 2, 1)) y = ifft2(y) y = np.transpose(y, (0, 2, 1)) y = fftshift(y, (1, 2)) return y else: raise ValueError("x must be 2D or 3D")
def decomposition_fft(X, filter, **kwargs): """Decompose a 2d input field into multiple spatial scales by using the Fast Fourier Transform (FFT) and a bandpass filter. Parameters ---------- X : array_like Two-dimensional array containing the input field. All values are required to be finite. filter : dict A filter returned by any method implemented in bandpass_filters.py. Other Parameters ---------------- MASK : array_like Optional mask to use for computing the statistics for the cascade levels. Pixels with MASK==False are excluded from the computations. Returns ------- out : ndarray A dictionary described in the module documentation. The parameter n is determined from the filter (see bandpass_filters.py). """ MASK = kwargs.get("MASK", None) if len(X.shape) != 2: raise ValueError("the input is not two-dimensional array") if MASK is not None and MASK.shape != X.shape: raise ValueError("dimension mismatch between X and MASK: X.shape=%s, MASK.shape=%s" % \ (str(X.shape), str(MASK.shape))) if X.shape != filter["weights_2d"].shape[1:3]: raise ValueError( "dimension mismatch between X and filter: X.shape=%s, filter['weights_2d'].shape[1:3]=%s" % (str(X.shape), str(filter["weights_2d"].shape[1:3]))) if np.any(~np.isfinite(X)): raise ValueError("X contains non-finite values") result = {} means = [] stds = [] F = fft.fftshift(fft.fft2(X, **fft_kwargs)) X_decomp = [] for k in range(len(filter["weights_1d"])): W_k = filter["weights_2d"][k, :, :] X_ = np.real(fft.ifft2(fft.ifftshift(F * W_k), **fft_kwargs)) X_decomp.append(X_) if MASK is not None: X_ = X_[MASK] means.append(np.mean(X_)) stds.append(np.std(X_)) result["cascade_levels"] = np.stack(X_decomp) result["means"] = means result["stds"] = stds return result
def downsample(stack, n, mask=None, stack_in_fourier=False): """ Use Fourier methods to change the sample interval and/or aspect ratio of any dimensions of the input image 'img'. If the optional argument stack is set to True, then the *first* dimension of 'img' is interpreted as the index of each image in the stack. The size argument side is an integer, the size of the output images. Let the size of a stack of 2D images 'img' be n1 x n1 x k. The size of the output will be side x side x k. If the optional mask argument is given, this is used as the zero-centered Fourier mask for the re-sampling. The size of mask should be the same as the output image size. For example for downsampling an n0 x n0 image with a 0.9 x nyquist filter, do the following: msk = fuzzymask(n,2,.45*n,.05*n) out = downsample(img, n, 0, msk) The size of the mask must be the size of output. The optional fx output argument is the padded or cropped, masked, FT of in, with zero frequency at the origin. """ size_in = np.square(stack.shape[1]) size_out = np.square(n) mask = 1 if mask is None else mask num_images = stack.shape[0] output = np.zeros((num_images, n, n), dtype='float32') images_batches = np.array_split(np.arange(num_images), 10) for batch in images_batches: curr_batch = np.array(stack[batch]) curr_batch = curr_batch if stack_in_fourier else fft2(curr_batch) fx = common.crop(np.fft.fftshift(curr_batch, axes=(-2, -1)), (-1, n, n)) * mask output[batch] = ifft2(np.fft.ifftshift( fx, axes=(-2, -1))) * (size_out / size_in) print('finished {}/{}'.format(batch[-1] + 1, num_images)) return output
def propcoad(ref_dir, make_fits=False): """ Proper coaddition function """ if len(ref_dir)==1: F = glob.glob(ref_dir[0]+'/*.fits') # collect images you want to coadd print(F) else: F = ref_dir print(F) psf_dat, psf_hed, sexcat1, psfcat1 = get_psf(F[0]) pool = multiprocessing.Pool(len(F)-1) tmp_array = pool.starmap(coad_func, [(F[i+1] , psfcat1, sexcat1) for i in range(len(F)-1)]) Nomin = sum(x[0] for x in tmp_array) Denom = sum(x[1] for x in tmp_array) Denom = np.sqrt(Denom) if np.any(Denom==0): print('ZEROS') R_hat = Nomin/Denom R = np.real(fft.ifft2(R_hat)) subprocess.call(['rm', sexcat1, psfcat1, sexcat1.replace('.fits', '.psf')]) if make_fits == True: hed = fits.getheader(F[0]) hed['COMMENT'] = 'ZO coaddition from ZiP' hed['COMMENT'] = 'List of stacked fits' hed['COMMENT'] = ', '.join(F) fits.writeto('ZOcoadd'+F[0], R, hed, overwrite = True) return('ZOcoadd'+F[0], R) else: return(R)
def phrt(im, plan, method=4, nr_threads=2): """Process a tomographic projection image with the selected phase retrieval algorithm. Parameters ---------- im : array_like Flat corrected image data as numpy array. plan : structure Structure with pre-computed data (see prepare_plan function) method : int Phase retrieval filter {1 = TIE (default), 2 = CTF, 3 = CTF first-half sine, 4 = Quasiparticle, 5 = Quasiparticle first half sine}. nr_threads : int Number of threads to be used in the computation of FFT by PyFFTW Credits ------- Julian Moosmann, KIT (Germany) is acknowledged for this code """ # Extract plan values: dim0_o = plan['dim0'] dim1_o = plan['dim1'] n_pad0 = plan['npad0'] n_pad1 = plan['npad1'] marg0 = (n_pad0 - dim0_o) / 2 marg1 = (n_pad1 - dim1_o) / 2 filter = plan['filter'] # Pad image (if required): im = padImage(im, n_pad0, n_pad1) # (Un)comment the following two lines to use PyFFTW: n_byte_align(im, simd_alignment) im = fft2(im - 1, threads=nr_threads) # (Un)comment the following line to use NumPy: #im = fft2(im - 1) # Apply phase retrieval filter: im = filter * im # (Un)comment the following two lines to use PyFFTW: n_byte_align(im, simd_alignment) im = real(ifft2(im, threads=nr_threads)) # (Un)comment the following line to use NumPy: #im = real(ifft2(im)) # Return the negative: im = - im # Return cropped output: return im[marg0:dim0_o + marg0, marg1:dim1_o + marg1]
def phrt(im, plan, method=4, nr_threads=2): """Process a tomographic projection image with the selected phase retrieval algorithm. Parameters ---------- im : array_like Flat corrected image data as numpy array. plan : structure Structure with pre-computed data (see prepare_plan function) method : int Phase retrieval filter {1 = TIE (default), 2 = CTF, 3 = CTF first-half sine, 4 = Quasiparticle, 5 = Quasiparticle first half sine}. nr_threads : int Number of threads to be used in the computation of FFT by PyFFTW Credits ------- Julian Moosmann, KIT (Germany) is acknowledged for this code """ # Extract plan values: dim0_o = plan['dim0'] dim1_o = plan['dim1'] n_pad0 = plan['npad0'] n_pad1 = plan['npad1'] marg0 = (n_pad0 - dim0_o) / 2 marg1 = (n_pad1 - dim1_o) / 2 filter = plan['filter'] # Pad image (if required): im = padImage(im, n_pad0, n_pad1) # (Un)comment the following two lines to use PyFFTW: n_byte_align(im, simd_alignment) im = fft2(im - 1, threads=nr_threads) # (Un)comment the following line to use NumPy: #im = fft2(im - 1) # Apply phase retrieval filter: im = filter * im # (Un)comment the following two lines to use PyFFTW: n_byte_align(im, simd_alignment) im = real(ifft2(im, threads=nr_threads)) # (Un)comment the following line to use NumPy: #im = real(ifft2(im)) # Return the negative: im = -im # Return cropped output: return im[marg0:dim0_o + marg0, marg1:dim1_o + marg1]
def generate_signals( shape: tuple, specs: list[CompanionSpec], template: np.ndarray, derotation_angles: np.ndarray = None, template_scale_factors: Union[np.ndarray, float, None] = None, ): '''Inject signals for companions specified using optional derotation angles to *counter*rotate the coordinates before injection, such that rotation CCW (e.g. by `derotate_cube`) aligns 0 deg PA with +Y Parameters ---------- shape : tuple[int,int,int] specs : list[CompanionSpec] template : np.ndarray derotation_angles : Optional[np.ndarray] template_scale_factors : Union[np.ndarray,float,None] Scale factor relative to 1.0 being the average brightness of the primary over the observation, used to scale the template image to reflect particularly sharp or poor AO correction Returns ------- outcube : np.ndarray ''' outcube = np.zeros(shape, dtype=template.dtype) n_obs = shape[0] template = improc.shift2(template, 0, 0, output_shape=shape[1:]) ft_template = fft.fft2(template) xfreqs = fft.fftfreq(shape[2]) yfreqs = fft.fftfreq(shape[1]) if template_scale_factors is None: template_scale_factors = np.ones(n_obs) if np.isscalar(template_scale_factors): template_scale_factors = np.repeat(np.array([template_scale_factors]), n_obs) if derotation_angles is None: derotation_angles = np.zeros(n_obs) for spec in specs: theta = np.deg2rad(90 + spec.pa_deg - derotation_angles) for i in range(n_obs): dx = spec.r_px * np.cos(theta[i]) dy = spec.r_px * np.sin(theta[i]) shifter = np.exp(2j * np.pi * ((-dx * xfreqs[np.newaxis, :]) + (-dy * yfreqs[:, np.newaxis]))) cube_contribution = fft.ifft2(ft_template * shifter).real cube_contribution *= template_scale_factors[i] * spec.scale outcube[i] += cube_contribution return outcube
def translation(im0, im1): """Return translation vector to register images.""" shape = im0.shape f0 = fft2(im0) f1 = fft2(im1) ir = abs(ifft2((f0 * f1.conjugate()) / (abs(f0) * abs(f1)))) t0, t1 = np.unravel_index(np.argmax(ir), shape) if t0 > shape[0] // 2: t0 -= shape[0] if t1 > shape[1] // 2: t1 -= shape[1] return [t0, t1]
def im_cov(im): # image isotrope covariance function based on fft cov = np.real(fft.fftshift(fft.ifft2(np.absolute(fft.fft2(im))** 2))) / im.size print(cov.shape) cov1 = cov[int(np.shape(cov)[0] / 2 + 0.5), int(np.shape(cov)[1] / 2):] cov2 = np.flip(cov[int(np.shape(cov)[0] / 2 + 0.5), :int(np.shape(cov)[1] / 2 + 0.5)]) cov3 = cov[int(np.shape(cov)[1] / 2):, int(np.shape(cov)[0] / 2 + 0.5)] cov4 = np.flip(cov[:int(np.shape(cov)[1] / 2 + 0.5), int(np.shape(cov)[0] / 2 + 0.5)]) cov = (cov1 + cov2 + cov3 + cov4) / 4 lags = np.arange(np.shape(cov)[0]) + 0.5 return lags, cov # x and y of the cov function
def fresnel_propagate_numpy(wavefront, energy_ev, psize_cm, dist_cm): lmbda_nm = 1240. / energy_ev lmbda_cm = 0.000124 / energy_ev psize_nm = psize_cm * 1e7 dist_nm = dist_cm * 1e7 if dist_cm == 'inf': wavefront = np_fftshift(fft2(wavefront)) else: n = np.mean(wavefront.shape) z_crit_cm = (psize_cm * n)**2 / (lmbda_cm * n) algorithm = 'TF' if dist_cm < z_crit_cm else 'IR' if algorithm == 'TF': h = get_kernel(dist_nm, lmbda_nm, [psize_nm, psize_nm], wavefront.shape) wavefront = ifft2(np_ifftshift(np_fftshift(fft2(wavefront)) * h)) else: h = get_kernel_ir(dist_nm, lmbda_nm, [psize_nm, psize_nm], wavefront.shape) wavefront = np_ifftshift(ifft2(np_fftshift(fft2(wavefront)) * h)) return wavefront
def worker(): # generate Gaussian white noise field, multiply it with the standard # deviation of the observed field and apply the precipitation mask N = randstates[k].randn(R.shape[0], R.shape[1]) N = np.real(fft.ifft2(fft.fft2(N) * R_fft)) N = N / np.std(N) * sigma + mu N[~MASK] = R_thr_2 # subtract the mean and decompose the masked noise field into a # cascade N -= mu decomp_N = decomp_method(N, F, MASK=MASK_) return decomp_N["stds"]
def _phase_correlation(im0, im1, callback=None, *args): """ Computes phase correlation between im0 and im1 Args: im0 im1 callback (function): Process the cross-power spectrum (i.e. choose coordinates of the best element, usually of the highest one). Defaults to :func:`imreg_dft.utils.argmax2D` Returns: tuple: The translation vector (Y, X). Translation vector of (0, 0) means that the two images match. """ if callback is None: callback = utils._argmax2D # TODO: Implement some form of high-pass filtering of PHASE correlation f0, f1 = [ fft.fft2(arr, threads=4, overwrite_input=True, auto_align_input=True, auto_contiguous=True, planner_effort='FFTW_ESTIMATE') for arr in (im0, im1) ] # spectrum can be filtered (already), # so we have to take precaution against dividing by 0 eps = abs(f1).max() * 1e-15 # cps == cross-power spectrum of im0 and im1 cps = abs( fft.ifft2((f0 * f1.conjugate()) / (abs(f0) * abs(f1) + eps), threads=4, overwrite_input=True, auto_align_input=True, auto_contiguous=True, planner_effort='FFTW_ESTIMATE')) # scps = shifted cps scps = fft.fftshift(cps) (t0, t1), success = callback(scps, *args) ret = np.array((t0, t1)) # _compensate_fftshift is not appropriate here, this is OK. t0 -= f0.shape[0] // 2 t1 -= f0.shape[1] // 2 ret -= np.array(f0.shape, int) // 2 return ret, success
def ft_shift2(image: np.ndarray, dy: float, dx: float, flux_tol: Union[None, float] = 1e-15, output_shape=None): """ Fast Fourier subpixel shifting Parameters ---------- dy : float Translation in +Y direction (i.e. a feature at (x, y) moves to (x, y + dy)) dx : float Translation in +X direction (i.e. a feature at (x, y) moves to (x + dx, y)) flux_tol : float Fractional flux change permissible ``(sum(output) - sum(image)) / sum(image) < flux_tol`` (default: 1e-15) output_shape : tuple shape of output array (default: same as input) """ if output_shape is None: output_shape = image.shape xfreqs = fft.fftfreq(output_shape[1]) yfreqs = fft.fftfreq(output_shape[0]) xform = fft.fft2(image, s=output_shape) if output_shape is not None: # compute center-to-center displacement such that # supplying dx == dy == 0.0 will be a no-op (aside # from changing shape) orig_ctr_x, orig_ctr_y = (image.shape[1] - 1) / 2, (image.shape[0] - 1) / 2 new_ctr_x, new_ctr_y = (output_shape[1] - 1) / 2, (output_shape[0] - 1) / 2 base_dx, base_dy = new_ctr_x - orig_ctr_x, new_ctr_y - orig_ctr_y else: base_dx = base_dy = 0 modified_xform = xform * np.exp(2j * np.pi * ( (-(dx + base_dx) * xfreqs)[np.newaxis, :] + (-(dy + base_dy) * yfreqs)[:, np.newaxis])) new_image = fft.ifft2(modified_xform).real frac_diff_flux = (np.sum(image) - np.sum(new_image)) / np.sum(image) if flux_tol is not None and frac_diff_flux > flux_tol: raise RuntimeError( f"Flux conservation violated by {frac_diff_flux} fractional difference (more than {flux_tol})" ) return new_image
def velocity_abs_exp(alpha, timestep, spacing, wavenumber): """ Absorption coefficient exponent. :param alpha: :math:`\\alpha_{\\xi}` :param timestep: Timestep :math:`\\Delta t` This value is calculated according to .. math:: e^{-\\alpha_{\\xi} \\Delta t / 2} However, since the velocity field is shifted by half a spacing, a correction needs to be applied. .. math:: \\mathcal{F}^{-1} \\left[ e^{+j k_{\\xi} \\Delta \\xi / 2} \\mathcal{F} \\left( e^{-\\alpha_{\\xi} \Delta t / 2} \\right) \\right] """ j = 1j return ifft2(ne.evaluate("exp(+j * wavenumber*spacing/2.0)") * fft2(ne.evaluate("exp(-alpha * timestep / 2.0)")))
def _phase_correlation(im0, im1, callback=None, *args): """ Computes phase correlation between im0 and im1 Args: im0 im1 callback (function): Process the cross-power spectrum (i.e. choose coordinates of the best element, usually of the highest one). Defaults to :func:`imreg_dft.utils.argmax2D` Returns: tuple: The translation vector (Y, X). Translation vector of (0, 0) means that the two images match. """ if callback is None: callback = utils._argmax2D # TODO: Implement some form of high-pass filtering of PHASE correlation f0, f1 = [fft.fft2(arr) for arr in (im0, im1)] print(im0) # spectrum can be filtered (already), # so we have to take precaution against dividing by 0 eps = abs(f1).max() * 1e-15 # cps == cross-power spectrum of im0 and im1 cps = abs(fft.ifft2((f0 * f1.conjugate()) / (abs(f0) * abs(f1) + eps))) #scps = shifted cps scps = fft.fftshift(cps) # imshow(scps*1000,scps*1000,scps*1000) imshow(scps, scps, scps) #scps = cps (t0, t1), success = callback(scps, *args) ret = np.array((t0, t1)) # _compensate_fftshift is not appropriate here, this is OK. # t0 -= f0.shape[0] // 2 # t1 -= f0.shape[1] // 2 print("ret_oir", ret) ret = ret - np.array(f0.shape, int) // 2 print("ret", ret) return ret, success
def fresnel_propagate(wavefield, energy_ev, psize_cm, dist_cm, fresnel_approx=True, pad=0, sign_convention=1): """ Perform Fresnel propagation on a batch of wavefields. :param wavefield: complex wavefield with shape [n_batches, n_y, n_x]. :param energy_ev: float. :param psize_cm: size-3 vector with pixel size ([dy, dx, dz]). :param dist_cm: propagation distance. :return: """ minibatch_size = wavefield.shape[0] if pad > 0: wavefield = np.pad(wavefield, [[0, 0], [pad, pad], [pad, pad]], mode='edge') grid_shape = wavefield.shape[1:] if len(psize_cm) == 1: psize_cm = [psize_cm] * 3 voxel_nm = np.array(psize_cm) * 1.e7 lmbda_nm = 1240. / energy_ev mean_voxel_nm = np.prod(voxel_nm)**(1. / 3) size_nm = np.array(grid_shape) * voxel_nm dist_nm = dist_cm * 1e7 h = get_kernel(dist_nm, lmbda_nm, voxel_nm, grid_shape, fresnel_approx=fresnel_approx, sign_convention=sign_convention) wavefield = ifft2( ifftshift(fftshift(fft2(wavefield), axes=[1, 2]) * h, axes=[1, 2])) if pad > 0: wavefield = wavefield[:, pad:-pad, pad:-pad] return wavefield
def multidistance_ctf(prj_ls, dist_cm_ls, psize_cm, energy_kev, kappa=50, sigma_cut=0.01, alpha_1=5e-4, alpha_2=1e-16): prj_ls = np.array(prj_ls) dist_cm_ls = np.array(dist_cm_ls) dist_nm_ls = dist_cm_ls * 1.e7 lmbda_nm = 1.24 / energy_kev psize_nm = psize_cm * 1.e7 prj_shape = prj_ls.shape[1:] u_max = 1. / (2. * psize_nm) v_max = 1. / (2. * psize_nm) u, v = gen_mesh([v_max, u_max], prj_shape) xi_mesh = PI * lmbda_nm * (u**2 + v**2) xi_ls = np.zeros([len(dist_cm_ls), *prj_shape]) for i in range(len(dist_cm_ls)): xi_ls[i] = xi_mesh * dist_nm_ls[i] abs_nu = np.sqrt(u**2 + v**2) nu_cut = 0.6 * u_max f = 0.5 * (1 - erf((abs_nu - nu_cut) / sigma_cut)) alpha = alpha_1 * f + alpha_2 * (1 - f) # plt.imshow(abs(np.log(np_fftshift(fft2(prj_ls[0] - 1, axes=(-2, -1)), axes=(-2, -1))))) # plt.imshow(alpha) # plt.show() # alpha = 0 phase = np.sum( np_fftshift(fft2(prj_ls - 1, axes=(-2, -1)), axes=(-2, -1)) * (np.sin(xi_ls) + 1. / kappa * np.cos(xi_ls)), axis=0) phase /= ( np.sum(2 * (np.sin(xi_ls) + 1. / kappa * np.cos(xi_ls))**2, axis=0) + alpha) phase = ifft2(np_ifftshift(phase, axes=(-2, -1)), axes=(-2, -1)) return np.abs(phase)
def free_propagate_spherical_numpy(wavefront, dist_cm, r_cm, wavelen_nm, probe_size, theta_max=PI / 18, phi_max=PI / 18): dist_nm = dist_cm * 1.e7 r_nm = r_cm * 1.e7 k_theta = PI / theta_max * (np.arange(probe_size[0]) - float(probe_size[0] - 1) / 2) k_phi = PI / phi_max * (np.arange(probe_size[1]) - float(probe_size[1] - 1) / 2) k_phi, k_theta = np.meshgrid(k_phi, k_theta) k = 2 * PI / wavelen_nm wavefront = np_fftshift(fft2(wavefront), axes=[-1, -2]) wavefront *= np.exp(-1j / (2 * k) * (k_theta**2 + k_phi**2) * (1. / (r_nm + dist_nm) - 1. / r_nm)) wavefront = ifft2(np_ifftshift(wavefront, axes=[-1, -2])) return wavefront
def _coarseFindCarrier(self, band0, band1, mask): otf_exclude_min_radius = self.eta / 2 # Min Radius of the circular region around DC that is to be excluded from the cross-correlation calculation maskhpf = fft.fftshift(self._kr > otf_exclude_min_radius) band0_common = fft.ifft2(fft.fft2(band0) * maskhpf) # band1_common = fft.ifft2(fft.fft2(band1)*maskhpf) ix = band0_common * band1 ixf = np.abs(fft.fftshift(fft.fft2(fft.fftshift(ix)))) if self.debug: plt.figure() plt.title('Find carrier') plt.imshow(ixf, cmap=plt.get_cmap('gray')) # pyc0, pxc0 = self._findPeak((ixf - gaussian_filter(ixf, 20)) * mask) pyc0, pxc0 = self._findPeak(ixf * mask) kx = self._dk * (pxc0 - self.N / 2) ky = self._dk * (pyc0 - self.N / 2) return kx, ky
def fft_filter(im, rois, shift, plot=True, imshow_kwargs={}): f0 = fft.fftshift(fft.fft2(im)) f1 = f0.copy() if plot: kw = dict(vmin=-0.05, vmax=1) kw.update(imshow_kwargs) fig, axes = plt.subplots(2, 2, figsize=(9, 4), sharex='col', sharey='col') (ax, ax_f), (ax1, ax1_f) = axes ax.imshow(im, **kw) ax_f.imshow(abs(f0), norm=LogNorm()) dx, dy = shift for j, roi in enumerate(rois): print(j, roi) roi_t = translate_roi(roi, dx, dy) roi_symm = symmetric_roi(roi, im.shape) roi_symm_t = symmetric_roi(roi_t, im.shape) f1[roi] = f1[roi_t] f1[roi_symm] = f1[roi_symm_t] if plot: roi_to_rect(roi, ax=ax_f, index=j) roi_to_rect(roi, ax=ax1_f) roi_to_rect(roi_t, ax=ax_f, ec='r') roi_to_rect(roi_symm, ax=ax_f) roi_to_rect(roi_symm, ax=ax1_f) roi_to_rect(roi_symm_t, ax=ax_f, ec='r') im1 = fft.ifft2(fft.ifftshift(f1)).real if plot: ax1.imshow(im1, **kw) ax1_f.imshow(abs(f1), norm=LogNorm()) return im1
def generate_noise_2d_fft_filter(F, randstate=np.random, seed=None): """Produces a field of correlated noise using global Fourier filtering. Parameters ---------- F : array-like Two-dimensional array containing the input filter. It can be computed by related methods. All values are required to be finite. randstate : mtrand.RandomState Optional random generator to use. If set to None, use numpy.random. seed : int Value to set a seed for the generator. None will not set the seed. Returns ------- N : array-like A two-dimensional numpy array of stationary correlated noise. """ if len(F.shape) != 2: raise ValueError("the input is not two-dimensional array") if np.any(~np.isfinite(F)): raise ValueError("F contains non-finite values") # set the seed if seed is not None: randstate.seed(seed) # produce fields of white noise N = randstate.randn(F.shape[0], F.shape[1]) # apply the global Fourier filter to impose a correlation structure fN = fft.fft2(N, **fft_kwargs) fN *= F N = np.array(fft.ifft2(fN, **fft_kwargs).real) N = (N - N.mean()) / N.std() return N
def usecorrelation( im1, im2 ): """Assess the offset (to be used for e.g. the assessment of the center of rotation or the ovarlap) by computation the peak of the correlation between the two input images. Parameters ---------- im1 : array_like Image data as numpy array. im2 : array_like Image data as numpy array. Return value ---------- An integer value of the location of the maximum peak correlation. """ # Fourier transform both images: f_im1 = fft2( im1.astype(float32), threads=2 ); f_im2 = fft2( im2.astype(float32), threads=2 ); # Perform phase correlation (amplitude is normalized): fc = f_im1 * ( f_im2.conjugate() ); fcn = fc / abs(fc); # Inverse fourier of peak correlation matrix and max location: peak_correlation_matrix = real( ifft2( fcn, threads=2 )); # Calculate actual translation: max_ix = argmax(peak_correlation_matrix.flatten()) (row, col) = unravel_index(max_ix, peak_correlation_matrix.shape) if ( col < (peak_correlation_matrix.shape[1]/2) ): col = - (col - 1); else: col = peak_correlation_matrix.shape[1] - (col - 1); return col / 2;
def _phase_correlation(im0, im1, callback=None, *args): """ Computes phase correlation between im0 and im1 Args: im0 im1 callback (function): Process the cross-power spectrum (i.e. choose coordinates of the best element, usually of the highest one). Defaults to :func:`imreg_dft.utils.argmax2D` Returns: tuple: The translation vector (Y, X). Translation vector of (0, 0) means that the two images match. """ if callback is None: callback = utils._argmax2D # TODO: Implement some form of high-pass filtering of PHASE correlation f0, f1 = [fft.fft2(arr) for arr in (im0, im1)] # spectrum can be filtered (already), # so we have to take precaution against dividing by 0 eps = abs(f1).max() * 1e-15 # cps == cross-power spectrum of im0 and im1 cps = abs(fft.ifft2((f0 * f1.conjugate()) / (abs(f0) * abs(f1) + eps))) # scps = shifted cps scps = fft.fftshift(cps) (t0, t1), success = callback(scps, *args) ret = np.array((t0, t1)) # _compensate_fftshift is not appropriate here, this is OK. t0 -= f0.shape[0] // 2 t1 -= f0.shape[1] // 2 ret -= np.array(f0.shape, int) // 2 return ret, success
def phaseflip_star_file(star_file, pixel_size=None, return_in_fourier=False): """ :param star_file: :param pixel_size: :param return_in_fourier: To save computation can skip the ifft. :return: """ # star is a list of star lines describing projections star_records = read_star(star_file)['__root__'] dir_path = os.path.dirname(star_file) stack_info = organize_star_records(star_records) # Initializing projections stack_name = star_records[0].rlnImageName.split('@')[1] mrc_path = os.path.join(os.path.dirname(star_file), stack_name) stack = load_stack_from_file(mrc_path) resolution = stack.shape[1] num_projections = len(star_records) # rfft_resolution = resolution // 2 + 1 # imhat_stack = np.zeros((num_projections, resolution, rfft_resolution), dtype='complex64') # Todo - add temporary dir if return_in_fourier: projections = np.memmap('tmp_phaseflipped_projections.dat', dtype='complex64', mode='w+', shape=(num_projections, resolution, resolution)) else: projections = np.memmap('tmp_phaseflipped_projections.dat', dtype='float32', mode='w+', shape=(num_projections, resolution, resolution)) # Initializing pixel_size if pixel_size is None: tmp = Box(cryo_parse_Relion_CTF_struct(star_records[0])) if tmp.tmppixA != -1: pixel_size = tmp.tmppixA else: raise ValueError( "Pixel size not provided and does not appear in STAR file") # Initializing parameter for cryo_CTF_Relion_fast # a, b, c = precompute_cryo_CTF_Relion_fast(resolution, r=True) a, b, c = precompute_cryo_CTF_Relion_fast(resolution, r=False) num_finished = 0 for stack_name in stack_info: mrc_path = os.path.join(dir_path, stack_name) stack = load_stack_from_file(mrc_path) pos_in_stack = stack_info[stack_name].pos_in_stack pos_in_records = stack_info[stack_name].pos_in_records tic = time.time() for i, j in zip(pos_in_stack, pos_in_records): curr_image = stack[i] curr_records = Box(cryo_parse_Relion_CTF_struct(star_records[j])) curr_records.pixel_size = pixel_size # reference code # h = cryo_CTF_Relion(resolution, curr_records) # imhat = np.fft.fftshift(fft2(curr_image)) # imhat *= np.sign(h) # pfim = ifft2(np.fft.ifftshift(imhat)) # Instead of shifting and shifting back im, shift h (when computing h it is already shifted). h = cryo_CTF_Relion_fast(a, b, c, curr_records) imhat = fft2(curr_image) if return_in_fourier: np.multiply(imhat, np.sign(h), out=projections[j]) else: imhat *= np.sign(h) projections[j] = ifft2(imhat) # Use real fft instead. # h = cryo_CTF_Relion_fast(a, b, c, curr_records) # np.multiply(imhat, np.sign(h), out=imhat_stack[j]) # can irfft2 back the whole stack # imhat2 *= np.sign(h) # pfim2 = irfft2(imhat2) num_finished += 1 toc = time.time() print( 'Finished {} images in {} seconds. In total finished {}/{}'.format( len(pos_in_stack), toc - tic, num_finished, len(star_records))) return projections
def phase_retrieval(im, plan, method=1, nr_threads=2): """Process a tomographic projection image with the selected phase retrieval algorithm. Parameters ---------- im : array_like Flat corrected image data as numpy array. plan : structure Structure with pre-computed data (see prepare_plan function) method : int Phase retrieval algorithm {1 = TIE (default), 2 = Born, 3 = Rytov, 4 = Wu} nr_threads : int Number of threads to be used in the computation of FFT by PyFFTW """ # Extract plan values: dim0_o = plan['dim0'] dim1_o = plan['dim1'] n_pad0 = plan['npad0'] n_pad1 = plan['npad1'] marg0 = (n_pad0 - dim0_o) / 2 marg1 = (n_pad1 - dim1_o) / 2 den = plan['den'] mu = plan['mu'] # Pad image (if required): im = padImage(im, n_pad0, n_pad1) # (Un)comment the following two lines to use PyFFTW: n_byte_align(im, simd_alignment) im = fft2(im, threads=nr_threads) # (Un)comment the following line to use NumPy: #im = fft2(im) # Apply formula: if method == 1: im = im / den # (Un)comment the following two lines to use PyFFTW: n_byte_align(im, simd_alignment) im = ifft2(im, threads=nr_threads) # (Un)comment the following line to use NumPy: #im = ifft2(im) im = im.astype(complex64) im = real(im) im = im.astype(float32) im = -1 / mu * nplog(im) # # WARNING: The following methods are not tested # elif method == 2: im = real(ifft2((im - 1.0) / 2.0) / den) elif method == 3: im = real(ifft2(nplog(im) / 2.0) / den) elif method == 4: im = real(ifft2(im / den)) im = -1 / 2 * (delta / beta) * nplog(im) # Return cropped output: return im[marg0:dim0_o + marg0, marg1:dim1_o + marg1]
def similarity(im0, im1): """Return similarity transformed image im1 and transformation parameters. Transformation parameters are: isotropic scale factor, rotation angle (in degrees), and translation vector. A similarity transformation is an affine transformation with isotropic scale and without shear. Limitations: Image shapes must be equal and square. All image areas must have same scale, rotation, and shift. Scale change must be less than 1.8. No subpixel precision. """ if im0.shape != im1.shape: raise ValueError("Images must have same shapes.") elif len(im0.shape) != 2: raise ValueError("Images must be 2 dimensional.") f0 = fftshift(abs(fft2(im0))) f1 = fftshift(abs(fft2(im1))) h = highpass(f0.shape) f0 *= h f1 *= h del h f0, log_base = logpolar(f0) f1, log_base = logpolar(f1) f0 = fft2(f0) f1 = fft2(f1) r0 = abs(f0) * abs(f1) ir = abs(ifft2((f0 * f1.conjugate()) / r0)) i0, i1 = np.unravel_index(np.argmax(ir), ir.shape) angle = 180.0 * i0 / ir.shape[0] scale = log_base ** i1 if scale > 1.8: ir = abs(ifft2((f1 * f0.conjugate()) / r0)) i0, i1 = np.unravel_index(np.argmax(ir), ir.shape) angle = -180.0 * i0 / ir.shape[0] scale = 1.0 / (log_base ** i1) if scale > 1.8: raise ValueError("Images are not compatible. Scale change > 1.8") if angle < -90.0: angle += 180.0 elif angle > 90.0: angle -= 180.0 im2 = ndii.zoom(im1, 1.0 / scale) im2 = ndii.rotate(im2, angle) if im2.shape < im0.shape: t = np.zeros_like(im0) t[:im2.shape[0], :im2.shape[1]] = im2 im2 = t elif im2.shape > im0.shape: im2 = im2[:im0.shape[0], :im0.shape[1]] f0 = fft2(im0) f1 = fft2(im2) ir = abs(ifft2((f0 * f1.conjugate()) / (abs(f0) * abs(f1)))) t0, t1 = np.unravel_index(np.argmax(ir), ir.shape) if t0 > f0.shape[0] // 2: t0 -= f0.shape[0] if t1 > f0.shape[1] // 2: t1 -= f0.shape[1] im2 = ndii.shift(im2, [t0, t1]) # correct parameters for ndimage's internal processing if angle > 0.0: d = int((int(im1.shape[1] / scale) * math.sin(math.radians(angle)))) t0, t1 = t1, d + t0 elif angle < 0.0: d = int((int(im1.shape[0] / scale) * math.sin(math.radians(angle)))) t0, t1 = d + t1, d + t0 scale = (im1.shape[1] - 1) / (int(im1.shape[1] / scale) - 1) return im2, scale, angle, [-t0, -t1]
def sync_steps(p, v, p_fft, k, kappa, spacing, timestep, density, soundspeed, abs_exp_p, abs_exp_v, source_p, source_v): v = velocity_with_pml(v, ifft2(to_pressure_gradient(p_fft, k, kappa, spacing)), timestep, density, abs_exp_v, source_v) p = pressure_with_pml(p, ifft2(to_velocity_gradient(fft2(v), k, kappa, spacing)), timestep, density, soundspeed, abs_exp_p, source_p) return p, v
def fft_convolve2d(x, f, mode='same', boundary='constant', fft_filter=False): r""" Performs fast 2d convolution in the frequency domain convolving each image channel with its corresponding filter channel. Parameters ---------- x : ``(channels, height, width)`` `ndarray` Image. f : ``(channels, height, width)`` `ndarray` Filter. mode : str {`full`, `same`, `valid`}, optional Determines the shape of the resulting convolution. boundary: str {`constant`, `symmetric`}, optional Determines how the image is padded. fft_filter: `bool`, optional If `True`, the filter is assumed to be defined on the frequency domain. If `False` the filter is assumed to be defined on the spatial domain. Returns ------- c: ``(channels, height, width)`` `ndarray` Result of convolving each image channel with its corresponding filter channel. """ if fft_filter: # extended shape is filter shape ext_shape = np.asarray(f.shape[-2:]) # extend image and filter ext_x = pad(x, ext_shape, boundary=boundary) # compute ffts of extended image fft_ext_x = fft2(ext_x) fft_ext_f = f else: # extended shape x_shape = np.asarray(x.shape[-2:]) f_shape = np.asarray(f.shape[-2:]) f_half_shape = (f_shape / 2).astype(int) ext_shape = x_shape + f_half_shape - 1 # extend image and filter ext_x = pad(x, ext_shape, boundary=boundary) ext_f = pad(f, ext_shape) # compute ffts of extended image and extended filter fft_ext_x = fft2(ext_x) fft_ext_f = fft2(ext_f) # compute extended convolution in Fourier domain fft_ext_c = fft_ext_f * fft_ext_x # compute ifft of extended convolution ext_c = np.real(ifftshift(ifft2(fft_ext_c), axes=(-2, -1))) if mode is 'full': return ext_c elif mode is 'same': return crop(ext_c, x_shape) elif mode is 'valid': return crop(ext_c, x_shape - f_half_shape + 1) else: raise ValueError( "mode={}, is not supported. The only supported " "modes are: 'full', 'same' and 'valid'.".format(mode))
def single_grating_harmonic_images(img, harmonicPeriod, searchRegion=10, plotFlag=False, verbose=False): """ Auxiliary function to process the data of single 2D grating Talbot imaging. It obtain the (real space) harmonic images 00, 01 and 10. Parameters ---------- img : ndarray – Data (data_exchange format) Experimental image, whith proper blank image, crop and rotation already applied. harmonicPeriod : list of integers in the format [periodVert, periodHor] ``periodVert`` and ``periodVert`` are the period of the harmonics in the reciprocal space in pixels. For the checked board grating, periodVert = sqrt(2) * pixel Size / grating Period * number of rows in the image. For 1D grating, set one of the values to negative or zero (it will set the period to number of rows or colunms). searchRegion: int search for the peak will be in a region of harmonicPeriod/searchRegion around the theoretical peak position. See also `:py:func:`wavepy.grating_interferometry.plot_harmonic_grid` plotFlag: boolean Flag to plot the image in the reciprocal space and to show the position of the found peaked and the limits of the harmonic image verbose: Boolean verbose flag. Returns ------- three 2D ndarray data Images obtained from the harmonics 00, 01 and 10. """ imgFFT = np.fft.fftshift(fft2(img, norm='ortho')) if plotFlag: plot_harmonic_grid(imgFFT, harmonicPeriod=harmonicPeriod, isFFT=True) plt.show(block=False) imgFFT00 = extract_harmonic(imgFFT, harmonicPeriod=harmonicPeriod, harmonic_ij='00', searchRegion=searchRegion, isFFT=True, plotFlag=plotFlag, verbose=verbose) imgFFT01 = extract_harmonic(imgFFT, harmonicPeriod=harmonicPeriod, harmonic_ij=['0', '1'], searchRegion=searchRegion, isFFT=True, plotFlag=plotFlag, verbose=verbose) imgFFT10 = extract_harmonic(imgFFT, harmonicPeriod=harmonicPeriod, harmonic_ij=['1', '0'], searchRegion=searchRegion, isFFT=True, plotFlag=plotFlag, verbose=verbose) # Plot Fourier image (intensity) if plotFlag: # Intensity is Fourier Space intFFT00 = np.log10(np.abs(imgFFT00)) intFFT01 = np.log10(np.abs(imgFFT01)) intFFT10 = np.log10(np.abs(imgFFT10)) fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(14, 5)) for dat, ax, textTitle in zip([intFFT00, intFFT01, intFFT10], axes.flat, ['FFT 00', 'FFT 01', 'FFT 10']): # The vmin and vmax arguments specify the color limits im = ax.imshow(dat, cmap='inferno', vmin=np.min(intFFT00), vmax=np.max(intFFT00), extent=wpu.extent_func(dat)) ax.set_title(textTitle) if textTitle == 'FFT 00': ax.set_ylabel('Pixels') ax.set_xlabel('Pixels') # Make an axis for the colorbar on the right side cax = fig.add_axes([0.92, 0.1, 0.03, 0.8]) fig.colorbar(im, cax=cax) plt.suptitle('FFT subsets - Intensity', fontsize=18, weight='bold') plt.show(block=False) img00 = ifft2(np.fft.ifftshift(imgFFT00), norm='ortho') # non existing harmonics will return NAN, so here we check NAN if np.all(np.isfinite(imgFFT01)): img01 = ifft2(np.fft.ifftshift(imgFFT01), norm='ortho') else: img01 = imgFFT01 if np.all(np.isfinite(imgFFT10)): img10 = ifft2(np.fft.ifftshift(imgFFT10), norm='ortho') else: img10 = imgFFT10 return (img00, img01, img10)
def _get_ang_scale(ims, bgval, exponent='inf', constraints=None, reports=None): """ Given two images, return their scale and angle difference. Args: ims (2-tuple-like of 2D ndarrays): The images bgval: We also pad here in the :func:`map_coordinates` exponent (float or 'inf'): The exponent stuff, see :func:`similarity` constraints (dict, optional) reports (optional) Returns: tuple: Scale, angle. Describes the relationship of the subject image to the first one. """ assert len(ims) == 2, \ "Only two images are supported as input" shape = ims[0].shape ims_apod = [utils._apodize(im) for im in ims] dfts = [fft.fftshift(fft.fft2(im)) for im in ims_apod] filt = _logpolar_filter(shape) dfts = [dft * filt for dft in dfts] # High-pass filtering used to be here, but we have moved it to a higher # level interface pcorr_shape = _get_pcorr_shape(shape) log_base = _get_log_base(shape, pcorr_shape[1]) stuffs = [_logpolar(np.abs(dft), pcorr_shape, log_base) for dft in dfts] (arg_ang, arg_rad), success = _phase_correlation( stuffs[0], stuffs[1], utils.argmax_angscale, log_base, exponent, constraints, reports) angle = -np.pi * arg_ang / float(pcorr_shape[0]) angle = np.rad2deg(angle) angle = utils.wrap_angle(angle, 360) scale = log_base ** arg_rad angle = - angle scale = 1.0 / scale if reports is not None: reports["shape"] = filt.shape reports["base"] = log_base if reports.show("spectra"): reports["dfts_filt"] = dfts if reports.show("inputs"): reports["ims_filt"] = [fft.ifft2(np.fft.ifftshift(dft)) for dft in dfts] if reports.show("logpolar"): reports["logpolars"] = stuffs if reports.show("scale_angle"): reports["amas-result-raw"] = (arg_ang, arg_rad) reports["amas-result"] = (scale, angle) reports["amas-success"] = success extent_el = pcorr_shape[1] / 2.0 reports["amas-extent"] = ( log_base ** (-extent_el), log_base ** extent_el, -90, 90 ) if not 0.5 < scale < 2: raise ValueError( "Images are not compatible. Scale change %g too big to be true." % scale) return scale, angle
def upgradePixelPitch(m, N=1, threads=1): """ Go to finer pixels with fourier interpolation. Parameters ---------- m : liteMap The liteMap object holding the data to upgrade the pixel size of. N : int, optional Go to 2^N times smaller pixels. Default is 1. threads : int, optional Number of threads to use in pyFFTW calculations. Default is 1. Returns ------- mNew : liteMap The map with smaller pixels. """ if N < 1: return m.copy() Ny = m.Ny * 2 ** N Nx = m.Nx * 2 ** N npix = Ny * Nx if have_pyFFTW: ft = fft2(m.data, threads=threads) else: ft = fft2(m.data) ftShifted = fftshift(ft) newFtShifted = numpy.zeros((Ny, Nx), dtype=numpy.complex128) # From the numpy.fft.fftshift help: # """ # Shift zero-frequency component to center of spectrum. # # This function swaps half-spaces for all axes listed (defaults to all). # If len(x) is even then the Nyquist component is y[0]. # """ # # So in the case that we have an odd dimension in our map, we want to put # the extra zero at the beginning if m.Nx % 2 != 0: offsetX = (Nx - m.Nx) / 2 + 1 else: offsetX = (Nx - m.Nx) / 2 if m.Ny % 2 != 0: offsetY = (Ny - m.Ny) / 2 + 1 else: offsetY = (Ny - m.Ny) / 2 newFtShifted[offsetY : offsetY + m.Ny, offsetX : offsetX + m.Nx] = ftShifted del ftShifted ftNew = ifftshift(newFtShifted) del newFtShifted # Finally, deconvolve by the pixel window mPix = numpy.copy(numpy.real(ftNew)) mPix[:] = 0.0 mPix[ mPix.shape[0] / 2 - (2 ** (N - 1)) : mPix.shape[0] / 2 + (2 ** (N - 1)), mPix.shape[1] / 2 - (2 ** (N - 1)) : mPix.shape[1] / 2 + (2 ** (N - 1)), ] = (1.0 / (2.0 ** N) ** 2) if have_pyFFTW: ftPix = fft2(mPix, threads=threads) else: ftPix = fft2(mPix) del mPix inds = numpy.where(ftNew != 0) ftNew[inds] /= numpy.abs(ftPix[inds]) if have_pyFFTW: newData = ifft2(ftNew, threads=threads) * (2 ** N) ** 2 else: newData = ifft2(ftNew) * (2 ** N) ** 2 del ftNew del ftPix x0_new, y0_new = m.pixToSky(0, 0) m = m.copy() # don't overwrite original m.wcs.header.update("NAXIS1", 2 ** N * m.wcs.header["NAXIS1"]) m.wcs.header.update("NAXIS2", 2 ** N * m.wcs.header["NAXIS2"]) m.wcs.header.update("CDELT1", m.wcs.header["CDELT1"] / 2.0 ** N) m.wcs.header.update("CDELT2", m.wcs.header["CDELT2"] / 2.0 ** N) m.wcs.updateFromHeader() p_x, p_y = m.skyToPix(x0_new, y0_new) m.wcs.header.update("CRPIX1", m.wcs.header["CRPIX1"] - p_x) m.wcs.header.update("CRPIX2", m.wcs.header["CRPIX2"] - p_y) m.wcs.updateFromHeader() mNew = liteMapFromDataAndWCS(numpy.real(newData), m.wcs) mNew.data[:] = numpy.real(newData[:]) return mNew
def fillWithGRFFromTemplate(self, twodPower, bufferFactor=1, threads=1): """ Generate a Gaussian random field from an input power spectrum specified as a 2d powerMap Notes ----- BufferFactor = 1 means the map will have periodic boundary function, while BufferFactor > 1 means the map will be genrated on a patch bufferFactor times larger in each dimension and then cut out so as to have non-periodic boundary conditions. Fills the data field of the map with the GRF realization. """ ft = fftTools.fftFromLiteMap(self, threads=threads) Ny = self.Ny * bufferFactor Nx = self.Nx * bufferFactor bufferFactor = int(bufferFactor) assert bufferFactor >= 1 realPart = numpy.zeros([Ny, Nx]) imgPart = numpy.zeros([Ny, Nx]) ly = fftfreq(Ny, d=self.pixScaleY) * (2 * numpy.pi) lx = fftfreq(Nx, d=self.pixScaleX) * (2 * numpy.pi) # print ly modLMap = numpy.zeros([Ny, Nx]) iy, ix = numpy.mgrid[0:Ny, 0:Nx] modLMap[iy, ix] = numpy.sqrt(ly[iy] ** 2 + lx[ix] ** 2) # divide out area factor area = twodPower.Nx * twodPower.Ny * twodPower.pixScaleX * twodPower.pixScaleY twodPower.powerMap *= (twodPower.Nx * twodPower.Ny) ** 2 / area if bufferFactor > 1 or twodPower.Nx != Nx or twodPower.Ny != Ny: lx_shifted = fftshift(twodPower.lx) ly_shifted = fftshift(twodPower.ly) twodPower_shifted = fftshift(twodPower.powerMap) f_interp = interp2d(lx_shifted, ly_shifted, twodPower_shifted) # ell = numpy.ravel(twodPower.modLMap) # Cell = numpy.ravel(twodPower.powerMap) # print ell # print Cell # s = splrep(ell,Cell,k=3) # # # ll = numpy.ravel(modLMap) # kk = splev(ll,s) kk = f_interp(fftshift(lx), fftshift(ly)) kk = ifftshift(kk) # id = numpy.where(modLMap > ell.max()) # kk[id] = 0. # add a cosine ^2 falloff at the very end # id2 = numpy.where( (ll> (ell.max()-500)) & (ll<ell.max())) # lEnd = ll[id2] # kk[id2] *= numpy.cos((lEnd-lEnd.min())/(lEnd.max() -lEnd.min())*numpy.pi/2) # pylab.loglog(ll,kk) area = Nx * Ny * self.pixScaleX * self.pixScaleY # p = numpy.reshape(kk,[Ny,Nx]) /area * (Nx*Ny)**2 p = kk # / area * (Nx*Ny)**2 else: area = Nx * Ny * self.pixScaleX * self.pixScaleY p = twodPower.powerMap # /area*(Nx*Ny)**2 realPart = numpy.sqrt(p) * numpy.random.randn(Ny, Nx) imgPart = numpy.sqrt(p) * numpy.random.randn(Ny, Nx) kMap = realPart + 1j * imgPart if have_pyFFTW: data = numpy.real(ifft2(kMap, threads=threads)) else: data = numpy.real(ifft2(kMap)) b = bufferFactor self.data = data[(b - 1) / 2 * self.Ny : (b + 1) / 2 * self.Ny, (b - 1) / 2 * self.Nx : (b + 1) / 2 * self.Nx]
def mapFromFFT(self, kFilter=None, kFilterFromList=None, showFilter=False, setMeanToZero=False, returnFFT=False, threads=1): """ Perform the inverse fft (map from FFT) with an optional filter. Parameters ---------- kFilter : array_like, optional 2D array specifying a k-space filter to apply. If applied, resulting map is IFFT(fft*kFilter). Default is ``None``. kFilterFromList : tuple, optional Tuple of length 2 specifying 1D filter (ell, F_ell), which will be interpolated into a 2D array. Default is ``None``. showFilter : bool, optional Whether to plot the filter. Default is ``False``. setMeanToZero : bool, optional Whether to set the ell = 0 pixel to zero (zeroes mean in real space). Default is ``False``. returnFFT : bool, optional Whether to return the fftTools.fft2D class as well. Default is ``False``. threads : int, optional Number of threads to use in pyFFTW calculations. Default is 1. Returns ------- data : array_like The (optinally filtered) 2D real space data array ftMap : fftTools.fft2D, optional The fft2D object, returned if returnFFT = ``True``. """ kMap = self.kMap.copy() kFilter0 = numpy.real(kMap.copy())*0.+ 1. if kFilter != None: kFilter0 *= kFilter if kFilterFromList != None: kFilter = kMap.copy()*0. l = kFilterFromList[0] Fl = kFilterFromList[1] FlSpline = splrep(l,Fl,k=3) ll = numpy.ravel(self.modLMap) kk = (splev(ll,FlSpline)) kFilter = numpy.reshape(kk,[self.Ny,self.Nx]) kFilter0 *= kFilter if setMeanToZero: id = numpy.where(self.modLMap == 0.) kFilter0[id] = 0. if showFilter: pylab.semilogy(l,Fl,'r',ll,kk,'b.') pylab.matshow(fftshift(kFilter0),origin="down",extent=[numpy.min(self.lx),\ numpy.max(self.lx),\ numpy.min(self.ly),\ numpy.max(self.ly)]) pylab.show() kMap[:,:] *= kFilter0[:,:] if have_pyFFTW: data = numpy.real(ifft2(kMap, threads=threads)) else: data = numpy.real(ifft2(kMap)) if returnFFT: ftMap = self.copy() ftMap.kMap = kMap.copy() return data, ftMap else: return data