test_xrft.py 文件源码

def test_power_spectrum_dask():
    """Test the power spectrum function on dask data"""
    N = 16
    dim = ['x','y']
    da = xr.DataArray(np.random.rand(2,N,N), dims=['time','x','y'],
                      coords={'time':range(2),'x':range(N),
                              'y':range(N)}).chunk({'time': 1}
                     )
    ps = xrft.power_spectrum(da, dim=dim, density=False)
    daft = xrft.dft(da, dim=['x','y'])
    npt.assert_almost_equal(ps.values, (daft * np.conj(daft)).real.values)

    ps = xrft.power_spectrum(da, dim=dim, window=True, detrend='constant')
    daft = xrft.dft(da, dim=dim, window=True, detrend='constant')
    coord = list(daft.coords)
    test = (daft * np.conj(daft)).real/N**4
    for i in dim:
        test /= daft['freq_' + i + '_spacing']
    npt.assert_almost_equal(ps.values, test)
    npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)