def run(self): max_shape = self._find_max_shape() # compute FT, assuming they are the same size fft1 = cp.asarray(self.image1, dtype=cp.complex64) fft2 = cp.asarray(self.image2, dtype=cp.complex64) plan = get_fft_plan(fft1, value_type="C2C") fft1 = fftn(fft1, overwrite_x=True, plan=plan) fft2 = fftn(fft2, overwrite_x=True, plan=plan) print(f"shape: {fft1.shape}, dtype: {fft1.dtype}") @cp.fuse def normalize(fft_image): re, im = cp.real(fft_image), cp.imag(fft_image) length = cp.sqrt(re * re + im * im) return fft_image / length fft1 = normalize(fft1) fft2 = cp.conj(normalize(fft2)) # phase correlation spectrum pcm = fft1 * fft2 pcm = ifftn(pcm, overwrite_x=True, plan=plan) pcm = cp.real(pcm) from skimage.morphology import disk from skimage.filters import median pcm = cp.asnumpy(pcm) pcm = median(pcm, disk(3)) pcm = cp.asarray(pcm) peak_list = self._extract_correlation_peaks(pcm)
def test_fftn_multiple_plan_error(self, dtype): import cupy import cupyx.scipy.fftpack as fftpack x = testing.shaped_random(self.shape, cupy, dtype) # hack: avoid testing the cases when getting a cuFFT plan is impossible if _default_fft_func(x, s=self.s, axes=self.axes) is not _fftn: return plan = fftpack.get_fft_plan(x, shape=self.s, axes=self.axes) with pytest.raises(RuntimeError) as ex, plan: fftpack.fftn(x, shape=self.s, axes=self.axes, plan=plan) assert 'Use the cuFFT plan either as' in str(ex.value)
def wiener_deconvolution(image, psf, snr=30, add_pad=0): """ A GPU accelerated implementation of a linear Wiener filter. Some effort is made to allow processing even relatively large images, but some kind of block-based processing (as in the RL implementation) may be required in some cases.""" assert isinstance(image, Image) assert isinstance(psf, Image) image_s = Image(image.copy(), image.spacing) orig_shape = image.shape if image.ndim != psf.ndim: raise ValueError("Image and psf dimensions do not match") if psf.spacing != image.spacing: psf = imops.zoom_to_spacing(psf, image.spacing) if add_pad != 0: new_shape = list(i + 2 * add_pad for i in image_s.shape) image_s = imops.zero_pad_to_shape(image_s, new_shape) if psf.shape != image_s.shape: psf = imops.zero_pad_to_shape(psf, image_s.shape) psf /= psf.max() psf = fftshift(psf) psf_dev = cp.asarray(psf.astype(np.complex64)) with get_fft_plan(psf_dev): psf_dev = fftn(psf_dev, overwrite_x=True) below = cp.asnumpy(psf_dev) psf_abs = cp.abs(psf_dev)**2 psf_abs /= (psf_abs + snr) above = cp.asnumpy(psf_abs) psf_abs = None psf_dev = None image_dev = cp.asarray(image_s.astype(np.complex64)) with get_fft_plan(image_dev): image_dev = fftn(image_dev, overwrite_x=True) wiener_dev = cp.asarray(arrayops.safe_divide(above, below)) image_dev *= wiener_dev result = cp.asnumpy(cp.abs(ifftn(image_dev, overwrite_x=True)).real) result = Image(result, image.spacing) return imops.remove_zero_padding(result, orig_shape)
def transform(self, data): if self.f_contiguous: _dat = data.T else: _dat = data _dat[:] = cufftp.fftn(_dat, axes=self._ax, plan=self._fftplan)[:] """The transform is done inplace"""
def __get_fourier_psfs(self): """ Pre-calculates the PSFs during image fusion process. """ print("Pre-calculating PSFs") padded_block_size = tuple(self.block_size + 2 * self.options.block_pad) memmap_shape = (self.n_views, ) + padded_block_size if self.options.disable_fft_psf_memmap: self.psfs_fft = numpy.zeros(memmap_shape, dtype=numpy.complex64) self.adj_psfs_fft = numpy.zeros(memmap_shape, dtype=numpy.complex64) else: psfs_fft_f = os.path.join(self.memmap_directory, "psf_fft_f.dat") self.psfs_fft = numpy.memmap(psfs_fft_f, dtype='complex64', mode='w+', shape=memmap_shape) adj_psfs_fft_f = os.path.join(self.memmap_directory, "adj_psf_fft_f.dat") self.adj_psfs_fft = numpy.memmap(adj_psfs_fft_f, dtype='complex64', mode='w+', shape=memmap_shape) for idx in range(self.n_views): self.psfs_fft[idx] = ops_array.expand_to_shape( self.psfs[idx], padded_block_size).astype(numpy.complex64) self.adj_psfs_fft[idx] = ops_array.expand_to_shape( self.adj_psfs[idx], padded_block_size).astype(numpy.complex64) self.psfs_fft[idx] = numpy.fft.fftshift(self.psfs_fft[idx]) self.adj_psfs_fft[idx] = numpy.fft.fftshift(self.adj_psfs_fft[idx]) self.psfs_fft[idx] = cp.asnumpy( fftpack.fftn(cp.asarray(self.psfs_fft[idx]), plan=self._fft_plan)) self.adj_psfs_fft[idx] = cp.asnumpy( fftpack.fftn(cp.asarray(self.adj_psfs_fft[idx]), plan=self._fft_plan))
def _fft_convolve(self, h_data, h_kernel): """ Calculate a convolution on GPU, using FFTs. :param h_data: a Numpy array with the data to convolve :param h_kernel: a Numpy array with the convolution kernel. The kernel should already be in Fourier domain (To avoid repeating the transform at every iteration.) """ #todo: See whether to add back streams. I removed them on Cupy refactor. d_data = cp.asarray(h_data) d_data = fftpack.fftn(d_data, overwrite_x=True, plan=self._fft_plan) d_kernel = cp.asarray(h_kernel) d_data *= d_kernel d_data = fftpack.ifftn(d_data, overwrite_x=True, plan=self._fft_plan) return cp.asnumpy(d_data)
def phasecorr_gpu(X, cfRefImg, lcorr): ''' not being used - no speed up - may be faster with cuda.jit''' nimg, Ly, Lx = X.shape ly, lx = cfRefImg.shape[-2:] lyhalf = int(np.floor(ly / 2)) lxhalf = int(np.floor(lx / 2)) # put on GPU ref_gpu = cp.asarray(cfRefImg) x_gpu = cp.asarray(X) # phasecorrelation x_gpu = fftn(x_gpu, axes=(1, 2), overwrite_x=True) * np.sqrt(Ly - 1) * np.sqrt(Lx - 1) for t in range(x_gpu.shape[0]): tmp = x_gpu[t, :, :] tmp = cp.multiply(tmp, ref_gpu) tmp = cp.divide(tmp, cp.absolute(tmp) + 1e-5) x_gpu[t, :, :] = tmp x_gpu = ifftn(x_gpu, axes=(1, 2), overwrite_x=True) * np.sqrt(Ly - 1) * np.sqrt(Lx - 1) x_gpu = cp.fft.fftshift(cp.real(x_gpu), axes=(1, 2)) # get max index x_gpu = x_gpu[cp.ix_(np.arange(0, nimg, 1, int), np.arange(lyhalf - lcorr, lyhalf + lcorr + 1, 1, int), np.arange(lxhalf - lcorr, lxhalf + lcorr + 1, 1, int))] ix = cp.argmax(cp.reshape(x_gpu, (nimg, -1)), axis=1) cmax = x_gpu[np.arange(0, nimg, 1, int), ix] ymax, xmax = cp.unravel_index(ix, (2 * lcorr + 1, 2 * lcorr + 1)) cmax = cp.asnumpy(cmax).flatten() ymax = cp.asnumpy(ymax) xmax = cp.asnumpy(xmax) ymax, xmax = ymax - lcorr, xmax - lcorr return ymax, xmax, cmax
def soft_mask(v, voxel_size, num_subunit_residues, helical_repeat_distance=None, repeats_to_include=0, filter_resolution=20, expansion_factor=1.2, expansion_radius=0, print_progress=True, return_mask=False): full_expansion_radius = expansion_radius + filter_resolution/2 # avg AA mol wt. in g/mol, density in g/cm3 avg_aa_molwt = 110 protein_density = 1.4 # 2 print helical_repeat_distance v_thresh = np.zeros(v.shape) v_thresh[:] = v[:] sz = np.array(v.shape).astype(int) total_molwt = num_subunit_residues*avg_aa_molwt/6.023e23 if helical_repeat_distance != None: total_molwt = total_molwt * sz[2]*voxel_size / helical_repeat_distance total_vol = np.prod(sz) * voxel_size**3 # vol in A3 mol_vol = total_molwt/protein_density / (1.0e-24) # vol in A3 mol_vol_frac = mol_vol/total_vol target_vol_frac = mol_vol_frac*expansion_factor thresh = find_binary_threshold(v_thresh, target_vol_frac) true_frac = (0.0 + np.sum(v_thresh >= thresh)) / v_thresh.size if repeats_to_include != 0: zdim = np.round(repeats_to_include * helical_repeat_distance/voxel_size) else: zdim = sz[2] if zdim > sz[2] - 4*np.ceil(filter_resolution/voxel_size): zdim = sz[2] - 4*np.ceil(filter_resolution/voxel_size) zdim = zdim.astype(int) v_thresh[:,:,0:np.floor(sz[2]/2).astype(int) - np.floor(zdim/2).astype(int)] = 0 v_thresh[:,:,np.floor(sz[2]/2).astype(int) - np.floor(zdim/2).astype(int) + 1 + zdim - 1:] = 0 v_thresh[v_thresh < thresh] = 0 if print_progress: print 'Target volume fraction: {}'.format(target_vol_frac) print 'Achieved volume fraction: {}'.format(true_frac) print 'Designated threshold: {}'.format(thresh) progress_bar = tqdm(total=5) v_thresh = fftpack.fftn(v_thresh) progress_bar.update(1) # 3 cosmask_filter = np.fft.fftshift(spherical_cosmask(sz, 0, np.ceil(filter_resolution/voxel_size))) cosmask_filter = fftpack.fftn(cosmask_filter) / np.sum(cosmask_filter) progress_bar.update(1) v_thresh = v_thresh * cosmask_filter v_thresh = fftpack.ifftn(v_thresh) progress_bar.update(1) v_thresh = np.real(v_thresh) v_thresh[np.abs(v_thresh) < 10*np.finfo(type(v_thresh.ravel()[0])).eps] = 0 v_thresh[v_thresh != 0] = 1 # The extent of blurring is equal to the diameter of the cosmask sphere; # if we want this to equal the expected falloff for filter_resolution, # we therefore need to divide filter_res by 4 to get the # desired radius for spherical_cosmask. v_thresh = fftpack.fftn(v_thresh) progress_bar.update(1) v_thresh = v_thresh * cosmask_filter v_thresh = fftpack.ifftn(v_thresh) progress_bar.update(1) v_thresh = np.real(v_thresh) v_thresh[np.abs(v_thresh) < 10*np.finfo(type(v_thresh.ravel()[0])).eps] = 0 if return_mask: v[:,:,:] = v_thresh else: v *= v_thresh return v_thresh
def fftconvolve(in1, in2, mode="full", axes=None): """Convolve two N-dimensional arrays using FFT. Convolve `in1` and `in2` using the fast Fourier transform method, with the output size determined by the `mode` argument. This is generally much faster than `convolve` for large arrays (n > ~500), but can be slower when only a few output values are needed, and can only output float arrays (int or object array inputs will be cast to float). As of v0.19, `convolve` automatically chooses this method or the direct method based on an estimation of which is faster. Parameters ---------- in1 : array_like First input. in2 : array_like Second input. Should have the same number of dimensions as `in1`. mode : str {'full', 'valid', 'same'}, optional A string indicating the size of the output: ``full`` The output is the full discrete linear convolution of the inputs. (Default) ``valid`` The output consists only of those elements that do not rely on the zero-padding. In 'valid' mode, either `in1` or `in2` must be at least as large as the other in every dimension. ``same`` The output is the same size as `in1`, centered with respect to the 'full' output. axis : tuple, optional axes : int or array_like of ints or None, optional Axes over which to compute the convolution. The default is over all axes. Returns ------- out : array An N-dimensional array containing a subset of the discrete linear convolution of `in1` with `in2`. Examples -------- Autocorrelation of white noise is an impulse. >>> import cusignal >>> import cupy as cp >>> import numpy as np >>> sig = cp.random.randn(1000) >>> autocorr = cusignal.fftconvolve(sig, sig[::-1], mode='full') >>> import matplotlib.pyplot as plt >>> fig, (ax_orig, ax_mag) = plt.subplots(2, 1) >>> ax_orig.plot(cp.asnumpy(sig)) >>> ax_orig.set_title('White noise') >>> ax_mag.plot(np.arange(-len(sig)+1,len(sig)), autocorr) >>> ax_mag.set_title('Autocorrelation') >>> fig.tight_layout() >>> fig.show() Gaussian blur implemented using FFT convolution. Notice the dark borders around the image, due to the zero-padding beyond its boundaries. The `convolve2d` function allows for other types of image boundaries, but is far slower. >>> from scipy import misc >>> face = misc.face(gray=True) >>> kernel = cp.outer(cusignal.gaussian(70, 8), cusignal.gaussian(70, 8)) >>> blurred = cusignal.fftconvolve(face, kernel, mode='same') >>> fig, (ax_orig, ax_kernel, ax_blurred) = plt.subplots(3, 1, ... figsize=(6, 15)) >>> ax_orig.imshow(face, cmap='gray') >>> ax_orig.set_title('Original') >>> ax_orig.set_axis_off() >>> ax_kernel.imshow(cp.asnumpy(kernel), cmap='gray') >>> ax_kernel.set_title('Gaussian kernel') >>> ax_kernel.set_axis_off() >>> ax_blurred.imshow(cp.asnumpy(blurred), cmap='gray') >>> ax_blurred.set_title('Blurred') >>> ax_blurred.set_axis_off() >>> fig.show() """ in1 = cp.ascontiguousarray(in1) in2 = cp.ascontiguousarray(in2) noaxes = axes is None if in1.ndim == in2.ndim == 0: # scalar inputs return in1 * in2 elif in1.ndim != in2.ndim: raise ValueError("in1 and in2 should have the same dimensionality") elif in1.size == 0 or in2.size == 0: # empty arrays return cp.array([]) _, axes = _init_nd_shape_and_axes_sorted(in1, shape=None, axes=axes) # axes needs to be numpy type for proper execution in FFT axes = cp.asnumpy(axes) if not noaxes and not axes.size: raise ValueError("when provided, axes cannot be empty") if noaxes: other_axes = np.array([], dtype=cp.intc) else: other_axes = np.setdiff1d(np.arange(in1.ndim), axes) s1 = np.array(in1.shape) s2 = np.array(in2.shape) if not np.all((s1[other_axes] == s2[other_axes]) | (s1[other_axes] == 1) | (s2[other_axes] == 1)): raise ValueError("incompatible shapes for in1 and in2:" " {0} and {1}".format(in1.shape, in2.shape)) complex_result = np.issubdtype(in1.dtype, np.complexfloating) or np.issubdtype( in2.dtype, cp.complexfloating) shape = np.maximum(s1, s2) shape[axes] = s1[axes] + s2[axes] - 1 # Check that input sizes are compatible with 'valid' mode if _inputs_swap_needed(mode, s1, s2): # Convolution is commutative; order doesn't have any effect on output in1, s1, in2, s2 = in2, s2, in1, s1 # Speed up FFT by padding to optimal size for FFTPACK fshape = [next_fast_len(d) for d in shape[axes]] fslice = tuple([slice(sz) for sz in shape]) if not complex_result: if (str(fshape), str(axes), "R2C") in _cupy_fft_cache: rplan = _cupy_fft_cache[(str(fshape), str(axes), "R2C")] else: rplan = _cupy_fft_cache[(str(fshape), str(axes), "R2C")] = fftpack.get_fft_plan( in1, fshape, axes=axes, value_type="R2C") try: with rplan: sp1 = cp.fft.rfftn(in1, fshape, axes=axes) sp2 = cp.fft.rfftn(in2, fshape, axes=axes) except Exception: sp1 = cp.fft.rfftn(in1, fshape, axes=axes) sp2 = cp.fft.rfftn(in2, fshape, axes=axes) ret = cp.fft.irfftn(sp1 * sp2, fshape, axes=axes)[fslice].copy() else: # Need to move to cupyx.scipy.fft with CuPy v8 if (str(fshape), str(axes)) in _cupy_fft_cache: plan = _cupy_fft_cache[(str(fshape), str(axes))] else: plan = _cupy_fft_cache[(str(fshape), str(axes))] = fftpack.get_fft_plan( in1, fshape, axes=axes) try: with plan: sp1 = fftpack.fftn(in1, fshape, axes=axes) sp2 = fftpack.fftn(in2, fshape, axes=axes) ret = fftpack.ifftn(sp1 * sp2, axes=axes)[fslice].copy() except Exception: sp1 = fftpack.fftn(in1, fshape, axes=axes) sp2 = fftpack.fftn(in2, fshape, axes=axes) ret = fftpack.ifftn(sp1 * sp2, axes=axes)[fslice].copy() if mode == "full": return ret elif mode == "same": return _centered(ret, s1) elif mode == "valid": shape_valid = shape.copy() shape_valid[axes] = s1[axes] - s2[axes] + 1 return _centered(ret, shape_valid) else: raise ValueError("acceptable mode flags are \ 'valid'," " 'same', or 'full'")