test_xrft.py 文件源码

def test_dft_3d_dask():
    """Test the discrete Fourier transform on 3D dask array data"""
    N=16
    da = xr.DataArray(np.random.rand(N,N,N), dims=['time','x','y'],
                      coords={'time':range(N),'x':range(N),
                              'y':range(N)}
                     )
    daft = xrft.dft(da.chunk({'time': 1}), dim=['x','y'], shift=False)
    # assert hasattr(daft.data, 'dask')
    npt.assert_almost_equal(daft.values,
                        np.fft.fftn(da.chunk({'time': 1}).values, axes=[1,2])
                           )

    with pytest.raises(ValueError):
        xrft.dft(da.chunk({'time': 1, 'x': 1}), dim=['x'])

    daft = xrft.dft(da.chunk({'x': 1}), dim=['time'],
                    shift=False, detrend='linear')
    # assert hasattr(daft.data, 'dask')
    da_prime = sps.detrend(da.chunk({'x': 1}), axis=0)
    npt.assert_almost_equal(daft.values,
                        np.fft.fftn(da_prime, axes=[0])
                           )