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,
)
