python类fft2()的实例源码

def _shift_fft(array, shift_value):
    Ndim  = array.ndim
    dims  = array.shape
    dtype = array.dtype.kind

    if (dtype != 'f'):
        raise ValueError('Array must be float')

    shifted = array
    if (Ndim == 1):
        Nx = dims[0]

        x_ramp = np.arange(Nx, dtype=array.dtype) - Nx//2

        tilt = (2*np.pi/Nx) * (shift_value[0]*x_ramp)

        cplx_tilt = np.cos(tilt) + 1j*np.sin(tilt)
        cplx_tilt = fft.fftshift(cplx_tilt)
        narray    = fft.fft(fft.ifft(array) * cplx_tilt)
        shifted   = narray.real
    elif (Ndim == 2):
        Nx = dims[0]
        Ny = dims[1]

        x_ramp = np.outer(np.full(Nx, 1.), np.arange(Ny, dtype=array.dtype)) - Nx//2
        y_ramp = np.outer(np.arange(Nx, dtype=array.dtype), np.full(Ny, 1.)) - Ny//2

        tilt = (2*np.pi/Nx) * (shift_value[0]*x_ramp+shift_value[1]*y_ramp)

        cplx_tilt = np.cos(tilt) + 1j*np.sin(tilt)        
        cplx_tilt = fft.fftshift(cplx_tilt)

        narray    = fft.fft2(fft.ifft2(array) * cplx_tilt)
        shifted   = narray.real
    else:
        raise ValueError('This function can shift only 1D or 2D arrays')

    return shifted

def op(self, img):
        """ This method calculates the masked Fourier transform of a 2-D image.

        Parameters
        ----------
        img: np.ndarray
            input 2D array with the same shape as the mask.

        Returns
        -------
        x: np.ndarray
            masked Fourier transform of the input image.
        """
        return self._mask * pfft.fft2(img)

def filters_bank(M, N, J, L=8):
    filters = {}
    filters['psi'] = []


    offset_unpad = 0
    for j in range(J):
        for theta in range(L):
            psi = {}
            psi['j'] = j
            psi['theta'] = theta
            psi_signal = morlet_2d(M, N, 0.8 * 2**j, (int(L-L/2-1)-theta) * np.pi / L, 3.0 / 4.0 * np.pi /2**j,offset=offset_unpad) # The 5 is here just to match the LUA implementation :)
            psi_signal_fourier = fft.fft2(psi_signal)
            for res in range(j + 1):
                psi_signal_fourier_res = crop_freq(psi_signal_fourier, res)
                psi[res]=torch.FloatTensor(np.stack((np.real(psi_signal_fourier_res), np.imag(psi_signal_fourier_res)), axis=2))
                # Normalization to avoid doing it with the FFT!
                psi[res].div_(M*N// 2**(2*j))
            filters['psi'].append(psi)

    filters['phi'] = {}
    phi_signal = gabor_2d(M, N, 0.8 * 2**(J-1), 0, 0, offset=offset_unpad)
    phi_signal_fourier = fft.fft2(phi_signal)
    filters['phi']['j'] = J
    for res in range(J):
        phi_signal_fourier_res = crop_freq(phi_signal_fourier, res)
        filters['phi'][res]=torch.FloatTensor(np.stack((np.real(phi_signal_fourier_res), np.imag(phi_signal_fourier_res)), axis=2))
        filters['phi'][res].div_(M*N // 2 ** (2 * J))

    return filters

def fourier(self):
        F1 = fftpack.fft2(self.image)
        F2 = fftpack.fftshift(F1)
        return F2