def test_dft_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.dft(da, real='x')
npt.assert_almost_equal(daft.values,
np.fft.rfftn(da.transpose('y','x')).transpose())
npt.assert_almost_equal(daft.values,
xrft.dft(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_dft_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.dft(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.dft(da, real='y', detrend='constant')
def test_dft_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.dft(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.dft(
da[:, 0].drop("z"),
dim=dim,
shift=False,
detrend="linear",
).values,
np.fft.fftn(da_prime),
)
def test_dft_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.dft(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.dft(da, dim=['x'])
da = da.chunk({'time': N})
daft = xrft.dft(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]))
def test_dft_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.dft(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.dft(da, real="y", detrend="constant")
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.dft(da, dim=dim, shift=True, detrend='constant',
window=True)
daft2 = xrft.dft(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, np.real(daft*np.conj(daft2)))
npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 0.)
cs = xrft.cross_spectrum(da, da2, dim=dim, shift=True, window=True,
detrend='constant')
test = (daft * np.conj(daft2)).real.values/N**4
dk = np.diff(np.fft.fftfreq(N, 1.))[0]
test /= dk**2
npt.assert_almost_equal(cs.values, test)
npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 0.)
def test_dft_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.dft(da, real="x")
npt.assert_almost_equal(
daft.values,
np.fft.rfftn(da.transpose("y", "x")).transpose())
npt.assert_almost_equal(daft.values, xrft.dft(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_cross_spectrum():
"""Test the cross spectrum function"""
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)
})
cs = xrft.cross_spectrum(da,
da2,
window=True,
density=False,
detrend='constant')
daft = xrft.dft(da, dim=None, shift=True, detrend='constant', window=True)
daft2 = xrft.dft(da2,
dim=None,
shift=True,
detrend='constant',
window=True)
npt.assert_almost_equal(cs.values, np.real(daft * np.conj(daft2)))
npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 0.)
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.dft(da, dim=dim, shift=True, detrend='constant',
window=True)
daft2 = xrft.dft(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, np.real(daft*np.conj(daft2)))
npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 0.)
cs = xrft.cross_spectrum(da, da2, dim=dim, shift=True, window=True,
detrend='constant')
test = (daft * np.conj(daft2)).real.values/N**4
dk = np.diff(np.fft.fftfreq(N, 1.))[0]
test /= dk**2
npt.assert_almost_equal(cs.values, test)
npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 0.)
def test_dft_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.dft(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.dft(da, dim=['x'])
da = da.chunk({'time':N})
daft = xrft.dft(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])
)
def test_dft_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.dft(da, shift=False)
npt.assert_almost_equal(ft.values, np.fft.fftn(da.values))
ft = xrft.dft(da, shift=False, window=True, detrend="constant")
dim = da.dims
window = np.hanning(N) * np.hanning(N)[:, 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_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.dft(da,
dim=dim,
shift=True,
detrend="constant",
window=True)
daft2 = xrft.dft(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, np.real(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)).real.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.dft(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.)
ps = xrft.power_spectrum(da,
dim=['y'],
real='x',
window=True,
density=False,
detrend='constant')
daft = xrft.dft(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.dft(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]
test /= dk**2
npt.assert_almost_equal(ps.values, test)
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)
### Remove least-square fit
ps = xrft.power_spectrum(da,
dim=['y', 'x'],
window=True,
density=False,
detrend='linear')
daft = xrft.dft(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.)
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.dft(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.dft(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.dft(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.dft(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_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.)
def test_dim_dft(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.dft(da, dim='y', shift=False).values,
xrft.dft(da, dim=['y'], shift=False).values
)
assert xrft.dft(da, dim='y').dims == ('x','freq_y')
def test_reversed_coordinates():
"""Reversed coordinates should not impact dft with true_phase = True"""
N = 20
s = xr.DataArray(
np.random.rand(N) + 1j * np.random.rand(N),
dims="x",
coords={"x": np.arange(N // 2, -N // 2, -1) + 2},
)
s2 = s.sortby("x")
xrt.assert_allclose(xrft.dft(s, dim="x", true_phase=True),
xrft.dft(s2, dim="x", true_phase=True))
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
#----
# Perform forward DFT
if real:
wave = xrft.dft(da, dim=dim, shift=shift, real=dim, **kwargs)
else:
wave = xrft.dft(da, dim=dim, shift=shift, real=None, **kwargs)
freq = wave.coords[kdim]
#----
#----
# Multiply amplitudes by 2πk
ikF = (pow(1.j * 2. * np.pi, n) * wave * freq)
#----
#----
# Perform inverse DFT
if real:
deriv = np.fft.irfft(ikF, axis=wave.dims.index(kdim))
else:
deriv = np.fft.ifft(ikF, axis=wave.dims.index(kdim)).real
#----
#----
# Repackage into dataarray
newdims = [
dim if dim != kdim else kdim.replace("freq_", "") for dim in wave.dims
]
deriv = xr.DataArray(deriv, coords=da.coords, dims=newdims)
return deriv.transpose(*da.dims)
def test_dim_dft(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.dft(da, dim="y", shift=False).values,
xrft.dft(da, dim=["y"], shift=False).values,
)
assert xrft.dft(da, dim="y").dims == ("x", "freq_y")
def test_constant_coordinates():
"""error should be raised if coordinates are constant"""
N = 20
s = xr.DataArray(
np.random.rand(N) + 1j * np.random.rand(N),
dims="freq_x",
coords={"freq_x": np.zeros(N)},
)
with pytest.raises(ValueError):
xrft.dft(s)
with pytest.raises(ValueError):
xrft.idft(s)
def test_dft_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.dft(da, shift=False)
npt.assert_almost_equal(ft.values, np.fft.fftn(da.values))
ft = xrft.dft(da, shift=False, window=True, detrend='constant')
dim = da.dims
window = np.hanning(N) * np.hanning(N)[:, np.newaxis]
da_prime = (da - da.mean(dim=dim)).values
npt.assert_almost_equal(ft.values, np.fft.fftn(da_prime*window))
def test_idft_dft():
"""
Testing idft(dft(s)) == s
"""
N = 40
dx = np.random.rand()
s = xr.DataArray(
np.random.rand(N) + 1j * np.random.rand(N),
dims="x",
coords={
"x":
dx * (np.arange(-N // 2, -N // 2 + N) +
np.random.randint(-N // 2, N // 2))
},
)
FTs = xrft.dft(s, true_phase=True, true_amplitude=True)
mean_lag = float(s["x"][{
"x": s.sizes["x"] // 2
}]) # lag ensure IFTs to be on the same coordinate range than s
# lag is set manually
IFTs = xrft.idft(FTs,
shift=True,
true_phase=True,
true_amplitude=True,
lag=mean_lag)
xrt.assert_allclose(s, IFTs)
# lag is set automatically
IFTs = xrft.idft(FTs, shift=True, true_phase=True, true_amplitude=True)
xrt.assert_allclose(s, IFTs)
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] + .001
da3 = xr.DataArray(np.random.rand(Nx), coords=[x], dims=['x'])
# This shouldn't raise an error
xrft.dft(da3, spacing_tol=1e-1)
# But this should
with pytest.raises(ValueError):
xrft.dft(da3, spacing_tol=1e-4)
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.dft(da3, spacing_tol=1e-1)
# But this should
with pytest.raises(ValueError):
xrft.dft(da3, spacing_tol=1e-4)
def test_dft_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.dft(data, dim=["time"])
ps = xrft.power_spectrum(data, dim=["time"])
def test_dft_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.dft(da, shift=False)
npt.assert_almost_equal(ft.values, np.fft.fftn(da.values))
ft = xrft.dft(da, shift=False, window=True, detrend='constant')
dim = da.dims
window = np.hanning(N) * np.hanning(N)[:, np.newaxis]
da_prime = (da - da.mean(dim=dim)).values
npt.assert_almost_equal(ft.values, np.fft.fftn(da_prime * window))
def test_dft_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.dft(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.dft(da, dim=["x"])
with pytest.raises(ValueError):
xrft.dft(da, dim="x")
da = da.chunk({"time": N})
daft = xrft.dft(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.dft(da, dim="time", shift=False, detrend="linear"))
def test_real_dft_true_phase():
"""Test if real transform is half the total transform when signal is real and true_phase=True"""
Nx = 40
dx = np.random.rand()
s = xr.DataArray(
np.random.rand(Nx),
dims="x",
coords={
"x":
dx * (np.arange(-Nx // 2, -Nx // 2 + Nx) +
np.random.randint(-Nx // 2, Nx // 2))
},
)
s1 = xrft.dft(s, dim="x", true_phase=True, shift=True)
s2 = xrft.dft(s, real_dim="x", true_phase=True, shift=True)
s1 = np.conj(s1[{"freq_x": slice(None, s1.sizes["freq_x"] // 2 + 1)}])
s1 = s1.assign_coords(freq_x=-s1["freq_x"]).sortby("freq_x")
xrt.assert_allclose(s1, s2)
def test_dft_1d():
"""Test the discrete Fourier transform function on one-dimensional data."""
Nx = 16
Lx = 1.0
x = np.linspace(0, Lx, Nx)
dx = x[1] - x[0]
da = xr.DataArray(np.random.rand(Nx), coords=[x], dims=['x'])
# 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]
# 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)
# modify data to be non-evenly spaced
da2 = da.copy()
da2[-1] = np.nan
with pytest.raises(ValueError):
ft = xrft.dft(da2)
da['x'].values[-1] *= 2
with pytest.raises(ValueError):
ft = xrft.dft(da)
def test_dft_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)}
)
with pytest.raises(ValueError):
xrft.dft(da.chunk({'time':8}), dim=['y','x'], detrend='linear')
ft = xrft.dft(da, shift=False)
npt.assert_almost_equal(ft.values, np.fft.fftn(da.values))
da_prime = xrft.detrendn(da[:,0].values, [0,1,2]) # cubic detrend over time, y, and x
npt.assert_almost_equal(xrft.dft(da[:,0].drop('z'),
dim=['time','y','x'],
shift=False, detrend='linear'
).values,
np.fft.fftn(da_prime))
def test_dft_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)}
)
with pytest.raises(ValueError):
xrft.dft(da.chunk({'time':8}), dim=['y','x'], detrend='linear')
ft = xrft.dft(da, shift=False)
npt.assert_almost_equal(ft.values, np.fft.fftn(da.values))
da_prime = xrft.detrendn(da[:,0].values, [0,1,2]) # cubic detrend over time, y, and x
npt.assert_almost_equal(xrft.dft(da[:,0].drop('z'),
dim=['time','y','x'],
shift=False, detrend='linear'
).values,
np.fft.fftn(da_prime))
def periodogram_plot(self, plot_data, title):
dat = xrft.dft((plot_data - plot_data.mean()).chunk(
(plot_data - plot_data.mean()).size))
i = (np.multiply(dat, np.conj(dat)) / dat.size).real
i = np.log10(i[2:int(dat.size / 2) + 1])
freqs = np.array(range(1, int(dat.size / 2))) / dat.size
mpl.pyplot.subplots(1, 1, sharey=True, tight_layout=True)
mpl.pyplot.plot(freqs, i)
mpl.pyplot.title(f'periodogram: {title}')
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.dft(da.chunk(chunks=((20,N,N),(N-20,N,N))), dim=['time'],
detrend='linear', chunks_to_segments=True)
ft = xrft.dft(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.dft(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)).real.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.dft(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)).real.values,
)
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.dft(da.chunk(chunks=((20,N,N),(N-20,N,N))), dim=['time'],
detrend='linear', chunks_to_segments=True)
ft = xrft.dft(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.dft(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)).real.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.dft(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)).real.values,
)
def test_dft_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.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)
def test_dft_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.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)
def test_dft_1d_time(self):
"""Test the discrete Fourier transform function on timeseries data."""
time = pd.date_range('2000-01-01', '2001-01-01', closed='left')
Nt = len(time)
da = xr.DataArray(np.random.rand(Nt), coords=[time], dims=['time'])
ft = xrft.dft(da)
# check that frequencies are correct
dt = (time[1] - time[0]).total_seconds()
freq_time_expected = np.fft.fftshift(np.fft.fftfreq(Nt, dt))
npt.assert_allclose(ft['freq_time'], freq_time_expected)
def test_dft_1d_time(self):
"""Test the discrete Fourier transform function on timeseries data."""
time = pd.date_range('2000-01-01', '2001-01-01', closed='left')
Nt = len(time)
da = xr.DataArray(np.random.rand(Nt), coords=[time], dims=['time'])
ft = xrft.dft(da)
# check that frequencies are correct
dt = (time[1] - time[0]).total_seconds()
freq_time_expected = np.fft.fftshift(np.fft.fftfreq(Nt, dt))
npt.assert_allclose(ft['freq_time'], freq_time_expected)
def test_cross_spectrum_dask():
"""Test the power spectrum function on dask data"""
N = 16
dim = ['x', 'y']
da = xr.DataArray(np.random.rand(2, N, N),
dims=['time', 'x', 'y'],
coords={
'time': range(2),
'x': range(N),
'y': range(N)
}).chunk({'time': 1})
da2 = xr.DataArray(np.random.rand(2, N, N),
dims=['time', 'x', 'y'],
coords={
'time': range(2),
'x': range(N),
'y': range(N)
}).chunk({'time': 1})
cs = xrft.cross_spectrum(da, da2, dim=dim, density=False)
daft = xrft.dft(da, dim=dim)
daft2 = xrft.dft(da2, dim=dim)
npt.assert_almost_equal(cs.values, (daft * np.conj(daft2)).real.values)
cs = xrft.cross_spectrum(da,
da2,
dim=dim,
shift=True,
window=True,
detrend='constant')
daft = xrft.dft(da, dim=dim, shift=True, window=True, detrend='constant')
daft2 = xrft.dft(da2, dim=dim, shift=True, window=True, detrend='constant')
coord = list(daft.coords)
test = (daft * np.conj(daft2)).real.values / N**4
# for i in dim:
# test /= daft['freq_' + i + '_spacing']
dk = np.diff(np.fft.fftfreq(N, 1.))[0]
test /= dk**2
npt.assert_almost_equal(cs.values, test)
npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 0.)
def test_true_phase_preservation(chunk):
"""Test if dft is (phase-) preserved when signal is at same place but coords range is changed"""
x = np.arange(-15, 15)
y = np.random.rand(len(x))
N1 = np.random.randint(30) + 5
N2 = np.random.randint(30) + 5
l = np.arange(-N1, 0) + np.min(x)
r = np.arange(1, N2 + 1) + np.max(x)
s1 = xr.DataArray(
np.concatenate([np.zeros(N1), y, np.zeros(N2)]),
dims=("x", ),
coords={"x": np.concatenate([l, x, r])},
)
if chunk:
s1 = s1.chunk()
S1 = xrft.dft(s1, dim="x", true_phase=True)
assert s1.chunks == S1.chunks
N3 = N1
while N3 == N1:
N3 = np.minimum(np.random.randint(30), N1 + N2)
N4 = N1 + N2 - N3
l = np.arange(-N3, 0) + np.min(x)
r = np.arange(1, N4 + 1) + np.max(x)
s2 = xr.DataArray(
np.concatenate([np.zeros(N3), y, np.zeros(N4)]),
dims=("x", ),
coords={"x": np.concatenate([l, x, r])},
)
if chunk:
s2 = s2.chunk()
S2 = xrft.dft(s2, dim="x", true_phase=True)
assert s2.chunks == S2.chunks
xrt.assert_allclose(S1, S2)
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.dft(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.)
ps = xrft.power_spectrum(da, dim=['y'], real='x', window=True,
density=False, detrend='constant')
daft = xrft.dft(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.dft(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]
test /= dk**2
npt.assert_almost_equal(ps.values, test)
npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)
### Remove least-square fit
ps = xrft.power_spectrum(da, dim=['y','x'],
window=True, density=False, detrend='linear'
)
daft = xrft.dft(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.)
def test_spacing_tol_float_value(test_data_1d):
da = test_data_1d
with pytest.raises(TypeError):
xrft.dft(da, spacing_tol='string')
def test_dft_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.dft(data,dim=['time'])
ps = xrft.power_spectrum(data,dim=['time'])
