test_fft.py 文件源码

def test_fft_3d_single(self):
        from pyculib.fft import fft
        N3 = 2
        N2 = 2
        N1 = 32
        N = N1 * N2 * N3
        x = np.asarray(np.arange(N), dtype=np.float32).reshape(N3, N2, N1)
        xf = np.fft.fftn(x)

        xf_gpu = np.empty(shape=(N3, N2, N1//2 + 1), dtype=np.complex64)
        fft(x, xf_gpu)

        self.assertTrue( np.allclose(xf[:, :, 0:N1//2+1], xf_gpu, atol=1e-6) )