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'),