test_xrft.py 文件源码

def test_parseval():
    """Test whether the Parseval's relation is satisfied."""

    N = 16
    da = xr.DataArray(np.random.rand(N,N),
                    dims=['x','y'], coords={'x':range(N), 'y':range(N)})
    da2 = xr.DataArray(np.random.rand(N,N),
                    dims=['x','y'], coords={'x':range(N), 'y':range(N)})

    dim = da.dims
    delta_x = []
    for d in dim:
        coord = da[d]
        diff = np.diff(coord)
        delta = diff[0]
        delta_x.append(delta)

    window = np.hanning(N) * np.hanning(N)[:, np.newaxis]
    ps = xrft.power_spectrum(da, window=True, detrend='constant')
    da_prime = da.values - da.mean(dim=dim).values
    npt.assert_almost_equal(ps.values.sum(),
                            (np.asarray(delta_x).prod()
                            * ((da_prime*window)**2).sum()
                            ), decimal=5
                            )

    cs = xrft.cross_spectrum(da, da2, window=True, detrend='constant')
    da2_prime = da2.values - da2.mean(dim=dim).values
    npt.assert_almost_equal(cs.values.sum(),
                            (np.asarray(delta_x).prod()
                            * ((da_prime*window)
                            * (da2_prime*window)).sum()
                            ), decimal=5
                            )

    d3d = xr.DataArray(np.random.rand(N,N,N),
                    dims=['time','y','x'],
                    coords={'time':range(N), 'y':range(N), 'x':range(N)}
                      ).chunk({'time':1})
    ps = xrft.power_spectrum(d3d, dim=['x','y'], window=True, detrend='linear')
    npt.assert_almost_equal(ps[0].values.sum(),
                            (np.asarray(delta_x).prod()
                            * ((numpy_detrend(d3d[0].values)*window)**2).sum()
                            ), decimal=5
                           )