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 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
lonbox=lon[jmin:jmax,imin:imax] print('Select dates') u_JFM=dsu.sel(time_counter=slice('2013-01-01','2013-03-31'))['vozocrtx'] v_JFM=dsv.sel(time_counter=slice('2013-01-01','2013-03-31'))['vomecrty'] print('Select box area') u_JFM_box=u_JFM[:,jmin:jmax,imin:imax].chunk(chunks={'time_counter':-1,'x':-1,'y':-1}) v_JFM_box=v_JFM[:,jmin:jmax,imin:imax].chunk(chunks={'time_counter':-1,'x':-1,'y':-1}) # - 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']
print('Select dates') u_1 = dsu1['sozocrtx'] u_2 = dsu2['sozocrtx'] u_3 = dsu3['sozocrtx'] u_JFM = xr.concat([u_1, u_2, u_3], dim='time_counter') print('Select box area') u_JFM_box = u_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]) #... 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,
################### Main computation ###################### u = xr.open_mfdataset(Ufile, chunks={ 'x': 100, 'time_counter': 1 })['vozocrtx'][:, box.jmin:box.jmax, box.imin:box.imax] v = xr.open_mfdataset(Vfile, chunks={ 'x': 100, 'time_counter': 1 })['vomecrty'][:, box.jmin:box.jmax, box.imin:box.imax] # - remove NaN u = u.interpolate_na(dim='y') v = v.interpolate_na(dim='y') # - get dx and dy dx, dy = wfs.get_dx_dy(u[0]) #... Detrend data in all dimension ... print('Detrend data in all dimension') u = wfs.detrendn(u, axes=[0, 1, 2]) v = wfs.detrendn(v, axes=[0, 1, 2]) #... Apply hanning windowing ...') print('Apply hanning windowing') u = wfs.apply_window(u, u.dims, window_type='hanning') v = wfs.apply_window(v, v.dims, window_type='hanning') # - get derivatives derivatives = wfs.velocity_derivatives(u, v, xdim='x',
print('Select box area') u_JAS_box = u_JAS[:, jmin:jmax, imin:imax].chunk({ 'time_counter': -1, 'x': 120, 'y': 120 }) v_JAS_box = v_JAS[:, jmin:jmax, imin:imax].chunk({ 'time_counter': -1, 'x': 120, 'y': 120 }) # - get dx and dy print('get dx and dy') dx_JAS, dy_JAS = wfs.get_dx_dy(u_JAS_box[0], lonbox, latbox) #... 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'),