def test_detrend():
N = 16
x = np.arange(N + 1)
y = np.arange(N - 1)
t = np.linspace(-int(N / 2), int(N / 2), N - 6)
z = np.arange(int(N / 2))
d4d = (t[:, np.newaxis, np.newaxis, np.newaxis] +
z[np.newaxis, :, np.newaxis, np.newaxis] +
y[np.newaxis, np.newaxis, :, np.newaxis] +
x[np.newaxis, np.newaxis, np.newaxis, :])
da4d = xr.DataArray(d4d,
dims=['time', 'z', 'y', 'x'],
coords={
'time': range(len(t)),
'z': range(len(z)),
'y': range(len(y)),
'x': range(len(x))
})
func = xrft.detrend_wrap(xrft.detrendn)
#########
# Chunk along the `time` axis
#########
da = da4d.chunk({'time': 1})
with pytest.raises(ValueError):
func(da.data, axes=[0]).compute
with pytest.raises(ValueError):
func(da.data, axes=[0, 1, 2, 3]).compute()
da_prime = func(da.data, axes=[2]).compute()
npt.assert_allclose(da_prime[0, 0], sps.detrend(d4d[0, 0], axis=0))
da_prime = func(da.data, axes=[1, 2, 3]).compute()
npt.assert_allclose(da_prime[0], xrft.detrendn(d4d[0], axes=[0, 1, 2]))
#########
# Chunk along the `time` and `z` axes
#########
da = da4d.chunk({'time': 1, 'z': 1})
with pytest.raises(ValueError):
func(da.data, axes=[1, 2]).compute()
with pytest.raises(ValueError):
func(da.data, axes=[2, 2]).compute()
da_prime = func(da.data, axes=[2, 3]).compute()
npt.assert_allclose(da_prime[0, 0], xrft.detrendn(d4d[0, 0], axes=[0, 1]))
s = np.arange(2)
d5d = d4d[np.newaxis, :, :, :, :] + s[:, np.newaxis, np.newaxis,
np.newaxis, np.newaxis]
da5d = xr.DataArray(d5d,
dims=['s', 'time', 'z', 'y', 'x'],
coords={
's': range(len(s)),
'time': range(len(t)),
'z': range(len(z)),
'y': range(len(y)),
'x': range(len(x))
})
da = da5d.chunk({'time': 1})
with pytest.raises(ValueError):
func(da.data).compute()
def test_dft_4d():
"""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))
with pytest.raises(NotImplementedError):
xrft.dft(da, detrend='linear')
with pytest.raises(NotImplementedError):
xrft.dft(da, dim=['time', 'y', 'x'], detrend='linear')
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_detrend():
N = 16
x = np.arange(N+1)
y = np.arange(N-1)
t = np.linspace(-int(N/2), int(N/2), N-6)
z = np.arange(int(N/2))
d4d = (t[:,np.newaxis,np.newaxis,np.newaxis]
+ z[np.newaxis,:,np.newaxis,np.newaxis]
+ y[np.newaxis,np.newaxis,:,np.newaxis]
+ x[np.newaxis,np.newaxis,np.newaxis,:]
)
da4d = xr.DataArray(d4d, dims=['time','z','y','x'],
coords={'time':range(len(t)),'z':range(len(z)),'y':range(len(y)),
'x':range(len(x))}
)
func = xrft.detrend_wrap(xrft.detrendn)
#########
# Chunk along the `time` axis
#########
da = da4d.chunk({'time': 1})
with pytest.raises(ValueError):
func(da.data, axes=[0]).compute
with pytest.raises(ValueError):
func(da.data, axes=[0,1,2,3]).compute()
da_prime = func(da.data, axes=[2]).compute()
npt.assert_allclose(da_prime[0,0], sps.detrend(d4d[0,0], axis=0))
da_prime = func(da.data, axes=[1,2,3]).compute()
npt.assert_allclose(da_prime[0],
xrft.detrendn(d4d[0], axes=[0,1,2]))
#########
# Chunk along the `time` and `z` axes
#########
da = da4d.chunk({'time':1, 'z':1})
with pytest.raises(ValueError):
func(da.data, axes=[1,2]).compute()
with pytest.raises(ValueError):
func(da.data, axes=[2,2]).compute()
da_prime = func(da.data, axes=[2,3]).compute()
npt.assert_allclose(da_prime[0,0],
xrft.detrendn(d4d[0,0], axes=[0,1]))
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 _apply_detrend(da, axis_num):
"""Wrapper function for applying detrending"""
'''
copied from xrft.py -- not sure why but the functions
in xrft.py starting with '_' cannot be called directly!
So I just copied here directly -- look for the reason later!!!!
'''
if da.chunks:
func = xrft.detrend_wrap(xrft.detrendn)
da = xr.DataArray(func(da.data, axes=axis_num),
dims=da.dims,
coords=da.coords)
else:
if da.ndim == 1:
da = xr.DataArray(sps.detrend(da), dims=da.dims, coords=da.coords)
else:
da = xrft.detrendn(da, axes=axis_num)
# else:
# raise ValueError("Data should be dask array.")
return da
