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)