def test_fft_real_2d(self): """ Test the real discrete Fourier transform function on one-dimensional data. Non-trivial because we need to keep only some of the negative frequencies. """ Nx, Ny = 16, 32 da = xr.DataArray( np.random.rand(Nx, Ny), dims=["x", "y"], coords={"x": range(Nx), "y": range(Ny)}, ) dx = float(da.x[1] - da.x[0]) dy = float(da.y[1] - da.y[0]) daft = xrft.fft(da, real="x") npt.assert_almost_equal( daft.values, np.fft.rfftn(da.transpose("y", "x")).transpose() ) npt.assert_almost_equal(daft.values, xrft.fft(da, dim=["y"], real="x")) actual_freq_x = daft.coords["freq_x"].values expected_freq_x = np.fft.rfftfreq(Nx, dx) npt.assert_almost_equal(actual_freq_x, expected_freq_x) actual_freq_y = daft.coords["freq_y"].values expected_freq_y = np.fft.fftfreq(Ny, dy) npt.assert_almost_equal(actual_freq_y, expected_freq_y)
def test_fft_4d(self): """Test the discrete Fourier transform on 2D data""" N = 16 da = xr.DataArray( np.random.rand(N, N, N, N), dims=["time", "z", "y", "x"], coords={ "time": range(N), "z": range(N), "y": range(N), "x": range(N) }, ) ft = xrft.fft(da, shift=False) npt.assert_almost_equal(ft.values, np.fft.fftn(da.values)) dim = ["time", "y", "x"] da_prime = xrft.detrend(da[:, 0], dim) # cubic detrend over time, y, and x npt.assert_almost_equal( xrft.fft( da[:, 0].drop("z"), dim=dim, shift=False, detrend="linear", ).values, np.fft.fftn(da_prime), )
def test_fft_2d(self): """Test the discrete Fourier transform on 2D data""" N = 16 da = xr.DataArray(np.random.rand(N, N), dims=["x", "y"], coords={ "x": range(N), "y": range(N) }) ft = xrft.fft(da, shift=False) npt.assert_almost_equal(ft.values, np.fft.fftn(da.values)) ft = xrft.fft(da, shift=False, window="hann", detrend="constant") dim = da.dims window = (sps.windows.hann(N, sym=False) * sps.windows.hann(N, sym=False)[:, np.newaxis]) da_prime = (da - da.mean(dim=dim)).values npt.assert_almost_equal(ft.values, np.fft.fftn(da_prime * window)) da = xr.DataArray( np.random.rand(N, N), dims=["x", "y"], coords={ "x": range(N, 0, -1), "y": range(N, 0, -1) }, ) assert (xrft.power_spectrum(da, shift=False, density=True) >= 0.0).all()
def test_fft_real_1d(self, 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.dims else 1 # defaults with no keyword args ft = xrft.fft(da, real="x", 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.rfftfreq(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.rfft(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) with pytest.raises(ValueError): xrft.fft(da, real="y", detrend="constant")
def test_dim_fft(self): N = 16 da = xr.DataArray( np.random.rand(N, N), dims=["x", "y"], coords={"x": range(N), "y": range(N)} ) npt.assert_array_equal( xrft.fft(da, dim="y", shift=False).values, xrft.fft(da, dim=["y"], shift=False).values, ) assert xrft.fft(da, dim="y").dims == ("x", "freq_y")
def test_spacing_tol(test_data_1d): da = test_data_1d da2 = da.copy().load() # Create improperly spaced data Nx = 16 Lx = 1.0 x = np.linspace(0, Lx, Nx) x[-1] = x[-1] + 0.001 da3 = xr.DataArray(np.random.rand(Nx), coords=[x], dims=["x"]) # This shouldn't raise an error xrft.fft(da3, spacing_tol=1e-1) # But this should with pytest.raises(ValueError): xrft.fft(da3, spacing_tol=1e-4)
def test_fft_nocoords(): # Julius' example # https://github.com/rabernat/xrft/issues/17 data = xr.DataArray(np.random.random([20, 30, 100]), dims=["time", "lat", "lon"]) dft = xrft.fft(data, dim=["time"]) ps = xrft.power_spectrum(data, dim=["time"])
def test_fft_3d_dask(self, 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)}, ) if dask: da = da.chunk({"time": 1}) daft = xrft.fft(da, dim=["x", "y"], shift=False) npt.assert_almost_equal( daft.values, np.fft.fftn(da.chunk({"time": 1}).values, axes=[1, 2]) ) da = da.chunk({"x": 1}) with pytest.raises(ValueError): xrft.fft(da, dim=["x"]) with pytest.raises(ValueError): xrft.fft(da, dim="x") da = da.chunk({"time": N}) daft = xrft.fft(da, dim=["time"], shift=False, detrend="linear") da_prime = sps.detrend(da, axis=0) npt.assert_almost_equal(daft.values, np.fft.fftn(da_prime, axes=[0])) npt.assert_array_equal( daft.values, xrft.fft(da, dim="time", shift=False, detrend="linear") )
def test_cross_spectrum(self, dask): """Test the cross spectrum function""" N = 16 dim = ["x", "y"] da = xr.DataArray( np.random.rand(2, N, N), dims=["time", "x", "y"], coords={ "time": np.array(["2019-04-18", "2019-04-19"], dtype="datetime64"), "x": range(N), "y": range(N), }, ) da2 = xr.DataArray( np.random.rand(2, N, N), dims=["time", "x", "y"], coords={ "time": np.array(["2019-04-18", "2019-04-19"], dtype="datetime64"), "x": range(N), "y": range(N), }, ) if dask: da = da.chunk({"time": 1}) da2 = da2.chunk({"time": 1}) daft = xrft.fft(da, dim=dim, shift=True, detrend="constant", window=True) daft2 = xrft.fft(da2, dim=dim, shift=True, detrend="constant", window=True) cs = xrft.cross_spectrum( da, da2, dim=dim, window=True, density=False, detrend="constant" ) npt.assert_almost_equal(cs.values, daft * np.conj(daft2)) npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 0.0) cs = xrft.cross_spectrum( da, da2, dim=dim, shift=True, window=True, detrend="constant" ) test = (daft * np.conj(daft2)).values / N ** 4 dk = np.diff(np.fft.fftfreq(N, 1.0))[0] test /= dk ** 2 npt.assert_almost_equal(cs.values, test) npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 0.0)
def test_power_spectrum(self, dask): """Test the power spectrum function""" N = 16 da = xr.DataArray( np.random.rand(2, N, N), dims=["time", "y", "x"], coords={ "time": np.array(["2019-04-18", "2019-04-19"], dtype="datetime64"), "y": range(N), "x": range(N), }, ) if dask: da = da.chunk({"time": 1}) ps = xrft.power_spectrum( da, dim=["y", "x"], window=True, density=False, detrend="constant" ) daft = xrft.fft(da, dim=["y", "x"], 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.0) ps = xrft.power_spectrum( da, dim=["y"], real="x", window=True, density=False, detrend="constant" ) daft = xrft.fft(da, dim=["y"], real="x", detrend="constant", window=True) npt.assert_almost_equal(ps.values, np.real(daft * np.conj(daft))) ### Normalized ps = xrft.power_spectrum(da, dim=["y", "x"], window=True, detrend="constant") daft = xrft.fft(da, dim=["y", "x"], window=True, detrend="constant") test = np.real(daft * np.conj(daft)) / N ** 4 dk = np.diff(np.fft.fftfreq(N, 1.0))[0] test /= dk ** 2 npt.assert_almost_equal(ps.values, test) npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.0) ### Remove least-square fit ps = xrft.power_spectrum( da, dim=["y", "x"], window=True, density=False, detrend="linear" ) daft = xrft.fft(da, dim=["y", "x"], window=True, detrend="linear") npt.assert_almost_equal(ps.values, np.real(daft * np.conj(daft))) npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.0)
def test_fft_1d(self, 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.dims else 1 # defaults with no keyword args ft = xrft.fft(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.fft(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.fft(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.fft(da)
def test_ifft_fft(): """ Testing ifft(fft(s.data)) == s.data """ N = 20 s = xr.DataArray( np.random.rand(N) + 1j * np.random.rand(N), dims="x", coords={"x": np.arange(0, N)}, ) FTs = xrft.fft(s) IFTs = xrft.ifft(FTs, shift=True) # Shift=True is mandatory for the assestion below npt.assert_allclose(s.data, IFTs.data)
def test_fft_1d_time(self, time_data): """Test the discrete Fourier transform function on timeseries data.""" time = time_data Nt = len(time) da = xr.DataArray(np.random.rand(Nt), coords=[time], dims=["time"]) ft = xrft.fft(da, shift=False) # check that frequencies are correct if pd.api.types.is_datetime64_dtype(time): dt = (time[1] - time[0]).total_seconds() else: dt = np.diff(time)[0].total_seconds() freq_time_expected = np.fft.fftfreq(Nt, dt) npt.assert_allclose(ft["freq_time"], freq_time_expected)
def diff_fft_xr(da, n=1, dim=None, shift=False, real=True, reshape=False, **kwargs): """ Calculate derivative of a DataArray using fft Notes: - Changing the zero frequency doesn't change the output as far as the derivative goes. - Normalizing the fourier coefficients by N, yields derivatives that are 1/N the expected value. - The nyquist frequency is generally pretty small so setting it to zero generally doesn't do anything. """ import numpy as np import xrft import xarray as xr kdim = 'freq_' + dim original_coords = da[dim] #---- # Perform forward DFT wave = xrft.fft(da, dim=dim, **kwargs) freq = wave.coords[kdim] #---- #---- # Multiply amplitudes by 2πk multiplied_amplitudes = pow(1j * 2 * np.pi * freq, n) * wave #---- #---- # Perform inverse DFT deriv = xrft.ifft(multiplied_amplitudes, dim=kdim, **kwargs) #---- #---- return deriv.real.assign_coords({dim: original_coords})
def test_power_spectrum(self, dask): """Test the power spectrum function""" N = 16 da = xr.DataArray( np.random.rand(N), dims=["x"], coords={ "x": range(N), }, ) f_scipy, p_scipy = sps.periodogram(da.values, window="rectangular", return_onesided=True) ps = xrft.power_spectrum(da, dim="x", real_dim="x", detrend="constant") npt.assert_almost_equal(ps.values, p_scipy) A = 20 fs = 1e4 n_segments = int(fs // 10) fsig = 300 ii = int(fsig * n_segments // fs) # Freq index of fsig tt = np.arange(fs) / fs x = A * np.sin(2 * np.pi * fsig * tt) for window_type in ["hann", "bartlett", "tukey", "flattop"]: # see https://github.com/scipy/scipy/blob/master/scipy/signal/tests/test_spectral.py#L485 x_da = xr.DataArray(x, coords=[tt], dims=["t"]).chunk({"t": n_segments}) ps = xrft.power_spectrum( x_da, dim="t", window=window_type, chunks_to_segments=True, window_correction=True, ).mean("t_segment") # Check the energy correction npt.assert_allclose( np.sqrt(np.trapz(ps.values, ps.freq_t.values)), A * np.sqrt(2) / 2, rtol=1e-3, ) ps = xrft.power_spectrum( x_da, dim="t", window=window_type, chunks_to_segments=True, scaling="spectrum", window_correction=True, ).mean("t_segment") # Check the amplitude correction # The factor of 0.5 is there because we're checking the two-sided spectrum npt.assert_allclose(ps.sel(freq_t=fsig), 0.5 * A**2 / 2.0) da = xr.DataArray( np.random.rand(2, N, N), dims=["time", "y", "x"], coords={ "time": np.array(["2019-04-18", "2019-04-19"], dtype="datetime64"), "y": range(N), "x": range(N), }, ) if dask: da = da.chunk({"time": 1}) ps = xrft.power_spectrum(da, dim=["y", "x"], window="hann", density=False, detrend="constant") daft = xrft.fft(da, dim=["y", "x"], detrend="constant", window="hann") npt.assert_almost_equal(ps.values, np.real(daft * np.conj(daft))) npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.0) ps = xrft.power_spectrum( da, dim=["y"], real_dim="x", window="hann", density=False, detrend="constant", ) daft = xrft.fft(da, dim=["y"], real_dim="x", detrend="constant", window="hann") ps_test = np.real(daft * np.conj(daft)) f = np.full(ps_test.sizes["freq_x"], 2.0) f[0], f[-1] = 1.0, 1.0 ps_test = ps_test * xr.DataArray(f, dims="freq_x") npt.assert_almost_equal(ps.values, ps_test.values) ### Normalized ps = xrft.power_spectrum(da, dim=["y", "x"], window="hann", detrend="constant") daft = xrft.fft(da, dim=["y", "x"], window="hann", detrend="constant") test = np.real(daft * np.conj(daft)) / N**4 dk = np.diff(np.fft.fftfreq(N, 1.0))[0] test /= dk**2 npt.assert_almost_equal(ps.values, test) npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.0) ### Remove least-square fit ps = xrft.power_spectrum(da, dim=["y", "x"], window="hann", density=False, detrend="linear") daft = xrft.fft(da, dim=["y", "x"], window="hann", detrend="linear") npt.assert_almost_equal(ps.values, np.real(daft * np.conj(daft))) npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.0) with pytest.raises(ValueError): xrft.power_spectrum(da, dim=["y", "x"], window=None, window_correction=True)
def test_spacing_tol_float_value(test_data_1d): da = test_data_1d with pytest.raises(TypeError): xrft.fft(da, spacing_tol="string")
def test_chunks_to_segments(): N = 32 da = xr.DataArray( np.random.rand(N, N, N), dims=["time", "y", "x"], coords={"time": range(N), "y": range(N), "x": range(N)}, ) with pytest.raises(ValueError): xrft.fft( da.chunk(chunks=((20, N, N), (N - 20, N, N))), dim=["time"], detrend="linear", chunks_to_segments=True, ) ft = xrft.fft( da.chunk({"time": 16}), dim=["time"], shift=False, chunks_to_segments=True ) assert ft.dims == ("time_segment", "freq_time", "y", "x") data = da.chunk({"time": 16}).data.reshape((2, 16, N, N)) npt.assert_almost_equal(ft.values, dsar.fft.fftn(data, axes=[1]), decimal=7) ft = xrft.fft( da.chunk({"y": 16, "x": 16}), dim=["y", "x"], shift=False, chunks_to_segments=True, ) assert ft.dims == ("time", "y_segment", "freq_y", "x_segment", "freq_x") data = da.chunk({"y": 16, "x": 16}).data.reshape((N, 2, 16, 2, 16)) npt.assert_almost_equal(ft.values, dsar.fft.fftn(data, axes=[2, 4]), decimal=7) ps = xrft.power_spectrum( da.chunk({"y": 16, "x": 16}), dim=["y", "x"], shift=False, density=False, chunks_to_segments=True, ) npt.assert_almost_equal( ps.values, (ft * np.conj(ft)).values, ) da2 = xr.DataArray( np.random.rand(N, N, N), dims=["time", "y", "x"], coords={"time": range(N), "y": range(N), "x": range(N)}, ) ft2 = xrft.fft( da2.chunk({"y": 16, "x": 16}), dim=["y", "x"], shift=False, chunks_to_segments=True, ) cs = xrft.cross_spectrum( da.chunk({"y": 16, "x": 16}), da2.chunk({"y": 16, "x": 16}), dim=["y", "x"], shift=False, density=False, chunks_to_segments=True, ) npt.assert_almost_equal( cs.values, (ft * np.conj(ft2)).values, )