test_fft.py 文件源码

def test_plan3d(self):
        from pyculib.fft.binding import Plan, CUFFT_C2C
        n = 32
        data = np.arange(n, dtype=np.complex64).reshape(2, 2, 8)

        orig = data.copy()
        d_data = cuda.to_device(data)
        fftplan = Plan.three(CUFFT_C2C, *data.shape)
        fftplan.forward(d_data, d_data)
        fftplan.inverse(d_data, d_data)
        d_data.copy_to_host(data)
        result = data / n
        self.assertTrue(np.allclose(orig, result.real))