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 rfft2(self,z): m = z.shape[0] if self._plan_fft2 is None: self._plan_fft2 = cpxFFT.get_fft_plan(z + 1j*cp.zeros_like(z), shape=z.shape) temp = cpxFFT.fft2(z.astype(cp.complex64),shape=z.shape,plan=self._plan_fft2,overwrite_x=True) temp = cp.fft.fftshift(temp,axes=0) return temp[m//2 -(self.basis_number-1):m//2 +self.basis_number,:self.basis_number ]/(2*self.extended_basis_number-1)
def __enter__(self): farplane = cp.empty( (self.nwaves, self.detector_shape, self.detector_shape), dtype='complex64') self.plan = get_fft_plan(farplane, axes=(-2, -1)) del farplane return self
def ifft2(x): global plan2D if plan2D == None and Plan: plan2D = fftpack.get_fft_plan(x, axes=(-2, -1)) # return return fftpack.ifft2(x, overwrite_x=owrite, plan=plan2D)
def test_fftn_error_on_wrong_plan(self, dtype, enable_nd): # This test ensures the context manager plan is picked up from cupyx.scipy.fftpack import get_fft_plan from cupy.fft import fftn assert config.enable_nd_planning == enable_nd # can't get a plan, so skip if self.axes is not None: if self.s is not None: if len(self.s) != len(self.axes): return elif len(self.shape) != len(self.axes): return a = testing.shaped_random(self.shape, cupy, dtype) bad_in_shape = tuple(2 * i for i in self.shape) if self.s is None: bad_out_shape = bad_in_shape else: bad_out_shape = tuple(2 * i for i in self.s) b = testing.shaped_random(bad_in_shape, cupy, dtype) plan_wrong = get_fft_plan(b, bad_out_shape, self.axes) with pytest.raises(ValueError) as ex, plan_wrong: fftn(a, s=self.s, axes=self.axes, norm=self.norm) # targeting a particular error assert 'The CUFFT plan and a.shape do not match' in str(ex.value)
def fft(x): global plan1D #print("Plan1D",plan1D) if plan1D == None and Plan: plan1D = fftpack.get_fft_plan(x, axes=(-1)) return fftpack.fft(x, overwrite_x=owrite, plan=plan1D)
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 _get_fft_plan(self, a, axes, **kwargs): """Cache multiple FFT plans at the same time.""" key = (*a.shape, *axes) if key in self.plan_cache: plan = self.plan_cache[key] else: plan = get_fft_plan(a, axes=axes) self.plan_cache[key] = plan return plan
def _get_fft_plan(self, a, axes=None, **kwargs): """Cache multiple FFT plans at the same time.""" axes = tuple(range(a.ndim)) if axes is None else axes key = (*a.shape, *axes) if key in self.plan_cache: plan = self.plan_cache[key] else: plan = get_fft_plan(a, axes=axes) self.plan_cache[key] = plan return plan
def cupy_prop(self, signal): if not signal.is_pol: raise Exception("only dp signal supported at this time") step_number = self.length / self.step_length step_number = int(np.floor(step_number)) temp = np.zeros_like(signal[:]) freq = fftfreq(signal.shape[1], 1 / signal.fs_in_fiber) freq_gpu = cp.asarray(freq) omeg = 2 * np.pi * freq_gpu D = -1j / 2 * self.beta2(signal.center_wavelength) * omeg**2 N = 8 / 9 * 1j * self.gamma atten = -self.alpha_lin / 2 time_x = cp.asarray(signal[0, :]) time_y = cp.asarray(signal[1, :]) plan = get_fft_plan(time_x) for i in range(step_number): time_x, time_y = self.linear_prop_cupy(D, time_x, time_y, self.step_length / 2, plan) time_x, time_y = self.nonlinear_prop_cupy(N, time_x, time_y) time_x = time_x * math.exp(atten * self.step_length) time_y = time_y * math.exp(atten * self.step_length) time_x, time_y = self.linear_prop_cupy(D, time_x, time_y, self.step_length / 2, plan) last_step = self.length - self.step_length * step_number last_step_eff = (1 - np.exp(-self.alpha_lin * last_step)) / self.alpha_lin if last_step == 0: time_x = cp.asnumpy(time_x) time_y = cp.asnumpy(time_y) temp[0, :] = time_x temp[1, :] = time_y return temp else: time_x, time_y = self.linear_prop_cupy(D, time_x, time_y, last_step / 2, plan) time_x, time_y = self.nonlinear_prop_cupy(N, time_x, time_y, last_step_eff) time_x = time_x * math.exp(atten * last_step) time_y = time_y * math.exp(atten * last_step) time_x, time_y = self.linear_prop_cupy(D, time_x, time_y, last_step / 2, plan) temp[0, :] = cp.asnumpy(time_x) temp[1, :] = cp.asnumpy(time_y) return temp
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 __init__(self, image, psf, writer, options): """ :param image: the image as a Image object :param psf: the psf as an Image object :param options: command line options that control the behavior of the fusion algorithm """ deconvolve.DeconvolutionRL.__init__(self, image, psf, writer, options) self._fft_plan = fftpack.get_fft_plan(cp.zeros(self.block_size, dtype=cp.complex64)) self.__get_fourier_psfs()
def test_fft_multiple_plan_error(self, dtype): # hack: avoid testing the cases when the output array is of size 0 # because cuFFT and numpy raise different kinds of exceptions if self.n == 0: return import cupy import cupyx.scipy.fftpack as fftpack x = testing.shaped_random(self.shape, cupy, dtype) plan = fftpack.get_fft_plan(x, shape=self.n, axes=self.axis) with pytest.raises(RuntimeError) as ex, plan: fftpack.fft(x, n=self.n, axis=self.axis, plan=plan) assert 'Use the cuFFT plan either as' in str(ex.value)
def irfft2(self,uHalf2D): """ Fourier transform of one dimensional signal ut = 1D signal num = Ut length - 1 dt = timestep (now using cp.fft.fft) in the implementation """ u2D = util.extend2D(util.symmetrize_2D(uHalf2D),self.extended_basis_number) if self._plan_ifft2 is None: self._plan_ifft2 = cpxFFT.get_fft_plan(u2D, shape=u2D.shape) u2D = cp.fft.ifftshift(u2D) temp = cpxFFT.ifft2(u2D,shape=u2D.shape,plan=self._plan_ifft2,overwrite_x=True) return temp.real*(2*self.extended_basis_number-1)
def test_ifftn(self, xp, dtype, enable_nd): assert config.enable_nd_planning == enable_nd a = testing.shaped_random(self.shape, xp, dtype) if xp == cupy: from cupyx.scipy.fftpack import get_fft_plan plan = get_fft_plan(a, self.s, self.axes) with plan: out = xp.fft.ifftn(a, s=self.s, axes=self.axes, norm=self.norm) else: out = xp.fft.ifftn(a, s=self.s, axes=self.axes, norm=self.norm) if xp == np and dtype is np.complex64: out = out.astype(np.complex64) return out
def test_rfftn(self, xp, dtype, enable_nd): assert config.enable_nd_planning == enable_nd a = testing.shaped_random(self.shape, xp, dtype) if xp is cupy: from cupyx.scipy.fftpack import get_fft_plan plan = get_fft_plan(a, self.s, self.axes, value_type='R2C') with plan: out = xp.fft.rfftn(a, s=self.s, axes=self.axes, norm=self.norm) else: out = xp.fft.rfftn(a, s=self.s, axes=self.axes, norm=self.norm) if xp is np and dtype in [np.float16, np.float32, np.complex64]: out = out.astype(np.complex64) return out
def __init__(self, data, writer, options): """ :param data: a ImageData object :param options: command line options that control the behavior of the fusion algorithm :param writer: a writer object that can save intermediate results """ fusion.MultiViewFusionRL.__init__(self, data, writer, options) padded_block_size = self.block_size + 2 * self.options.block_pad self._fft_plan = fftpack.get_fft_plan( cp.zeros(padded_block_size, dtype=cp.complex64)) self.__get_fourier_psfs()
def test_fft_error_on_wrong_plan(self, dtype): # This test ensures the context manager plan is picked up from cupyx.scipy.fftpack import get_fft_plan from cupy.fft import fft a = testing.shaped_random(self.shape, cupy, dtype) bad_shape = tuple(5 * i for i in self.shape) b = testing.shaped_random(bad_shape, cupy, dtype) plan_wrong = get_fft_plan(b) assert isinstance(plan_wrong, cupy.cuda.cufft.Plan1d) with pytest.raises(ValueError) as ex, plan_wrong: fft(a, n=self.n, norm=self.norm) # targeting a particular error assert 'Target array size does not match the plan.' in str(ex.value)
def test_ifft(self, xp, dtype): a = testing.shaped_random(self.shape, xp, dtype) if xp == cupy: from cupyx.scipy.fftpack import get_fft_plan shape = (self.n, ) if self.n is not None else None plan = get_fft_plan(a, shape=shape) assert isinstance(plan, cupy.cuda.cufft.Plan1d) with plan: out = xp.fft.ifft(a, n=self.n, norm=self.norm) else: out = xp.fft.ifft(a, n=self.n, norm=self.norm) if xp == np and dtype is np.complex64: out = out.astype(np.complex64) return out
def test_irfft(self, xp, dtype): a = testing.shaped_random(self.shape, xp, dtype) if xp is cupy: from cupyx.scipy.fftpack import get_fft_plan shape = (self.n,) if self.n is not None else None plan = get_fft_plan(a, shape=shape, value_type='C2R') assert isinstance(plan, cupy.cuda.cufft.Plan1d) with plan: out = xp.fft.irfft(a, n=self.n, norm=self.norm) else: out = xp.fft.irfft(a, n=self.n, norm=self.norm) if xp is np and dtype in [np.float16, np.float32, np.complex64]: out = out.astype(np.float32) return out
def __init__(self, context, data, axes): self.context = context self.axes = axes assert len(data.shape) > max(axes) from cupyx.scipy import fftpack as cufftp if data.flags.f_contiguous: self._ax = [data.ndim - 1 - aa for aa in axes] _dat = data.T self.f_contiguous = True else: self._ax = axes _dat = data self.f_contiguous = False self._fftplan = cufftp.get_fft_plan(_dat, axes=self._ax, value_type="C2C")
def ifft(x): global plan1D if plan1D == None and Plan: plan1D = fftpack.get_fft_plan(x, axes=(-1)) return fftpack.ifft(x, overwrite_x=owrite, plan=plan1D)
def channelize_poly(x, h, n_chans): """ Polyphase channelize signal into n channels Parameters ---------- x : array_like The input data to be channelized h : array_like The 1-D input filter; will be split into n channels of int number of taps n_chans : int Number of channels for channelizer Returns ---------- yy : channelized output matrix Notes ---------- Currently only supports simple channelizer where channel spacing is equivalent to the number of channels used (zero overlap). Number of filter taps (len of filter / n_chans) must be <=32. """ dtype = cp.promote_types(x.dtype, h.dtype) x = asarray(x, dtype=dtype) h = asarray(h, dtype=dtype) # number of taps in each h_n filter n_taps = int(len(h) / n_chans) if n_taps > 32: raise NotImplementedError( "The number of calculated taps ({}) in \ each filter is currently capped at 32. Please reduce filter \ length or number of channels".format( n_taps ) ) if n_taps > 32: raise NotImplementedError( "Number of taps ({}) must be less than (32).".format(n_taps) ) # number of outputs n_pts = int(len(x) / n_chans) if x.dtype == cp.float32 or x.dtype == cp.complex64: y = cp.empty((n_pts, n_chans), dtype=cp.complex64) elif x.dtype == cp.float64 or x.dtype == cp.complex128: y = cp.empty((n_pts, n_chans), dtype=cp.complex128) _channelizer(x, h, y, n_chans, n_taps, n_pts) # Remove with CuPy v8 if (x.dtype, n_pts, n_taps, n_chans) in _cupy_fft_cache: plan = _cupy_fft_cache[(x.dtype, n_pts, n_taps, n_chans)] else: plan = _cupy_fft_cache[ (x.dtype, n_pts, n_taps, n_chans) ] = fftpack.get_fft_plan(y, axes=-1) return cp.conj(fftpack.fft(y, overwrite_x=True, plan=plan)).T
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'")
def fiber_cupy(param: FiberParam, samples: np.ndarray, fs: float, center_wavelength: float, device: str): import cupy as np from cupyx.scipy.fftpack import fft as improved_fft from cupyx.scipy.fftpack import ifft as improved_ifft from cupyx.scipy.fftpack import get_fft_plan samples = np.array(samples) plan = get_fft_plan(samples[0]) def step(samples, step_length, step_eff): xpol = samples[0] ypol = samples[1] xpol, ypol = linear_prop(xpol, ypol, step_length / 2) xpol, ypol = nonlinear_prop(xpol, ypol, step_eff) xpol, ypol = atteunation(xpol, ypol, step_length) xpol, ypol = linear_prop(xpol, ypol, step_length / 2) return np.vstack((xpol, ypol)) def atteunation(xpol, ypol, step_length): xpol = xpol * np.exp(atten * step_length) ypol = ypol * np.exp(atten * step_length) return xpol, ypol def nonlinear_prop(xpol, ypol, step_length): amplitude_xpol = (xpol[:, 0]**2 + xpol[:, 1]**2) amplitude_ypol = (ypol[:, 0]**2 + ypol[:, 1]**2) phase_rotation_xpol = 8 / 9 * (amplitude_xpol + amplitude_ypol) * gamma * step_length phase_rotation_ypol = phase_rotation_xpol xpol = xpol * np.exp(1j * 2 * np.pi * phase_rotation_xpol) ypol = ypol * np.exp(1j * 2 * np.pi * phase_rotation_ypol) return xpol, ypol def linear_prop(xpol, ypol, length): frequency_domain_xpol = improved_fft(xpol, overwrite_x=True, plan=plan) frequency_domain_ypol = improved_fft(ypol, overwrite_x=True, plan=plan) frequency_domain_xpol = frequency_domain_xpol * np.exp(1j * D * length) frequency_domain_ypol = frequency_domain_ypol * np.exp(1j * D * length) xpol_time_domain = improved_ifft(frequency_domain_xpol, overwrite_x=True, plan=plan) ypol_time_domain = improved_ifft(frequency_domain_ypol, overwrite_x=True, plan=plan) return xpol_time_domain, ypol_time_domain gamma = param.gamma length = param.length step_length = param.step_length xpol = samples[0] freq = np.fft.fftfreq(len(xpol), 1 / fs) omeg = 2 * np.pi * freq D = -1 / 2 * param.beta2(center_wavelength) * omeg**2 atten = -param.alphalin / 2 step_number = length / step_length step_number = int(np.floor(step_number)) last_length = length - step_number * step_length step_eff = 1 - np.exp(-param.alphalin * step_length) step_eff = step_eff / param.alphalin last_length_eff = 1 - np.exp(-param.alphalin * last_length) last_length_eff = last_length_eff / param.alphalin for step_index in range(step_number): samples = step(samples, step_length, step_eff) if last_length_eff: samples = step(samples, last_length, last_length_eff) return samples