Python numpy.fft 模块，fftn() 实例源码

def run_test_c2c_impl(self, shape, axes, inverse=False, fftshift=False):
        shape = list(shape)
        shape[-1] *= 2 # For complex
        known_data = np.random.normal(size=shape).astype(np.float32).view(np.complex64)
        idata = bf.ndarray(known_data, space='cuda')
        odata = bf.empty_like(idata)
        fft = Fft()
        fft.init(idata, odata, axes=axes, apply_fftshift=fftshift)
        fft.execute(idata, odata, inverse)
        if inverse:
            if fftshift:
                known_data = np.fft.ifftshift(known_data, axes=axes)
            # Note: Numpy applies normalization while CUFFT does not
            norm = reduce(lambda a, b: a * b, [known_data.shape[d]
                                               for d in axes])
            known_result = gold_ifftn(known_data, axes=axes) * norm
        else:
            known_result = gold_fftn(known_data, axes=axes)
            if fftshift:
                known_result = np.fft.fftshift(known_result, axes=axes)
        x = (np.abs(odata.copy('system') - known_result) / known_result > RTOL).astype(np.int32)
        a = odata.copy('system')
        b = known_result
        compare(odata.copy('system'), known_result)

def correlate(x, y):

    from numpy import fft

    sx = numpy.array(x.shape)
    sy = numpy.array(y.shape)

    if (sx >= sy).sum():

        slices = [slice(None, sx[i] - sy[i] + 1) for i in range(len(sx))]

        X = fft.fftn(x)
        Y = fft.fftn(zerofill(y, sx))

    else:

        sf = sx + sy - 1
        slices = [slice(None, sf[i]) for i in range(len(sf))]

        X = fft.fftn(x, sf)
        Y = fft.fftn(zerofill(y, sf), sf)

    return fft.ifftn(X.conjugate() * Y)[slices].real

def crystal():
    # look at crystal and their ffts
    crystal3D = np.zeros((201,201,201), dtype= np.int32)
    crystal3D_fourier = np.zeros((201,201,201), dtype= np.complex64)

    dx = 1

    for row in range(60,140,dx):
        for col in range(80,120,dx):
            for time in range(90,110,dx):
                crystal3D[row,col,time] = 1

    crystal3D_fourier = fft.fftshift(fft.fftn(crystal3D))
    #del crystal3D
    diffPattern3D = (abs(crystal3D_fourier)**2)
    del crystal3D_fourier
    return diffPattern3D

def convolve(x, f):

    from numpy import fft, all

    sx = numpy.array(x.shape)
    sf = numpy.array(f.shape)

    if not all(sx >= sf): return convolve(f, x)

    y = fft.ifftn(fft.fftn(x) * fft.fftn(f, sx)).real
    slices = [slice(sf[i] - 1, sx[i]) for i in range(len(sf))]

    return y[slices]

def test_UnscaledFFT_3d(backend, M, N, K, B ):
    b = backend()

    # forward
    x = b.rand_array( (M*N*K, B) )
    y = b.rand_array( (M*N*K, B) )
    x_h = x.to_host().reshape( (M,N,K,B), order='F' )

    A = b.UnscaledFFT( (M,N,K), dtype=x.dtype )

    A.eval(y, x)

    y_exp = np.fft.fftn( x_h, axes=(0,1,2) )
    y_act = y.to_host().reshape( (M,N,K,B), order='F' )
    npt.assert_allclose(y_act, y_exp, rtol=1e-2)

    # adjoint
    x = b.rand_array( (M*N*K, B) )
    y = b.rand_array( (M*N*K, B) )
    x_h = x.to_host().reshape( (M,N,K,B), order='F' )

    A.H.eval(y, x)

    y_exp = np.fft.ifftn( x_h, axes=(0,1,2) ) * (M*N*K)
    y_act = y.to_host().reshape( (M,N,K,B), order='F' )
    npt.assert_allclose(y_act, y_exp, rtol=1e-2)

def test_UnscaledFFT_2d(backend, M, N, B ):
    b = backend()

    # forward
    x = b.rand_array( (M*N, B) )
    y = b.rand_array( (M*N, B) )
    x_h = x.to_host().reshape( (M,N,B), order='F' )

    A = b.UnscaledFFT( (M,N), dtype=x.dtype )

    A.eval(y, x)

    y_exp = np.fft.fftn( x_h, axes=(0,1) )
    y_act = y.to_host().reshape( (M,N,B), order='F' )
    npt.assert_allclose(y_act, y_exp, rtol=1e-2)

    # adjoint
    x = b.rand_array( (M*N, B) )
    y = b.rand_array( (M*N, B) )
    x_h = x.to_host().reshape( (M,N,B), order='F' )

    A.H.eval(y, x)

    y_exp = np.fft.ifftn( x_h, axes=(0,1) ) * (M*N)
    y_act = y.to_host().reshape( (M,N,B), order='F' )
    npt.assert_allclose(y_act, y_exp, rtol=1e-2)

def test_UnscaledFFT_1d(backend, M, B ):
    b = backend()

    # forward
    x = b.rand_array( (M, B) )
    y = b.rand_array( (M, B) )
    x_h = x.to_host().reshape( (M,B), order='F' )

    A = b.UnscaledFFT( (M,), dtype=x.dtype )

    A.eval(y, x)

    y_exp = np.fft.fftn( x_h, axes=(0,) )
    y_act = y.to_host().reshape( (M,B), order='F' )
    npt.assert_allclose(y_act, y_exp, rtol=1e-2)

    # adjoint
    x = b.rand_array( (M, B) )
    y = b.rand_array( (M, B) )
    x_h = x.to_host().reshape( (M,B), order='F' )

    A.H.eval(y, x)

    y_exp = np.fft.ifftn( x_h, axes=(0,) ) * M
    y_act = y.to_host().reshape( (M,B), order='F' )
    npt.assert_allclose(y_act, y_exp, rtol=1e-2)

def test_CenteredFFT(backend, M, N, K, B ):
    from numpy.fft import fftshift, ifftshift, fftn, ifftn

    b = backend()
    A = b.FFTc( (M,N,K), dtype=np.dtype('complex64') )

    # forward
    ax = (0,1,2)
    x = b.rand_array( (M*N*K,B) )
    y = b.rand_array( (M*N*K,B) )
    x_h = x.to_host().reshape( (M,N,K,B), order='F' )

    A.eval(y, x)

    y_act = y.to_host().reshape( (M,N,K,B), order='F' )
    y_exp = fftshift( fftn( ifftshift(x_h, axes=ax), axes=ax, norm='ortho'), axes=ax)
    npt.assert_allclose(y_act, y_exp, rtol=1e-2)

    # adjoint
    x = b.rand_array( (M*N*K,B) )
    y = b.rand_array( (M*N*K,B) )
    x_h = x.to_host().reshape( (M,N,K,B), order='F' )

    A.H.eval(y, x)

    y_act = y.to_host().reshape( (M,N,K,B), order='F' )
    y_exp = fftshift( ifftn( ifftshift(x_h, axes=ax), axes=ax, norm='ortho'), axes=ax)
    npt.assert_allclose(y_act, y_exp, rtol=1e-2)