Python numpy.fft() Examples

def _freq(N, delta_x, real, shift):
    # calculate frequencies from coordinates
    # coordinates are always loaded eagerly, so we use numpy
    if real is None:
        fftfreq = [np.fft.fftfreq]*len(N)
    else:
        # Discard negative frequencies from transform along last axis to be
        # consistent with np.fft.rfftn
        fftfreq = [np.fft.fftfreq]*(len(N)-1)
        fftfreq.append(np.fft.rfftfreq)

    k = [fftfreq(Nx, dx) for (fftfreq, Nx, dx) in zip(fftfreq, N, delta_x)]

    if shift:
        k = [np.fft.fftshift(l) for l in k]

    return k 

def make_node(self, frames, n, axis):
        """
        Compute an n-point fft of frames along given axis.

        """
        _frames = tensor.as_tensor(frames, ndim=2)
        _n = tensor.as_tensor(n, ndim=0)
        _axis = tensor.as_tensor(axis, ndim=0)
        if self.half and _frames.type.dtype.startswith('complex'):
            raise TypeError('Argument to HalfFFT must not be complex', frames)
        spectrogram = tensor.zmatrix()
        buf = generic()
        # The `buf` output is present for future work
        # when we call FFTW directly and re-use the 'plan' that FFTW creates.
        # In that case, buf would store a CObject encapsulating the plan.
        rval = Apply(self, [_frames, _n, _axis], [spectrogram, buf])
        return rval 

def perform(self, node, inp, out):
        frames, n, axis = inp
        spectrogram, buf = out
        if self.inverse:
            fft_fn = numpy.fft.ifft
        else:
            fft_fn = numpy.fft.fft

        fft = fft_fn(frames, int(n), int(axis))
        if self.half:
            M, N = fft.shape
            if axis == 0:
                if (M % 2):
                    raise ValueError(
                        'halfFFT on odd-length vectors is undefined')
                spectrogram[0] = fft[0:M / 2, :]
            elif axis == 1:
                if (N % 2):
                    raise ValueError(
                        'halfFFT on odd-length vectors is undefined')
                spectrogram[0] = fft[:, 0:N / 2]
            else:
                raise NotImplementedError()
        else:
            spectrogram[0] = fft 

def test_size_accuracy(self):
        # Sanity check for the accuracy for prime and non-prime sized inputs
        if self.rdt == np.float32:
            rtol = 1e-5
        elif self.rdt == np.float64:
            rtol = 1e-10

        for size in LARGE_COMPOSITE_SIZES + LARGE_PRIME_SIZES:
            np.random.seed(1234)
            x = np.random.rand(size).astype(self.rdt)
            y = ifft(fft(x))
            self.assertTrue(np.linalg.norm(x - y) < rtol*np.linalg.norm(x),
                            (size, self.rdt))
            y = fft(ifft(x))
            self.assertTrue(np.linalg.norm(x - y) < rtol*np.linalg.norm(x),
                            (size, self.rdt))

            x = (x + 1j*np.random.rand(size)).astype(self.cdt)
            y = ifft(fft(x))
            self.assertTrue(np.linalg.norm(x - y) < rtol*np.linalg.norm(x),
                            (size, self.rdt))
            y = fft(ifft(x))
            self.assertTrue(np.linalg.norm(x - y) < rtol*np.linalg.norm(x),
                            (size, self.rdt)) 

def test_fft():    
    b = numpy.random.random((3,3))
    a = afnumpy.array(b)
    fassert(afnumpy.fft.fft(a), numpy.fft.fft(b))

    b = numpy.random.random((3,2))
    a = afnumpy.array(b)
    fassert(afnumpy.fft.fft(a), numpy.fft.fft(b))

    b = numpy.random.random((5,3,2))
    a = afnumpy.array(b)
    fassert(afnumpy.fft.fft(a), numpy.fft.fft(b))

    b = numpy.random.random((5,3,2))+numpy.random.random((5,3,2))*1.0j
    a = afnumpy.array(b)
    fassert(afnumpy.fft.fft(a), numpy.fft.fft(b)) 

def test_ifft():    
    # Real to complex inverse fft not implemented in arrayfire
    # b = numpy.random.random((3,3))
    # a = afnumpy.array(b)
    # fassert(afnumpy.fft.ifft(a), numpy.fft.ifft(b))

    # b = numpy.random.random((3,2))
    # a = afnumpy.array(b)
    # fassert(afnumpy.fft.ifft(a), numpy.fft.ifft(b))

    # b = numpy.random.random((5,3,2))
    # a = afnumpy.array(b)
    # fassert(afnumpy.fft.ifft(a), numpy.fft.ifft(b))

    b = numpy.random.random((5,3,2))+numpy.random.random((5,3,2))*1.0j
#    b = numpy.ones((3,3))+numpy.zeros((3,3))*1.0j
    a = afnumpy.array(b)
    fassert(afnumpy.fft.ifft(a), numpy.fft.ifft(b)) 

def test_fft2():    
    b = numpy.random.random((3,3))
    a = afnumpy.array(b)
    fassert(afnumpy.fft.fft2(a), numpy.fft.fft2(b))

    b = numpy.random.random((3,2))
    a = afnumpy.array(b)
    fassert(afnumpy.fft.fft2(a), numpy.fft.fft2(b))

    b = numpy.random.random((5,3,2))
    a = afnumpy.array(b)
    fassert(afnumpy.fft.fft2(a), numpy.fft.fft2(b))

    b = numpy.random.random((5,3,2))+numpy.random.random((5,3,2))*1.0j
    a = afnumpy.array(b)
    fassert(afnumpy.fft.fft2(a), numpy.fft.fft2(b)) 

def test_ifft2():    
    # Real to complex inverse fft not implemented in arrayfire
    # b = numpy.random.random((3,3))
    # a = afnumpy.array(b)
    # fassert(afnumpy.fft.ifft2(a), numpy.fft.ifft2(b))

    # b = numpy.random.random((3,2))
    # a = afnumpy.array(b)
    # fassert(afnumpy.fft.ifft2(a), numpy.fft.ifft2(b))

    # b = numpy.random.random((5,3,2))
    # a = afnumpy.array(b)
    # fassert(afnumpy.fft.ifft2(a), numpy.fft.ifft2(b))

    b = numpy.random.random((5,3,2))+numpy.random.random((5,3,2))*1.0j
#    b = numpy.ones((3,3))+numpy.zeros((3,3))*1.0j
    a = afnumpy.array(b)
    fassert(afnumpy.fft.ifft2(a), numpy.fft.ifft2(b)) 

def test_shape_axes_argument(self):
        small_x = [[1,2,3],[4,5,6],[7,8,9]]
        large_x1 = array([[1,2,3,0],
                                  [4,5,6,0],
                                  [7,8,9,0],
                                  [0,0,0,0]])
        # Disable tests with shape and axes of different lengths
        # y = fftn(small_x,shape=(4,4),axes=(-1,))
        # for i in range(4):
        #    assert_array_almost_equal (y[i],fft(large_x1[i]))
        # y = fftn(small_x,shape=(4,4),axes=(-2,))
        # for i in range(4):
        #    assert_array_almost_equal (y[:,i],fft(large_x1[:,i]))
        y = fftn(small_x,shape=(4,4),axes=(-2,-1))
        assert_array_almost_equal(y,fftn(large_x1))
        y = fftn(small_x,shape=(4,4),axes=(-1,-2))
        assert_array_almost_equal(y,swapaxes(
            fftn(swapaxes(large_x1,-1,-2)),-1,-2)) 

def _fft_module(da):
    if da.chunks:
        return dsar.fft
    else:
        return np.fft 

def isotropize(ps, fftdim, nfactor=4):
    """
    Isotropize a 2D power spectrum or cross spectrum
    by taking an azimuthal average.

    .. math::
        \text{iso}_{ps} = k_r N^{-1} \sum_{N} |\mathbb{F}(da')|^2

    where :math:`N` is the number of azimuthal bins.

    Parameters
    ----------
    ps : `xarray.DataArray`
        The power spectrum or cross spectrum to be isotropized.
    fftdim : list
        The fft dimensions overwhich the isotropization must be performed.
    nfactor : int, optional
        Ratio of number of bins to take the azimuthal averaging with the
        data size. Default is 4.
    """

    # compute radial wavenumber bins
    k = ps[fftdim[1]]
    l = ps[fftdim[0]]
    N = [k.size, l.size]
    ki, kr = _radial_wvnum(k, l, min(N), nfactor)

    # average azimuthally
    ps = ps.assign_coords(freq_r=np.sqrt(k**2+l**2))
    iso_ps = (ps.groupby_bins('freq_r', bins=ki, labels=kr).mean()
              .rename({'freq_r_bins': 'freq_r'})
             )
    return iso_ps * iso_ps.freq_r 

def _fftn_blas(f, mesh):
    Gx = np.fft.fftfreq(mesh[0])
    Gy = np.fft.fftfreq(mesh[1])
    Gz = np.fft.fftfreq(mesh[2])
    expRGx = np.exp(np.einsum('x,k->xk', -2j*np.pi*np.arange(mesh[0]), Gx))
    expRGy = np.exp(np.einsum('x,k->xk', -2j*np.pi*np.arange(mesh[1]), Gy))
    expRGz = np.exp(np.einsum('x,k->xk', -2j*np.pi*np.arange(mesh[2]), Gz))
    out = np.empty(f.shape, dtype=np.complex128)
    buf = np.empty(mesh, dtype=np.complex128)
    for i, fi in enumerate(f):
        buf[:] = fi.reshape(mesh)
        g = lib.dot(buf.reshape(mesh[0],-1).T, expRGx, c=out[i].reshape(-1,mesh[0]))
        g = lib.dot(g.reshape(mesh[1],-1).T, expRGy, c=buf.reshape(-1,mesh[1]))
        g = lib.dot(g.reshape(mesh[2],-1).T, expRGz, c=out[i].reshape(-1,mesh[2]))
    return out.reshape(-1, *mesh) 

def _ifftn_blas(g, mesh):
    Gx = np.fft.fftfreq(mesh[0])
    Gy = np.fft.fftfreq(mesh[1])
    Gz = np.fft.fftfreq(mesh[2])
    expRGx = np.exp(np.einsum('x,k->xk', 2j*np.pi*np.arange(mesh[0]), Gx))
    expRGy = np.exp(np.einsum('x,k->xk', 2j*np.pi*np.arange(mesh[1]), Gy))
    expRGz = np.exp(np.einsum('x,k->xk', 2j*np.pi*np.arange(mesh[2]), Gz))
    out = np.empty(g.shape, dtype=np.complex128)
    buf = np.empty(mesh, dtype=np.complex128)
    for i, gi in enumerate(g):
        buf[:] = gi.reshape(mesh)
        f = lib.dot(buf.reshape(mesh[0],-1).T, expRGx, 1./mesh[0], c=out[i].reshape(-1,mesh[0]))
        f = lib.dot(f.reshape(mesh[1],-1).T, expRGy, 1./mesh[1], c=buf.reshape(-1,mesh[1]))
        f = lib.dot(f.reshape(mesh[2],-1).T, expRGz, 1./mesh[2], c=out[i].reshape(-1,mesh[2]))
    return out.reshape(-1, *mesh) 

def _fftn_wrapper(a):
            return np.fft.fftn(a, axes=(1,2,3)) 

def _ifftn_wrapper(a):
            return np.fft.ifftn(a, axes=(1,2,3)) 

def _fftn_wrapper(a):
        return np.fft.fftn(a, axes=(1,2,3)) 

def _fftn_wrapper(a):
        mesh = a.shape[1:]
        if mesh[0] in _EXCLUDE and mesh[1] in _EXCLUDE and mesh[2] in _EXCLUDE:
            return _fftn_blas(a, mesh)
        else:
            return np.fft.fftn(a, axes=(1,2,3)) 

def _ifftn_wrapper(a):
        mesh = a.shape[1:]
        if mesh[0] in _EXCLUDE and mesh[1] in _EXCLUDE and mesh[2] in _EXCLUDE:
            return _ifftn_blas(a, mesh)
        else:
            return np.fft.ifftn(a, axes=(1,2,3))

#?elif:  # 'FFTW+BLAS' 

def fft(f, mesh):
    '''Perform the 3D FFT from real (R) to reciprocal (G) space.

    After FFT, (u, v, w) -> (j, k, l).
    (jkl) is in the index order of Gv.

    FFT normalization factor is 1., as in MH and in `numpy.fft`.

    Args:
        f : (nx*ny*nz,) ndarray
            The function to be FFT'd, flattened to a 1D array corresponding
            to the index order of :func:`cartesian_prod`.
        mesh : (3,) ndarray of ints (= nx,ny,nz)
            The number G-vectors along each direction.

    Returns:
        (nx*ny*nz,) ndarray
            The FFT 1D array in same index order as Gv (natural order of
            numpy.fft).

    '''
    if f.size == 0:
        return np.zeros_like(f)

    f3d = f.reshape(-1, *mesh)
    assert(f3d.shape[0] == 1 or f[0].size == f3d[0].size)
    g3d = _fftn_wrapper(f3d)
    ngrids = np.prod(mesh)
    if f.ndim == 1 or (f.ndim == 3 and f.size == ngrids):
        return g3d.ravel()
    else:
        return g3d.reshape(-1, ngrids) 

def ifft(g, mesh):
    '''Perform the 3D inverse FFT from reciprocal (G) space to real (R) space.

    Inverse FFT normalization factor is 1./N, same as in `numpy.fft` but
    **different** from MH (they use 1.).

    Args:
        g : (nx*ny*nz,) ndarray
            The function to be inverse FFT'd, flattened to a 1D array
            corresponding to the index order of `span3`.
        mesh : (3,) ndarray of ints (= nx,ny,nz)
            The number G-vectors along each direction.

    Returns:
        (nx*ny*nz,) ndarray
            The inverse FFT 1D array in same index order as Gv (natural order
            of numpy.fft).

    '''
    if g.size == 0:
        return np.zeros_like(g)

    g3d = g.reshape(-1, *mesh)
    assert(g3d.shape[0] == 1 or g[0].size == g3d[0].size)
    f3d = _ifftn_wrapper(g3d)
    ngrids = np.prod(mesh)
    if g.ndim == 1 or (g.ndim == 3 and g.size == ngrids):
        return f3d.ravel()
    else:
        return f3d.reshape(-1, ngrids) 

def fftk(f, mesh, expmikr):
    r'''Perform the 3D FFT of a real-space function which is (periodic*e^{ikr}).

    fk(k+G) = \sum_r fk(r) e^{-i(k+G)r} = \sum_r [f(k)e^{-ikr}] e^{-iGr}
    '''
    return fft(f*expmikr, mesh) 

def frame2logspec(self, frame):
        frame = self.pre_emphasis(frame) * self.win
        fft = numpy.fft.rfft(frame, self.nfft)
        # Square of absolute value
        power = fft.real * fft.real + fft.imag * fft.imag
        return numpy.log(numpy.dot(power, self.filters).clip(1e-5,numpy.inf)) 

def test_numpy_fft():
    bad_results = check_dir(np.fft)
    assert bad_results == {} 

def psd(self):
        psd = 1 / (self.length * np.pi) * np.square(np.abs(np.fft.fft(self.h) * self.length / self.points_number))
        psd = psd[:len(psd) // 2]
        k = np.pi / self.length * np.linspace(0, self.points_number, self.points_number // 2)
        # k = k[len(k) // 2:]
        return (k, psd) 

def eval(self, method='mp'):
        # from ocelot.utils.xfel_utils import calc_wigner
        ds = self.s[1] - self.s[0]
        self.wig = calc_wigner(self.field, method=method, debug=1)
        phen = h_eV_s * (np.fft.fftfreq(self.s.size, d=ds / speed_of_light) + speed_of_light / self.xlamds)
        self.phen = np.fft.fftshift(phen, axes=0)
        # self.freq_lamd = h_eV_s * speed_of_light * 1e9 / freq_ev 

def ana_GetMagnitudeMax(x, Fx, Fs, fftlen):
    # Get maximal spectral magnitude of signal: x around position Fx Hz in dB
    # Fx can be a vector of frequency points
    # Note that the bigger fftlen the better the resolution

    if numpy.isnan(Fx):
        M = float('NaN')
        fmax = float('NaN')

    else:
        xlen = len(x)
        hamlen = min(fftlen, xlen)
    #X = fft(hamming(hamlen).*x(1:hamlen), fftlen);
    factor = 1; #/length(x); # not exactly Kosher but compatible to dfs_magn()
    X = numpy.fft(x,fftlen)
    for i in range(0, xlen):
        if X[i] == 0:
            X[i] = 0.000000001; # guard against log(0) warnings
    for i in range(1, fftlen/2):
        X[i] = 20 * numpy.log10(factor*abs(X[i]))
    fstep = Fs / fftlen
    lowf = (Fx-0.1)*Fx
    if lowf < 0:
        lowf = 0
    highf = Fx + 0.1*Fx
    if highf > Fs/2-fstep:
        highf = Fs/2-fstep

    for cnt in range(0,length(Fx)):
##        if X[cnt] == max(X):
##            break
        pos = cnt
        M = None   # This is a dummy value, see below

        # The code following the equals sign is untranslated Octave Code
        # M[cnt],pos = max(X(1+round(lowf[cnt]/fstep):1+round(highf(cnt)/fstep), :))
        fmax[cnt] = (pos-1+round(lowf(cnt)/fstep))*fstep

    return M, fmax 

def bench_random(self):
        from numpy.fft import fft as numpy_fft
        print()
        print('                 Fast Fourier Transform')
        print('=================================================')
        print('      |    real input     |   complex input    ')
        print('-------------------------------------------------')
        print(' size |  scipy  |  numpy  |  scipy  |  numpy ')
        print('-------------------------------------------------')
        for size,repeat in [(100,7000),(1000,2000),
                            (256,10000),
                            (512,10000),
                            (1024,1000),
                            (2048,1000),
                            (2048*2,500),
                            (2048*4,500),
                            ]:
            print('%5s' % size, end=' ')
            sys.stdout.flush()

            for x in [random([size]).astype(double),
                      random([size]).astype(cdouble)+random([size]).astype(cdouble)*1j
                      ]:
                if size > 500:
                    y = fft(x)
                else:
                    y = direct_dft(x)
                assert_array_almost_equal(fft(x),y)
                print('|%8.2f' % measure('fft(x)',repeat), end=' ')
                sys.stdout.flush()

                assert_array_almost_equal(numpy_fft(x),y)
                print('|%8.2f' % measure('numpy_fft(x)',repeat), end=' ')
                sys.stdout.flush()

            print(' (secs for %s calls)' % (repeat))
        sys.stdout.flush() 

def bench_random(self):
        from numpy.fft import ifft as numpy_ifft
        print()
        print('       Inverse Fast Fourier Transform')
        print('===============================================')
        print('      |     real input    |    complex input   ')
        print('-----------------------------------------------')
        print(' size |  scipy  |  numpy  |  scipy  |  numpy  ')
        print('-----------------------------------------------')
        for size,repeat in [(100,7000),(1000,2000),
                            (256,10000),
                            (512,10000),
                            (1024,1000),
                            (2048,1000),
                            (2048*2,500),
                            (2048*4,500),
                            ]:
            print('%5s' % size, end=' ')
            sys.stdout.flush()

            for x in [random([size]).astype(double),
                      random([size]).astype(cdouble)+random([size]).astype(cdouble)*1j
                      ]:
                if size > 500:
                    y = ifft(x)
                else:
                    y = direct_idft(x)
                assert_array_almost_equal(ifft(x),y)
                print('|%8.2f' % measure('ifft(x)',repeat), end=' ')
                sys.stdout.flush()

                assert_array_almost_equal(numpy_ifft(x),y)
                print('|%8.2f' % measure('numpy_ifft(x)',repeat), end=' ')
                sys.stdout.flush()

            print(' (secs for %s calls)' % (repeat))
        sys.stdout.flush() 

def bench_random(self):
        from numpy.fft import rfft as numpy_rfft
        print()
        print('Fast Fourier Transform (real data)')
        print('==================================')
        print(' size |  scipy  |  numpy  ')
        print('----------------------------------')
        for size,repeat in [(100,7000),(1000,2000),
                            (256,10000),
                            (512,10000),
                            (1024,1000),
                            (2048,1000),
                            (2048*2,500),
                            (2048*4,500),
                            ]:
            print('%5s' % size, end=' ')
            sys.stdout.flush()

            x = random([size]).astype(double)
            print('|%8.2f' % measure('rfft(x)',repeat), end=' ')
            sys.stdout.flush()

            print('|%8.2f' % measure('numpy_rfft(x)',repeat), end=' ')
            sys.stdout.flush()

            print(' (secs for %s calls)' % (repeat))
        sys.stdout.flush() 

def bench_random(self):
        from numpy.fft import irfft as numpy_irfft

        print()
        print('Inverse Fast Fourier Transform (real data)')
        print('==================================')
        print(' size |  scipy  |  numpy  ')
        print('----------------------------------')
        for size,repeat in [(100,7000),(1000,2000),
                            (256,10000),
                            (512,10000),
                            (1024,1000),
                            (2048,1000),
                            (2048*2,500),
                            (2048*4,500),
                            ]:
            print('%5s' % size, end=' ')
            sys.stdout.flush()

            x = random([size]).astype(double)
            x1 = zeros(size/2+1,dtype=cdouble)
            x1[0] = x[0]
            for i in range(1,size/2):
                x1[i] = x[2*i-1] + 1j * x[2*i]
            if not size % 2:
                x1[-1] = x[-1]
            y = irfft(x)

            print('|%8.2f' % measure('irfft(x)',repeat), end=' ')
            sys.stdout.flush()

            assert_array_almost_equal(numpy_irfft(x1,size),y)
            print('|%8.2f' % measure('numpy_irfft(x1,size)',repeat), end=' ')
            sys.stdout.flush()

            print(' (secs for %s calls)' % (repeat))

        sys.stdout.flush()