def test_fft_warning():
"""Test if a warning is raise if a wrong dimension is used
"""
a = np.mgrid[:5, :5, :5][0]
dummy_array = xr.DataArray(a, dims=['x', 'y', 'z'])
chunked_array = dummy_array.chunk(chunks={'x': 2, 'y': 2, 'z': 2})
with pytest.warns(UserWarning):
xfft.fft(chunked_array, dim=['x', 'y', 'time'])
def compute():
print('Select dates')
u_JFM=dsu.sel(time_counter=slice('2010-01-01','2010-03-31'))['sozocrtx']
v_JFM=dsv.sel(time_counter=slice('2010-01-01','2010-03-31'))['somecrty']
print('Select box area')
u_JFM_box=u_JFM[:,jmin:jmax,imin:imax].chunk({'time_counter':10,'x':120,'y':120})
v_JFM_box=v_JFM[:,jmin:jmax,imin:imax].chunk({'time_counter':10,'x':120,'y':120})
# - get dx and dy
print('get dx and dy')
dx_JFM,dy_JFM = wfs.get_dx_dy(u_JFM_box[0],lonbox,latbox)
#... Detrend data in all dimension ...
print('Detrend data in all dimension')
u_JFM = wfs.detrendn(u_JFM_box,axes=[0,1,2])
v_JFM = wfs.detrendn(v_JFM_box,axes=[0,1,2])
#... Apply hanning windowing ...')
print('Apply hanning windowing')
u_JFM = wfs.apply_window(u_JFM, u_JFM.dims, window_type='hanning')
v_JFM = wfs.apply_window(v_JFM, v_JFM.dims, window_type='hanning')
# - get derivatives
derivatives_JFM = wfs.velocity_derivatives(u_JFM, v_JFM, xdim='x', ydim='y', dx={'x': dx_JFM, 'y': dy_JFM})
dudx_JFM = derivatives_JFM['u_x']; dudy_JFM = derivatives_JFM['u_y']
dvdx_JFM = derivatives_JFM['v_x']; dvdy_JFM = derivatives_JFM['v_y']
# - compute terms
phi1_JFM = u_JFM*dudx_JFM + v_JFM*dudy_JFM
phi2_JFM = u_JFM*dvdx_JFM + v_JFM*dvdy_JFM
u_JFMhat = xfft.fft(u_JFM, dim=('time_counter', 'x', 'y'), dx={'x': dx_JFM, 'y': dx_JFM}, sym=True)
v_JFMhat = xfft.fft(v_JFM, dim=('time_counter', 'x', 'y'), dx={'x': dx_JFM, 'y': dx_JFM}, sym=True)
phi1_JFM_hat = xfft.fft(phi1_JFM, dim=('time_counter', 'x', 'y'), dx={'x': dx_JFM, 'y': dx_JFM}, sym=True)
phi2_JFM_hat = xfft.fft(phi2_JFM, dim=('time_counter', 'x', 'y'), dx={'x': dx_JFM, 'y': dx_JFM}, sym=True)
tm1_JFM = (u_JFMhat.conj())*phi1_JFM_hat
tm2_JFM = (v_JFMhat.conj())*phi2_JFM_hat
# - computer transfer
Nk_JFM,Nj_JFM,Ni_JFM = u_JFM.shape
transfer_2D_JFM = -1.0*(tm1_JFM + tm2_JFM)/np.square(Ni_JFM*Nj_JFM)
transfer_term_JFM = transfer_2D_JFM.real
#... Get frequency and wavenumber ...
print('Get frequency and wavenumber')
ffrequency_JFM = u_JFMhat.f_time_counter
kx_JFM = u_JFMhat.f_x
ky_JFM = u_JFMhat.f_y
#... Get istropic wavenumber ...
print('Get istropic wavenumber')
wavenumber_JFM,kradial_JFM = wfs.get_wavnum_kradial(kx_JFM,ky_JFM)
#... Get numpy array ...
print('Get numpy array')
var_psd_np_JFM = transfer_term_JFM.values
#... Get 2D frequency-wavenumber field ...
print('Get transfer')
transfer_JFM = wfs.get_f_k_in_2D(kradial_JFM,wavenumber_JFM,var_psd_np_JFM)
print('Get flux')
flux_JFM = wfs.get_flux_in_1D(kradial_JFM,wavenumber_JFM,var_psd_np_JFM)
def velocity_derivatives(u, v, xdim, ydim, dx):
uhat = xfft.fft(u, dim=(xdim, ydim), dx=dx, sym=True)
vhat = xfft.fft(v, dim=(xdim, ydim), dx=dx, sym=True)
k = uhat['f_%s' % xdim]
l = vhat['f_%s' % ydim]
u_x_hat = uhat * 2 * np.pi * 1j * k
u_y_hat = uhat * 2 * np.pi * 1j * l
v_x_hat = vhat * 2 * np.pi * 1j * k
v_y_hat = vhat * 2 * np.pi * 1j * l
ds_derivatives = xr.Dataset({'u_x': xfft.ifft(u_x_hat),
'u_y': xfft.ifft(u_y_hat),
'v_x': xfft.ifft(v_x_hat),
'v_y': xfft.ifft(v_y_hat)
})
return ds_derivatives
def test_parserval_complex_1d():
"""Test if the Parseval theorem is verified"""
a = np.exp(2j * np.pi * np.arange(8) / 8)
dummy_array = xr.DataArray(a, dims=['x'])
chunked_array = dummy_array.chunk(chunks={'x': 2})
spec = xfft.fft(chunked_array, dim=['x'], detrend='mean')
assert np.var(a) == xfft.ps(spec).sum()
def test_parserval_real_1d():
"""Test if the Parseval theorem is verified"""
a = [0, 1, 0, 1, 1, 0, 1]
dummy_array = xr.DataArray(a, dims=['x'])
chunked_array = dummy_array.chunk(chunks={'x': 2})
spec = xfft.fft(chunked_array, dim=['x'], detrend='mean')
assert np.isclose(np.var(a), xfft.ps(spec).sum())
def test_parseval_real_2d():
"""Test if the Parseval theorem is verified"""
a = np.mgrid[:5, :5, :5][0]
dummy_array = xr.DataArray(a, dims=['x', 'y', 'z'])
chunked_array = dummy_array.chunk(chunks={'x': 2, 'y': 2, 'z': 2})
spec = xfft.fft(chunked_array, dim=['y', 'z'], detrend='mean')
assert np.array_equal(np.var(a, axis=(1, 2)),
xfft.ps(spec).sum(dim=['f_y', 'f_z']))
def test_parseval_real_2d():
"""Test if the Parseval theorem is verified"""
a = np.mgrid[:5, :5, :5][0]
dummy_array = xr.DataArray(a, dims=['x', 'y', 'z'])
chunked_array = dummy_array.chunk(chunks={'x': 2, 'y': 2, 'z': 2})
chunked_array_zeromean = chunked_array - chunked_array.mean(dim=['y', 'z'])
spec = xfft.fft(chunked_array_zeromean, dim=['y', 'z'])
assert np.allclose(np.var(a, axis=(1, 2)),
xfft.ps(spec).sum(dim=['f_y', 'f_z']))
def test_parserval_complex_2d():
""" Compare the result from the spectrum.fft function to
numpy.fft.fftn
"""
a, b, c = np.meshgrid([0, 1, 0, 0], [0, 1j, 1j], [0, 1, 1, 1])
dummy_array = xr.DataArray(a * b * c, dims=['x', 'y', 'z'])
chunked_array = dummy_array.chunk(chunks={'x': 2, 'y': 2, 'z': 2})
spec = xfft.fft(chunked_array, dim=['y', 'z'], detrend='mean')
assert np.array_equal(np.var(a * b * c, axis=(1, 2)),
xfft.ps(spec).sum(dim=['f_y','f_z']))
def test_parserval_complex_2d():
""" Compare the result from the spectrum.fft function to
numpy.fft.fftn
"""
a, b, c = np.meshgrid([1j, 1, 1, 1j], [0, 1j, 1j], [0, 1, 1, 1])
dummy_array = xr.DataArray(a * b * c, dims=['x', 'y', 'z'])
chunked_array = dummy_array.chunk(chunks={'x': 2, 'y': 2, 'z': 2})
chunked_array_zeromean = chunked_array - chunked_array.mean(dim=['y', 'z'])
spec = xfft.fft(chunked_array_zeromean, dim=['y', 'z'], sym=True)
assert np.allclose(np.var(a * b * c, axis=(1, 2)),
xfft.ps(spec).sum(dim=['f_y', 'f_z']))
def compute_wk(data, lon, lat):
# - get dx and dy
dx, dy = wfs.get_dx_dy(data[0], lon, lat)
#... Detrend data in all dimension ...
ssh_detrended = wfs.detrendn(data, axes=[0, 1, 2])
#... Apply hanning windowing ...')
ssh_hanning = wfs.apply_window(ssh_detrended,
data.dims,
window_type='hanning')
#... Apply hanning windowing ...')
ssh_hat = xfft.fft(ssh_hanning,
dim=('time', 'x', 'y'),
dx={
'x': dx,
'y': dx
},
sym=True)
#... Apply hanning windowing ...')
ssh_psd = xfft.psd(ssh_hat)
#... Get frequency and wavenumber ...
frequency = ssh_hat.f_time
kx = ssh_hat.f_x
ky = ssh_hat.f_y
#... Get istropic wavenumber ...
wavenumber, kradial = wfs.get_wavnum_kradial(kx, ky)
#... Get numpy array ...
ssh_psd_np = ssh_psd.values
#... Get 2D frequency-wavenumber field ...
SSH_wavenum_freq_spectrum = wfs.get_f_k_in_2D(kradial, wavenumber,
ssh_psd_np)
return wavenumber, frequency, SSH_wavenum_freq_spectrum
#... Apply hanning windowing ...')
print('Apply hanning windowing')
u_JFM = wfs.apply_window(u_JFM, u_JFM.dims, window_type='hanning')
v_JFM = wfs.apply_window(v_JFM, v_JFM.dims, window_type='hanning')
# - get derivatives
derivatives_JFM = wfs.velocity_derivatives(u_JFM, v_JFM, xdim='x', ydim='y', dx={'x': dx_JFM, 'y': dy_JFM})
dudx_JFM = derivatives_JFM['u_x']; dudy_JFM = derivatives_JFM['u_y']
dvdx_JFM = derivatives_JFM['v_x']; dvdy_JFM = derivatives_JFM['v_y']
# - compute terms
phi1_JFM = u_JFM*dudx_JFM + v_JFM*dudy_JFM
phi2_JFM = u_JFM*dvdx_JFM + v_JFM*dvdy_JFM
u_JFMhat = xfft.fft(u_JFM, dim=('time_counter', 'x', 'y'), dx={'x': dx_JFM, 'y': dx_JFM}, sym=True)
v_JFMhat = xfft.fft(v_JFM, dim=('time_counter', 'x', 'y'), dx={'x': dx_JFM, 'y': dx_JFM}, sym=True)
phi1_JFM_hat = xfft.fft(phi1_JFM, dim=('time_counter', 'x', 'y'), dx={'x': dx_JFM, 'y': dx_JFM}, sym=True)
phi2_JFM_hat = xfft.fft(phi2_JFM, dim=('time_counter', 'x', 'y'), dx={'x': dx_JFM, 'y': dx_JFM}, sym=True)
tm1_JFM = (u_JFMhat.conj())*phi1_JFM_hat
tm2_JFM = (v_JFMhat.conj())*phi2_JFM_hat
# - computer transfer
Nk_JFM,Nj_JFM,Ni_JFM = u_JFM.shape
transfer_2D_JFM = -1.0*(tm1_JFM + tm2_JFM)/np.square(Ni_JFM*Nj_JFM)
transfer_term_JFM = transfer_2D_JFM.real
#... Get frequency and wavenumber ...
print('Get frequency and wavenumber')
dx_JFM, dy_JFM = wfs.get_dx_dy(u_JFM_box[0], lonbox, latbox)
#... Detrend data in all dimension ...
print('Detrend data in all dimension')
u_JFM = wfs.detrendn(u_JFM_box, axes=[0, 1, 2])
#... Apply hanning windowing ...')
print('Apply hanning windowing')
u_JFM = wfs.apply_window(u_JFM, u_JFM.dims, window_type='hanning')
#... Apply hanning windowing ...')
print('FFT ')
u_JFMhat = xfft.fft(u_JFM,
dim=('time_counter', 'x', 'y'),
dx={
'x': dx_JFM,
'y': dx_JFM
},
sym=True)
#... Apply hanning windowing ...')
print('PSD ')
u_JFM_psd = xfft.psd(u_JFMhat)
#... Get frequency and wavenumber ...
print('Get frequency and wavenumber')
frequency_JFM = u_JFMhat.f_time_counter
kx_JFM = u_JFMhat.f_x
ky_JFM = u_JFMhat.f_y
#... Get istropic wavenumber ...
'x': dx,
'y': dy
})
dudx = derivatives['u_x']
dudy = derivatives['u_y']
dvdx = derivatives['v_x']
dvdy = derivatives['v_y']
# - compute terms
phi1 = u * dudx + v * dudy
phi2 = u * dvdx + v * dvdy
uhat = xfft.fft(u,
dim=('time_counter', 'x', 'y'),
dx={
'x': dx,
'y': dx
},
sym=True)
vhat = xfft.fft(v,
dim=('time_counter', 'x', 'y'),
dx={
'x': dx,
'y': dx
},
sym=True)
phi1_hat = xfft.fft(phi1,
dim=('time_counter', 'x', 'y'),
dx={
'x': dx,
def test_spectrum_2d(tapering):
a = np.mgrid[:5, :5, :5][0]
dummy_array = xr.DataArray(a, dims=['x', 'y', 'z'])
chunked_array = dummy_array.chunk(chunks={'x': 2, 'y': 2, 'z': 2})
xfft.fft(chunked_array, dim=['y', 'z'], tapering=tapering).load()
def test_spectrum_1d(tapering):
a = [0, 1, 0, 0]
dummy_array = xr.DataArray(a, dims=['x'])
chunked_array = dummy_array.chunk(chunks={'x': 2})
xfft.fft(chunked_array, dim=['x'], tapering=tapering).load()
#... Detrend data in all dimension ...
print('Detrend data in all dimension')
u_JAS = wfs.detrendn(u_JAS_box, axes=[0, 1, 2])
v_JAS = wfs.detrendn(v_JAS_box, axes=[0, 1, 2])
#... Apply hanning windowing ...')
print('Apply hanning windowing')
u_JAS = wfs.apply_window(u_JAS, u_JAS.dims, window_type='hanning')
v_JAS = wfs.apply_window(v_JAS, v_JAS.dims, window_type='hanning')
#... Apply hanning windowing ...')
print('FFT ')
u_JAShat = xfft.fft(u_JAS,
dim=('time_counter', 'x', 'y'),
dx={
'x': dx_JAS,
'y': dx_JAS
},
sym=True)
v_JAShat = xfft.fft(v_JAS,
dim=('time_counter', 'x', 'y'),
dx={
'x': dx_JAS,
'y': dx_JAS
},
sym=True)
#... Apply hanning windowing ...')
print('PSD ')
u_JAS_psd = xfft.psd(u_JAShat)
v_JAS_psd = xfft.psd(v_JAShat)
