test_xrft.py 文件源码

def test_dft_1d(test_data_1d):
    """Test the discrete Fourier transform function on one-dimensional data."""
    da = test_data_1d
    Nx = len(da)
    dx = float(da.x[1] - da.x[0]) if 'x' in da else 1

    # defaults with no keyword args
    ft = xrft.dft(da, detrend='constant')
    # check that the frequency dimension was created properly
    assert ft.dims == ('freq_x',)
    # check that the coords are correct
    freq_x_expected = np.fft.fftshift(np.fft.fftfreq(Nx, dx))
    npt.assert_allclose(ft['freq_x'], freq_x_expected)
    # check that a spacing variable was created
    assert ft['freq_x_spacing'] == freq_x_expected[1] - freq_x_expected[0]
    # make sure the function is lazy
    assert isinstance(ft.data, type(da.data))
    # check that the Fourier transform itself is correct
    data = (da - da.mean()).values
    ft_data_expected = np.fft.fftshift(np.fft.fft(data))
    # because the zero frequency component is zero, there is a numerical
    # precision issue. Fixed by setting atol
    npt.assert_allclose(ft_data_expected, ft.values, atol=1e-14)


    # redo without removing mean
    ft = xrft.dft(da)
    ft_data_expected = np.fft.fftshift(np.fft.fft(da))
    npt.assert_allclose(ft_data_expected, ft.values)

    # redo with detrending linear least-square fit
    ft = xrft.dft(da, detrend='linear')
    da_prime = sps.detrend(da.values)
    ft_data_expected = np.fft.fftshift(np.fft.fft(da_prime))
    npt.assert_allclose(ft_data_expected, ft.values, atol=1e-14)

    if 'x' in da and not da.chunks:
        da['x'].values[-1] *= 2
        with pytest.raises(ValueError):
            ft = xrft.dft(da)