test_xrft.py 文件源码

def test_dft_4d():
    """Test the discrete Fourier transform on 2D data"""
    N = 16
    da = xr.DataArray(np.random.rand(N,N,N,N),
                    dims=['time','z','y','x'],
                    coords={'time':range(N),'z':range(N),
                            'y':range(N),'x':range(N)}
                     )
    ft = xrft.dft(da, shift=False)
    npt.assert_almost_equal(ft.values, np.fft.fftn(da.values))

    with pytest.raises(NotImplementedError):
        xrft.dft(da, detrend='linear')
    with pytest.raises(NotImplementedError):
        xrft.dft(da, dim=['time','y','x'], detrend='linear')

    da_prime = xrft.detrendn(da[:,0].values, [0,1,2]) # cubic detrend over time, y, and x
    npt.assert_almost_equal(xrft.dft(da[:,0].drop('z'),
                                    dim=['time','y','x'],
                                    shift=False, detrend='linear'
                                    ).values,
                            np.fft.fftn(da_prime))