test_operators.py 文件源码

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)