test_fft.py 文件源码

def test_against_fft_3d(self):
        from pyculib.fft.binding import Plan, CUFFT_R2C
        depth = 2
        colsize = 2
        rowsize = 64
        N = depth * colsize * rowsize
        x = np.arange(N, dtype=np.float32).reshape(depth, colsize, rowsize)
        xf = np.fft.fftn(x)
        d_x_gpu = cuda.to_device(x)
        xf_gpu = np.zeros(shape=(depth, colsize, rowsize//2 + 1), dtype=np.complex64)
        d_xf_gpu = cuda.to_device(xf_gpu)
        plan = Plan.many(x.shape, CUFFT_R2C)
        plan.forward(d_x_gpu, d_xf_gpu)
        d_xf_gpu.copy_to_host(xf_gpu)
        self.assertTrue(np.allclose(xf[:, :, 0:rowsize//2+1], xf_gpu, atol=1e-6))