def main(args): log = logging.getLogger('root') hdlr = logging.StreamHandler(sys.stdout) log.addHandler(hdlr) log.setLevel(logging.getLevelName(args.loglevel.upper())) pyfftw.interfaces.cache.enable() vol1 = mrc.read(args.volume1, inc_header=False, compat="relion") vol2 = mrc.read(args.volume2, inc_header=False, compat="relion") if args.mask is not None: mask = mrc.read(args.mask, inc_header=False, compat="relion") vol1 *= mask vol2 *= mask f3d1 = fft.rfftn(vol1, threads=args.threads) f3d2 = fft.rfftn(vol2, threads=args.threads) nside = 2**args.healpix_order x, y, z = pix2vec(nside, np.arange(12 * nside ** 2)) xhalf = x >= 0 hp = np.column_stack([x[xhalf], y[xhalf], z[xhalf]]) t0 = time.time() fcor = calc_dfsc(f3d1, f3d2, hp, np.deg2rad(args.arc)) log.info("Computed CFSC in %0.2f s" % (time.time() - t0)) fsc = calc_fsc(f3d1, f3d2) t0 = time.time() log.info("Computed GFSC in %0.2f s" % (time.time() - t0)) freqs = np.fft.rfftfreq(f3d1.shape[0]) np.save(args.output, np.row_stack([freqs, fsc, fcor])) return 0
def fftconvolve(in1, in2, mode="same", threads=1): """Same as above but with pyfftw added in""" in1 = np.asarray(in1) in2 = np.asarray(in2) if in1.ndim == in2.ndim == 0: # scalar inputs return in1 * in2 elif not 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 np.array([]) s1 = np.array(in1.shape) s2 = np.array(in2.shape) complex_result = (np.issubdtype(in1.dtype, complex) or np.issubdtype(in2.dtype, complex)) shape = s1 + s2 - 1 # Check that input sizes are compatible with 'valid' mode if sig._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 = [sig.fftpack.helper.next_fast_len(int(d)) for d in shape] fslice = tuple([slice(0, int(sz)) for sz in shape]) # Pre-1.9 NumPy FFT routines are not threadsafe. For older NumPys, make # sure we only call rfftn/irfftn from one thread at a time. if not complex_result and (sig._rfft_mt_safe or sig._rfft_lock.acquire(False)): try: sp1 = rfftn(in1, fshape, threads=threads) sp2 = rfftn(in2, fshape, threads=threads) ret = (irfftn(sp1 * sp2, fshape, threads=threads)[fslice].copy()) finally: if not sig._rfft_mt_safe: sig._rfft_lock.release() else: # If we're here, it's either because we need a complex result, or we # failed to acquire _rfft_lock (meaning rfftn isn't threadsafe and # is already in use by another thread). In either case, use the # (threadsafe but slower) SciPy complex-FFT routines instead. sp1 = fftn(in1, fshape, threads=threads) sp2 = fftn(in2, fshape, threads=threads) ret = ifftn(sp1 * sp2, threads=threads)[fslice].copy() if not complex_result: ret = ret.real if mode == "full": return ret elif mode == "same": return sig._centered(ret, s1) elif mode == "valid": return sig._centered(ret, s1 - s2 + 1) else: raise ValueError("Acceptable mode flags are 'valid'," " 'same', or 'full'.")
def fftconvolvefast(in1, in2): s1 = np.array(in1.shape) s2 = np.array(in2.shape) shape = s1 + s2 - 1 fshape = [fftpack.helper.next_fast_len(int(d)) for d in shape] fslice = tuple([slice(0, int(sz)) for sz in shape]) sp1 = rfftn(in1, fshape) sp2 = rfftn(in2, fshape) ret = (irfftn(sp1 * sp2, fshape)[fslice].copy()) return _centered(ret, s1 - s2 + 1)
def fft_gaussian_filter(img, sigma): """FFT gaussian convolution Parameters ---------- img : ndarray Image to convolve with a gaussian kernel sigma : int or sequence The sigma(s) of the gaussian kernel in _real space_ Returns ------- filt_img : ndarray The filtered image """ # This doesn't help agreement but it will make things faster # pull the shape s1 = np.array(img.shape) # s2 = np.array([int(s * 4) for s in _normalize_sequence(sigma, img.ndim)]) shape = s1 # + s2 - 1 # calculate a nice shape fshape = [sig.fftpack.helper.next_fast_len(int(d)) for d in shape] # pad out with reflection pad_img = fft_pad(img, fshape, "reflect") # calculate the padding padding = tuple(_calc_pad(o, n) for o, n in zip(img.shape, pad_img.shape)) # so that we can calculate the cropping, maybe this should be integrated # into `fft_pad` ... fslice = tuple( slice(s, -e) if e != 0 else slice(s, None) for s, e in padding) # fourier transfrom and apply the filter kimg = rfftn(pad_img, fshape) filt_kimg = fourier_gaussian(kimg, sigma, pad_img.shape[-1]) # inverse FFT and return. return irfftn(filt_kimg, fshape)[fslice]
def rfftn(a, axes=None): if _use_fftw: from pyfftw.interfaces.numpy_fft import rfftn _init_pyfftw() return rfftn(a, axes=axes, **_fft_extra_args) else: return np.fft.rfftn(a, axes=axes)
def urdftn(inarray, ndim=None, *args, **kwargs): """N-dim real unitary discrete Fourier transform This transform consider the Hermitian property of the transform on real input Parameters ---------- inarray : ndarray The array to transform. ndim : int, optional The `ndim` last axis along wich to compute the transform. All axes by default. Returns ------- outarray : array-like (the last ndim as N / 2 + 1 lenght) """ if not ndim: ndim = inarray.ndim return rfftn(inarray, axes=range(-ndim, 0), *args, **kwargs) / np.sqrt( np.prod(inarray.shape[-ndim:]))
def fft_forward(self): """Compute the direct Fourier transform of the input data.""" logger.debug('Computing FFT of input with %s threads', self.nthreads) data = self.data_raw # Note: we don't use the defaults argument for the FFT as we # don't want any planning effort going to this one self.data_raw_f = fft.rfftn(data, threads=self.nthreads)
def vol_ft(vol, pfac=2, threads=1, normfft=1): """ Returns a centered, Nyquist-limited, zero-padded, interpolation-ready 3D Fourier transform. :param vol: Volume to be Fourier transformed. :param pfac: Size factor for zero-padding. :param threads: Number of threads for pyFFTW. :param normfft: Normalization constant for Fourier transform. """ vol = grid_correct(vol.astype(np.float64), pfac=pfac, order=1) padvol = np.pad(vol, (vol.shape[0] * pfac - vol.shape[0]) // 2, "constant") ft = rfftn(np.fft.ifftshift(padvol), padvol.shape, threads=threads) ftc = np.zeros((ft.shape[0] + 3, ft.shape[1] + 3, ft.shape[2]), dtype=np.complex128) fill_ft(ft, ftc, vol.shape[0], normfft=normfft) return ftc
def fftconvolve_fast(data, kernel, **kwargs): """A faster version of fft convolution In this case the kernel ifftshifted before FFT but the data is not. This can be done because the effect of fourier convolution is to "wrap" around the data edges so whether we ifftshift before FFT and then fftshift after it makes no difference so we can skip the step entirely. """ # TODO: add error checking like in the above and add functionality # for complex inputs. Also could add options for different types of # padding. dshape = np.array(data.shape) kshape = np.array(kernel.shape) # find maximum dimensions maxshape = np.max((dshape, kshape), 0) # calculate a nice shape fshape = [sig.fftpack.helper.next_fast_len(int(d)) for d in maxshape] # pad out with reflection pad_data = fft_pad(data, fshape, "reflect") # calculate padding padding = tuple( _calc_pad(o, n) for o, n in zip(data.shape, pad_data.shape)) # so that we can calculate the cropping, maybe this should be integrated # into `fft_pad` ... fslice = tuple( slice(s, -e) if e != 0 else slice(s, None) for s, e in padding) if kernel.shape != pad_data.shape: # its been assumed that the background of the kernel has already been # removed and that the kernel has already been centered kernel = fft_pad(kernel, pad_data.shape, mode='constant') k_kernel = rfftn(ifftshift(kernel), pad_data.shape, **kwargs) k_data = rfftn(pad_data, pad_data.shape, **kwargs) convolve_data = irfftn(k_kernel * k_data, pad_data.shape, **kwargs) # return data with same shape as original data return convolve_data[fslice]
def vol_ft(vol, pfac=2, threads=1, normfft=1): """ Returns a centered, Nyquist-limited, zero-padded, interpolation-ready 3D Fourier transform. :param vol: Volume to be Fourier transformed. :param pfac: Size factor for zero-padding. :param threads: Number of threads for pyFFTW. :param normfft: Normalization constant for Fourier transform. """ vol = grid_correct(vol, pfac=pfac, order=1) padvol = np.pad(vol, int((vol.shape[0] * pfac - vol.shape[0]) // 2), "constant") ft = rfftn(np.fft.ifftshift(padvol), padvol.shape, threads=threads) # ft = np.fft.fftn(np.fft.ifftshift(padvol), padvol.shape ) ftc = np.zeros((ft.shape[0] + 3, ft.shape[1] + 3, ft.shape[2]), dtype=ft.dtype) vop.fill_ft(ft, ftc, vol.shape[0], normfft=normfft) print( ftc.shape ) plot2dimage(np.log10(np.abs(ftc[:][:][230])).astype(np.float)) sys.exit() return ftc
def from_array(cls, corr_array, is_cyclic=True): """Create an instance with the given correlations. Parameters ---------- corr_array: array_like The correlation of the first element of the domain with each other element. is_cyclic: bool Returns ------- HomogeneousIsotropicCorrelation """ corr_array = asarray(corr_array) shape = corr_array.shape ndims = corr_array.ndim if is_cyclic: computational_shape = shape self = cls(shape, computational_shape) corr_fourier = self._fft(corr_array) else: computational_shape = tuple(2 * (dim - 1) for dim in shape) self = cls(shape, computational_shape) for axis in reversed(range(ndims)): corr_array = concatenate( [corr_array, flip(corr_array[1:-1], axis)], axis=axis) # Advantages over dctn: guaranteed same format and gets # nice planning done for the later evaluations. corr_fourier = rfftn(corr_array, axes=arange(0, ndims, dtype=int), threads=NUM_THREADS, planner_effort=ADVANCE_PLANNER_EFFORT) # The fft axes need to be a single chunk for the dask ffts # It's in memory already anyway # TODO: create a from_spectrum to delegate to self._corr_fourier = (corr_fourier) self._fourier_near_zero = (corr_fourier < FOURIER_NEAR_ZERO) return self
def hartley(a, axes=None): # Check if the axes provided are valid given the shape if axes is not None and \ not all(axis < len(a.shape) for axis in axes): raise ValueError("Provided axes do not match array shape") if iscomplextype(a.dtype): raise TypeError("Hartley transform requires real-valued arrays.") tmp = rfftn(a, axes=axes) def _fill_array(tmp, res, axes): if axes is None: axes = tuple(range(tmp.ndim)) lastaxis = axes[-1] ntmplast = tmp.shape[lastaxis] slice1 = (slice(None),)*lastaxis + (slice(0, ntmplast),) np.add(tmp.real, tmp.imag, out=res[slice1]) def _fill_upper_half(tmp, res, axes): lastaxis = axes[-1] nlast = res.shape[lastaxis] ntmplast = tmp.shape[lastaxis] nrem = nlast - ntmplast slice1 = [slice(None)]*lastaxis + [slice(ntmplast, None)] slice2 = [slice(None)]*lastaxis + [slice(nrem, 0, -1)] for i in axes[:-1]: slice1[i] = slice(1, None) slice2[i] = slice(None, 0, -1) slice1 = tuple(slice1) slice2 = tuple(slice2) np.subtract(tmp[slice2].real, tmp[slice2].imag, out=res[slice1]) for i, ax in enumerate(axes[:-1]): dim1 = (slice(None),)*ax + (slice(0, 1),) axes2 = axes[:i] + axes[i+1:] _fill_upper_half(tmp[dim1], res[dim1], axes2) _fill_upper_half(tmp, res, axes) return res return _fill_array(tmp, np.empty_like(a), axes)
def fft_gaussian_filter(img, sigma): """FFT gaussian convolution Parameters ---------- img : ndarray Image to convolve with a gaussian kernel sigma : int or sequence The sigma(s) of the gaussian kernel in _real space_ Returns ------- filt_img : ndarray The filtered image """ # This doesn't help agreement but it will make things faster # pull the shape s1 = np.array(img.shape) # if any of the sizes is 32 or smaller, revert to proper filter if any(s1 < 33): warnings.warn(("Input is small along a dimension," " will revert to `gaussian_filter`")) return gaussian_filter(img, sigma) # s2 = np.array([int(s * 4) for s in _normalize_sequence(sigma, img.ndim)]) shape = s1 # + s2 - 1 # calculate a nice shape fshape = [sig.fftpack.helper.next_fast_len(int(d)) for d in shape] # pad out with reflection pad_img = fft_pad(img, fshape, "reflect") # calculate the padding padding = tuple(_calc_pad(o, n) for o, n in zip(img.shape, pad_img.shape)) # so that we can calculate the cropping, maybe this should be integrated # into `fft_pad` ... fslice = tuple(slice(s, -e) if e != 0 else slice(s, None) for s, e in padding) # fourier transfrom and apply the filter kimg = rfftn(pad_img, fshape) filt_kimg = fourier_gaussian(kimg, sigma, pad_img.shape[-1]) # inverse FFT and return. return irfftn(filt_kimg, fshape)[fslice]
def my_fftn_r2c(a, axes=None): # Check if the axes provided are valid given the shape if axes is not None and \ not all(axis < len(a.shape) for axis in axes): raise ValueError("Provided axes do not match array shape") if iscomplextype(a.dtype): raise TypeError("Transform requires real-valued input arrays.") tmp = rfftn(a, axes=axes) def _fill_complex_array(tmp, res, axes): if axes is None: axes = tuple(range(tmp.ndim)) lastaxis = axes[-1] ntmplast = tmp.shape[lastaxis] slice1 = [slice(None)]*lastaxis + [slice(0, ntmplast)] res[tuple(slice1)] = tmp def _fill_upper_half_complex(tmp, res, axes): lastaxis = axes[-1] nlast = res.shape[lastaxis] ntmplast = tmp.shape[lastaxis] nrem = nlast - ntmplast slice1 = [slice(None)]*lastaxis + [slice(ntmplast, None)] slice2 = [slice(None)]*lastaxis + [slice(nrem, 0, -1)] for i in axes[:-1]: slice1[i] = slice(1, None) slice2[i] = slice(None, 0, -1) # np.conjugate(tmp[slice2], out=res[slice1]) res[tuple(slice1)] = np.conjugate(tmp[tuple(slice2)]) for i, ax in enumerate(axes[:-1]): dim1 = tuple([slice(None)]*ax + [slice(0, 1)]) axes2 = axes[:i] + axes[i+1:] _fill_upper_half_complex(tmp[dim1], res[dim1], axes2) _fill_upper_half_complex(tmp, res, axes) return res return _fill_complex_array(tmp, np.empty_like(a, dtype=tmp.dtype), axes)
def urfftn(inarray, dim=None): """N-dimensional real unitary Fourier transform. This transform considers the Hermitian property of the transform on real-valued input. Parameters ---------- inarray : ndarray, shape (M, N, ..., P) The array to transform. dim : int, optional The last axis along which to compute the transform. All axes by default. Returns ------- outarray : ndarray, shape (M, N, ..., P / 2 + 1) The unitary N-D real Fourier transform of ``inarray``. Notes ----- The ``urfft`` functions assume an input array of real values. Consequently, the output has a Hermitian property and redundant values are not computed or returned. Examples -------- >>> input = np.ones((5, 5, 5)) >>> output = urfftn(input) >>> np.allclose(np.sum(input) / np.sqrt(input.size), output[0, 0, 0]) True >>> output.shape (5, 5, 3) """ if dim is None: dim = inarray.ndim outarray = rfftn(inarray, axes=range(-dim, 0)) return outarray / np.sqrt(np.prod(inarray.shape[-dim:]))
def generate_white_noise(self): """Generate a white noise with the relevant power spectrum. Returns ------- white_noise : np.ndarray """ # Compute the k grid d = self.Lbox / self.dimensions / (2 * np.pi) all_k = [fft.fftfreq(self.dimensions, d=d)] * (self.Ndim - 1) + [ fft.rfftfreq(self.dimensions, d=d) ] self.kgrid = kgrid = np.array(np.meshgrid(*all_k, indexing="ij")) self.knorm = knorm = np.sqrt(np.sum(kgrid**2, axis=0)) # Compute Pk Pk = np.zeros_like(knorm) mask = knorm > 0 Pk[mask] = self.Pk(knorm[mask] * 2) # Compute white noise (in Fourier space) mu = np.random.standard_normal([self.dimensions] * self.Ndim) muk = fft.rfftn(mu) deltak = muk * np.sqrt(Pk) # Compute field in real space white_noise = fft.irfftn(deltak) # Normalize variance deltak_smoothed = deltak * self.filter.W(knorm) field = fft.irfftn(deltak_smoothed) std = field.std() self.white_noise_fft = deltak * self.sigma8 / std self.white_noise = white_noise * self.sigma8 / std return self.white_noise
def pattern_params(my_pat, size=2): """Find stuff""" # REAL FFT! # note the limited shifting, we don't want to shift the last axis my_pat_fft = fftshift(rfftn(ifftshift(my_pat)), axes=tuple(range(my_pat.ndim))[:-1]) my_abs_pat_fft = abs(my_pat_fft) # find dc loc, center of FFT after shifting sizeky, sizekx = my_abs_pat_fft.shape # remember we didn't shift the last axis! dc_loc = (sizeky // 2, 0) # mask data and find next biggest peak dc_power = my_abs_pat_fft[dc_loc] my_abs_pat_fft[dc_loc] = 0 max_loc = np.unravel_index(my_abs_pat_fft.argmax(), my_abs_pat_fft.shape) # pull the 3x3 region around the peak and fit max_shift = localize_peak(my_abs_pat_fft[slice_maker(max_loc, 3)]) # calculate precise peak relative to dc peak = np.array(max_loc) + np.array(max_shift) - np.array(dc_loc) # correct location based on initial data shape peak_corr = peak / np.array(my_pat.shape) # calc angle preciseangle = np.arctan2(*peak_corr) # calc period precise_period = 1 / norm(peak_corr) # calc phase phase = np.angle(my_pat_fft[max_loc[0], max_loc[1]]) # calc modulation depth numerator = abs(my_pat_fft[slice_maker(max_loc, size)].sum()) mod = numerator / dc_power return {"period": precise_period, "angle": preciseangle, "phase": phase, "fft": my_pat_fft, "mod": mod, "max_loc": max_loc}
def rfft_to_fft(image, keep_rfft=False, use_pyfftw=False, threads=1, **pyfftw_kwargs): ''' Perform a RFFT on the image (2 or 3D) and return the absolute value in the same format as you would get with the fft (negative frequencies). This avoids ever having to have the full complex cube in memory. Inputs ------ image : numpy.ndarray 2 or 3D array. keep_rfft : bool, optional Return the rfft output instead of expanding to the negative frequencies of the full FFT. use_pyfftw : bool, optional Try using pyfftw for the FFT. threads : int, optional Number of threads to use when using pyfftw. Default is 1. pyfftw_kwargs : Passed to `~pyfftw.interfaces.numpy_fft.rfftn`. Outputs ------- fft_abs : absolute value of the fft. ''' ndim = len(image.shape) if ndim < 2 or ndim > 3: raise TypeError("Dimension of image must be 2D or 3D.") last_dim = image.shape[-1] if use_pyfftw: if PYFFTW_FLAG: fft_abs = np.abs(rfftn(image, **pyfftw_kwargs)) else: use_pyfftw = False warn("pyfftw is not installed") if not use_pyfftw: fft_abs = np.abs(np.fft.rfftn(image)) if keep_rfft: return fft_abs if ndim == 2: if last_dim % 2 == 0: fftstar_abs = fft_abs.copy()[:, -2:0:-1] else: fftstar_abs = fft_abs.copy()[:, -1:0:-1] fftstar_abs[1::, :] = fftstar_abs[:0:-1, :] return np.concatenate((fft_abs, fftstar_abs), axis=1) elif ndim == 3: if last_dim % 2 == 0: fftstar_abs = fft_abs.copy()[:, :, -2:0:-1] else: fftstar_abs = fft_abs.copy()[:, :, -1:0:-1] fftstar_abs[1::, :, :] = fftstar_abs[:0:-1, :, :] fftstar_abs[:, 1::, :] = fftstar_abs[:, :0:-1, :] return np.concatenate((fft_abs, fftstar_abs), axis=2)
def ir2tf(imp_resp, shape, dim=None, is_real=True): """Compute the transfer function of an impulse response (IR). This function makes the necessary correct zero-padding, zero convention, correct fft2, etc... to compute the transfer function of IR. To use with unitary Fourier transform for the signal (ufftn or equivalent). Parameters ---------- imp_resp : ndarray The impulse responses. shape : tuple of int A tuple of integer corresponding to the target shape of the transfer function. dim : int, optional The last axis along which to compute the transform. All axes by default. is_real : boolean (optional, default True) If True, imp_resp is supposed real and the Hermitian property is used with rfftn Fourier transform. Returns ------- y : complex ndarray The transfer function of shape ``shape``. See Also -------- ufftn, uifftn, urfftn, uirfftn Examples -------- >>> np.all(np.array([[4, 0], [0, 0]]) == ir2tf(np.ones((2, 2)), (2, 2))) True >>> ir2tf(np.ones((2, 2)), (512, 512)).shape == (512, 257) True >>> ir2tf(np.ones((2, 2)), (512, 512), is_real=False).shape == (512, 512) True Notes ----- The input array can be composed of multiple-dimensional IR with an arbitrary number of IR. The individual IR must be accessed through the first axes. The last ``dim`` axes contain the space definition. """ if not dim: dim = imp_resp.ndim # Zero padding and fill irpadded = np.zeros(shape) irpadded[tuple([slice(0, s) for s in imp_resp.shape])] = imp_resp # Roll for zero convention of the fft to avoid the phase # problem. Work with odd and even size. for axis, axis_size in enumerate(imp_resp.shape): if axis >= imp_resp.ndim - dim: irpadded = np.roll(irpadded, shift=-int(np.floor(axis_size / 2)), axis=axis) if is_real: return rfftn(irpadded, axes=range(-dim, 0)) else: return fftn(irpadded, axes=range(-dim, 0))
def from_function(cls, corr_func, shape, is_cyclic=True): """Create an instance to apply the correlation function. Parameters ---------- corr_func: callable(dist) -> float The correlation of the first element of the domain with each other element. shape: tuple of int The state is formally a vector, but the correlations are assumed to depend on the layout in some other shape, usually related to the physical layout. This is the other shape. is_cyclic: bool Whether to assume the domain is periodic in all directions. Returns ------- HomogeneousIsotropicCorrelation """ shape = np.atleast_1d(shape) if is_cyclic: computational_shape = tuple(shape) else: computational_shape = tuple( next_fast_len(2 * dim - 1) for dim in shape) self = cls(tuple(shape), computational_shape) shape = np.asarray(self._computational_shape) ndims = len(shape) broadcastable_shape = shape[:, newaxis] while broadcastable_shape.ndim < ndims + 1: broadcastable_shape = broadcastable_shape[..., newaxis] def corr_from_index(*index): """Correlation of index with zero. Turns a correlation function in terms of index distance into one in terms of indices on a periodic domain. Parameters ---------- index: tuple of int Returns ------- float[-1, 1] The correlation of the given index with the origin. See Also -------- DistanceCorrelationFunction.correlation_from_index """ comp2_1 = square(index) # Components of distance to shifted origin comp2_2 = square(broadcastable_shape - index) # use the smaller components to get the distance to the # closest of the shifted origins comp2 = fmin(comp2_1, comp2_2) return corr_func(sqrt(array_sum(comp2, axis=0))) corr_struct = fromfunction(corr_from_index, shape=tuple(shape), dtype=DTYPE) # I should be able to generate this sequence with a type-I DCT # For some odd reason complex/complex is faster than complex/real # This also ensures the format here is the same as in _matmat corr_fourier = rfftn(corr_struct, axes=arange(ndims, dtype=int), threads=NUM_THREADS, planner_effort=ADVANCE_PLANNER_EFFORT) self._corr_fourier = (corr_fourier) # This is also affected by roundoff abs_corr_fourier = abs(corr_fourier) self._fourier_near_zero = (abs_corr_fourier < FOURIER_NEAR_ZERO * abs_corr_fourier.max()) return self
def ir2fr(imp_resp, shape, center=None, real=True): """Return the frequency response from impulsionnal responses This function make the necessary correct zero-padding, zero convention, correct DFT etc. to compute the frequency response from impulsionnal responses (IR). The IR array is supposed to have the origin in the middle of the array. The Fourier transform is performed on the last `len(shape)` dimensions. Parameters ---------- imp_resp : ndarray The impulsionnal responses. shape : tuple of int A tuple of integer corresponding to the target shape of the frequency responses, without hermitian property. center : tuple of int, optional The origin index of the impulsionnal response. The middle by default. real : boolean (optionnal, default True) If True, imp_resp is supposed real, the hermissian property is used with rfftn DFT and the output has `shape[-1] / 2 + 1` elements on the last axis. Returns ------- y : ndarray The frequency responses of shape `shape` on the last `len(shape)` dimensions. Notes ----- - For convolution, the result have to be used with unitary discrete Fourier transform for the signal (udftn or equivalent). - DFT are always peformed on last axis for efficiency. - Results is always C-contiguous. See Also -------- udftn, uidftn, urdftn, uirdftn """ if len(shape) > imp_resp.ndim: raise ValueError("length of shape must inferior to imp_resp.ndim") if not center: center = [int(np.floor(length / 2)) for length in imp_resp.shape] if len(center) != len(shape): raise ValueError("center and shape must have the same length") # Place the provided IR at the beginning of the array irpadded = np.zeros(shape) irpadded[tuple([slice(0, s) for s in imp_resp.shape])] = imp_resp # Roll, or circshift to place the origin at 0 index, the # hypothesis of the DFT for axe, shift in enumerate(center): irpadded = np.roll(irpadded, -shift, imp_resp.ndim - len(shape) + axe) # Perform the DFT on the last axes if real: return np.ascontiguousarray( rfftn(irpadded, axes=list(range(imp_resp.ndim - len(shape), imp_resp.ndim)))) else: return np.ascontiguousarray( fftn(irpadded, axes=list(range(imp_resp.ndim - len(shape), imp_resp.ndim))))
def fft(*args, **kwargs): return fftw.rfftn(*args, **kwargs, threads=2)
def drift_plot(fit, title=None, dt=0.1, dx=130, lf=-np.inf, hf=np.inf, log=False, cmap='magma', xc='b', yc='r'): ''' Plotting utility to show drift curves nicely Parameters ---------- fit : pandas DataFrame Assumes that it has attributes x0 and y0 title : str (optional) Title of plot dt : float (optional) Sampling rate of data in seconds dx : pixel size (optional) Pixel size in nm lf : float (optional) Low frequency cutoff for fourier plot hf : float (optional) High frequency cutoff for fourier plot log : bool (optional) Take logarithm of FFT data before displaying cmap : string or matplotlib.colors.cmap instance Color map for scatter plot xc : string or `color` instance Color for x data yc : string or `color` instance Color for y data Returns ------- fig : figure object The figure axs : tuple of axes objects In the following order, Real axis, FFT axis, Scatter axis ''' # set up plot fig = plt.figure() fig.set_size_inches(8, 4) axreal = plt.subplot(221) axfft = plt.subplot(223) axscatter = plt.subplot(122) # label it if title is not None: fig.suptitle(title, y=1.02, fontweight='bold') # detrend mean ybar = fit.y0 - fit.y0.mean() ybar *= dx xbar = fit.x0 - fit.x0.mean() xbar *= dx # Plot Real space t = np.arange(len(fit)) * dt axreal.plot(t, xbar, xc, label=r"$x_0$") axreal.plot(t, ybar, yc, label=r"$y_0$") axreal.set_xlabel('Time (s)') axreal.set_ylabel('Displacement (nm)') # add legend to real axis axreal.legend(loc='best') # calc FFTs Y = rfftn(ybar) X = rfftn(xbar) # calc FFT freq k = rfftfreq(len(fit), dt) # limit FFT display range kg = np.logical_and(k > lf, k < hf) # Plot FFT if log: axfft.semilogy(k[kg], abs(X[kg]), xc) axfft.semilogy(k[kg], abs(Y[kg]), yc) else: axfft.plot(k[kg], abs(X[kg]), xc) axfft.plot(k[kg], abs(Y[kg]), yc) axfft.set_xlabel('Frequency (Hz)') # Plot scatter axscatter.scatter(xbar, ybar, c=t, cmap=cmap) axscatter.set_xlabel('x') axscatter.set_ylabel('y') # make sure the scatter plot is square lims = axreal.get_ylim() axscatter.set_ylim(lims) axscatter.set_xlim(lims) # tight layout axs = (axreal, axfft, axscatter) fig.tight_layout() # return fig, axs to user for further manipulation and/or saving if wanted return fig, axs
def cross_correlation_3d(pixels1, pixels2): '''Align the second image with the first using max cross-correlation returns the z,y,x offsets to add to image1's indexes to align it with image2 Many of the ideas here are based on the paper, "Fast Normalized Cross-Correlation" by J.P. Lewis (http://www.idiom.com/~zilla/Papers/nvisionInterface/nip.html) which is frequently cited when addressing this problem. ''' s = np.maximum(pixels1.shape, pixels2.shape) fshape = s*2 # # Calculate the # of pixels at a particular point # i,j,k = np.mgrid[-s[0]:s[0], -s[1]:s[1], -s[2]:s[2] ] unit = np.abs(i*j*k).astype(float) unit[unit<1]=1 # keeps from dividing by zero in some places # # Normalize the pixel values around zero which does not affect the # correlation, keeps some of the sums of multiplications from # losing precision and precomputes t(x-u,y-v) - t_mean # pixels1 = np.nan_to_num(pixels1-nanmean(pixels1)) pixels2 = np.nan_to_num(pixels2-nanmean(pixels2)) # # Lewis uses an image, f and a template t. He derives a normalized # cross correlation, ncc(u,v) = # sum((f(x,y)-f_mean(u,v))*(t(x-u,y-v)-t_mean),x,y) / # sqrt(sum((f(x,y)-f_mean(u,v))**2,x,y) * (sum((t(x-u,y-v)-t_mean)**2,x,y) # # From here, he finds that the numerator term, f_mean(u,v)*(t...) is zero # leaving f(x,y)*(t(x-u,y-v)-t_mean) which is a convolution of f # by t-t_mean. # fp1 = rfftn(pixels1.astype('float32'), fshape, axes=(0, 1, 2)) fp2 = rfftn(pixels2.astype('float32'), fshape, axes=(0, 1, 2)) corr12 = irfftn(fp1 * fp2.conj(), axes=(0, 1, 2)).real # # Use the trick of Lewis here - compute the cumulative sums # in a fashion that accounts for the parts that are off the # edge of the template. # # We do this in quadrants: # q0 q1 # q2 q3 # For the first, # q0 is the sum over pixels1[i:,j:] - sum i,j backwards # q1 is the sum over pixels1[i:,:j] - sum i backwards, j forwards # q2 is the sum over pixels1[:i,j:] - sum i forwards, j backwards # q3 is the sum over pixels1[:i,:j] - sum i,j forwards # # The second is done as above but reflected lr and ud # def get_cumsums(im, fshape): im_si = im.shape[0] im_sj = im.shape[1] im_sk = im.shape[2] im_sum = np.zeros(fshape) im_sum[:im_si,:im_sj,:im_sk] = cumsum_quadrant(im, False, False, False) im_sum[:im_si,:im_sj,-im_sk:] = cumsum_quadrant(im, False, False, True) im_sum[:im_si,-im_sj:,:im_sk] = cumsum_quadrant(im, False, True, True) im_sum[:im_si,-im_sj:,-im_sk:] = cumsum_quadrant(im, False, True, False) im_sum[-im_si:,:im_sj,:im_sk] = cumsum_quadrant(im, True, False, True) im_sum[-im_si:,:im_sj,-im_sk:] = cumsum_quadrant(im, True, False, False) im_sum[-im_si:,-im_sj:,:im_sk] = cumsum_quadrant(im, True, True, True) im_sum[-im_si:,-im_sj:,-im_sk:] = cumsum_quadrant(im, True, True, False) # # Divide the sum over the # of elements summed-over # return im_sum / unit p1_mean = get_cumsums(pixels1, fshape) p2_mean = get_cumsums(pixels2, fshape) # # Once we have the means for u,v, we can caluclate the # variance-like parts of the equation. We have to multiply # the mean^2 by the # of elements being summed-over # to account for the mean being summed that many times. # p1sd = np.sum(pixels1**2) - p1_mean**2 * np.product(s) p2sd = np.sum(pixels2**2) - p2_mean**2 * np.product(s) # # There's always chance of roundoff error for a zero value # resulting in a negative sd, so limit the sds here # sd = np.sqrt(np.maximum(p1sd * p2sd, 0)) with warnings.catch_warnings(): warnings.simplefilter("ignore") corrnorm = corr12 / sd # # There's not much information for points where the standard # deviation is less than 1/100 of the maximum. We exclude these # from consideration. # corrnorm[(unit < np.product(s) / 2) & (sd < np.mean(sd) / 100)] = 0 # Also exclude possibilites with few observed pixels. corrnorm[unit < np.product(s) / 4] = 0 return corrnorm
def cross_correlation_3d(pixels1, pixels2): '''Align the second image with the first using max cross-correlation returns the z,y,x offsets to add to image1's indexes to align it with image2 Many of the ideas here are based on the paper, "Fast Normalized Cross-Correlation" by J.P. Lewis (http://www.idiom.com/~zilla/Papers/nvisionInterface/nip.html) which is frequently cited when addressing this problem. ''' s = np.maximum(pixels1.shape, pixels2.shape) fshape = s*2 # # Calculate the # of pixels at a particular point # i, j, k = np.mgrid[-s[0]:s[0], -s[1]:s[1], -s[2]:s[2] ] unit = np.abs(i*j*k).astype(float) unit[unit < 1] = 1 # keeps from dividing by zero in some places # # Normalize the pixel values around zero which does not affect the # correlation, keeps some of the sums of multiplications from # losing precision and precomputes t(x-u,y-v) - t_mean # pixels1 = np.nan_to_num(pixels1-nanmean(pixels1)) pixels2 = np.nan_to_num(pixels2-nanmean(pixels2)) # # Lewis uses an image, f and a template t. He derives a normalized # cross correlation, ncc(u,v) = # sum((f(x,y)-f_mean(u,v))*(t(x-u,y-v)-t_mean),x,y) / # sqrt(sum((f(x,y)-f_mean(u,v))**2,x,y) * (sum((t(x-u,y-v)-t_mean)**2,x,y) # # From here, he finds that the numerator term, f_mean(u,v)*(t...) is zero # leaving f(x,y)*(t(x-u,y-v)-t_mean) which is a convolution of f # by t-t_mean. # fp1 = rfftn(pixels1.astype('float32'), fshape, axes=(0, 1, 2)) fp2 = rfftn(pixels2.astype('float32'), fshape, axes=(0, 1, 2)) corr12 = irfftn(fp1 * fp2.conj(), axes=(0, 1, 2)).real # # Use the trick of Lewis here - compute the cumulative sums # in a fashion that accounts for the parts that are off the # edge of the template. # # We do this in quadrants: # q0 q1 # q2 q3 # For the first, # q0 is the sum over pixels1[i:,j:] - sum i,j backwards # q1 is the sum over pixels1[i:,:j] - sum i backwards, j forwards # q2 is the sum over pixels1[:i,j:] - sum i forwards, j backwards # q3 is the sum over pixels1[:i,:j] - sum i,j forwards # # The second is done as above but reflected lr and ud # def get_cumsums(im, fshape): im_si = im.shape[0] im_sj = im.shape[1] im_sk = im.shape[2] im_sum = np.zeros(fshape) im_sum[:im_si, :im_sj, :im_sk] = cumsum_quadrant(im, False, False, False) im_sum[:im_si, :im_sj, -im_sk:] = cumsum_quadrant(im, False, False, True) im_sum[:im_si, -im_sj:, :im_sk] = cumsum_quadrant(im, False, True, True) im_sum[:im_si, -im_sj:, -im_sk:] = cumsum_quadrant(im, False, True, False) im_sum[-im_si:, :im_sj, :im_sk] = cumsum_quadrant(im, True, False, True) im_sum[-im_si:, :im_sj, -im_sk:] = cumsum_quadrant(im, True, False, False) im_sum[-im_si:, -im_sj:, :im_sk] = cumsum_quadrant(im, True, True, True) im_sum[-im_si:, -im_sj:, -im_sk:] = cumsum_quadrant(im, True, True, False) # # Divide the sum over the # of elements summed-over # return old_div(im_sum, unit) p1_mean = get_cumsums(pixels1, fshape) p2_mean = get_cumsums(pixels2, fshape) # # Once we have the means for u,v, we can caluclate the # variance-like parts of the equation. We have to multiply # the mean^2 by the # of elements being summed-over # to account for the mean being summed that many times. # p1sd = np.sum(pixels1**2) - p1_mean**2 * np.product(s) p2sd = np.sum(pixels2**2) - p2_mean**2 * np.product(s) # # There's always chance of roundoff error for a zero value # resulting in a negative sd, so limit the sds here # sd = np.sqrt(np.maximum(p1sd * p2sd, 0)) with warnings.catch_warnings(): warnings.simplefilter("ignore") corrnorm = old_div(corr12, sd) # # There's not much information for points where the standard # deviation is less than 1/100 of the maximum. We exclude these # from consideration. # corrnorm[(unit < old_div(np.product(s), 2)) & (sd < old_div(np.mean(sd), 100))] = 0 # Also exclude possibilites with few observed pixels. corrnorm[unit < old_div(np.product(s), 4)] = 0 return corrnorm
def fft(self, scalar): """ Performs fft of scalar field """ self.scalar_input_test(scalar) return fft.rfftn(scalar, threads=self.num_threads)
def fft(self, vector): """ Performs fft of vector field """ self.vector_input_test(vector) return fft.rfftn(vector, axes=(1, 2), threads=self.num_threads)