test_xrft.py 文件源码

def test_power_spectrum():
    """Test the power spectrum function"""
    N = 16
    da = xr.DataArray(np.random.rand(N,N), dims=['x','y'],
                    coords={'x':range(N),'y':range(N)}
                     )
    ps = xrft.power_spectrum(da, window=True, density=False,
                            detrend='constant')
    daft = xrft.dft(da,
                    dim=None, shift=True, detrend='constant',
                    window=True)
    npt.assert_almost_equal(ps.values, np.real(daft*np.conj(daft)))
    npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)

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

    ### Remove mean
    da = xr.DataArray(np.random.rand(5,20,30),
                  dims=['time', 'y', 'x'],
                  coords={'time': np.arange(5),
                        'y': np.arange(20), 'x': np.arange(30)})
    ps = xrft.power_spectrum(da, dim=['y', 'x'],
                            window=True, density=False, detrend='constant'
                            )
    daft = xrft.dft(da, dim=['y','x'], window=True, detrend='constant')
    npt.assert_almost_equal(ps.values, np.real(daft*np.conj(daft)))
    npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)

    ### Remove least-square fit
    da_prime = np.zeros_like(da.values)
    for t in range(5):
        da_prime[t] = numpy_detrend(da[t].values)
    da_prime = xr.DataArray(da_prime, dims=da.dims, coords=da.coords)
    ps = xrft.power_spectrum(da_prime, dim=['y', 'x'],
                            window=True, density=False, detrend='constant'
                            )
    daft = xrft.dft(da_prime, dim=['y','x'], window=True, detrend='constant')
    npt.assert_almost_equal(ps.values, np.real(daft*np.conj(daft)))
    npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)