This example requires the understanding of xgcm.grid and xmitgcm.open_mdsdataset.
In [1]:
import numpy as np
import xarray as xr
import os.path as op
import xrft
from dask.diagnostics import ProgressBar
from xmitgcm import open_mdsdataset
from xgcm.grid import Grid
from matplotlib import colors, ticker
import matplotlib.pyplot as plt
%matplotlib inline
In [2]:
ddir = '/swot/SUM05/takaya/MITgcm/channel/runs/'
One year of daily-averaged output from MITgcm.
In [3]:
ys20, dy20 = (60,1)
dt = 8e2
df = 108
ts = int(360*86400*ys20/dt+df)
te = int(360*86400*(ys20+dy20)/dt+df)
ds = open_mdsdataset(op.join(ddir,'zerores_20km_MOMbgc'), grid_dir=op.join(ddir,'20km_grid'),
iters=range(ts,te,df), prefix=['MOMtave'], delta_t=dt
).sel(YC=slice(5e5,15e5), YG=slice(5e5,15e5))
ds
Out[3]:
In [4]:
grid = Grid(ds, periodic=['X'])
In [5]:
u = ds.UVEL #zonal velocity
v = ds.VVEL #meridional velocity
w = ds.WVEL #vertical velocity
phi = ds.PHIHYD #hydrostatic pressure
In [6]:
b = grid.diff(phi,'Z',boundary='fill')/grid.diff(phi.Z,'Z',boundary='fill')
with ProgressBar():
what = xrft.dft(w.chunk({'time':1,'Zl':1}),
dim=['XC','YC'], detrend='linear', window=True).compute()
bhat = xrft.dft(b.chunk({'time':1,'Zl':1}),
dim=['XC','YC'], detrend='linear', window=True).compute()
bhat
Out[6]:
In [8]:
with ProgressBar():
uhat2 = xrft.power_spectrum(grid.interp(u,'X')[:,0].chunk({'time':1}),
dim=['XC','YC'], detrend='linear', window=True).compute()
vhat2 = xrft.power_spectrum(grid.interp(v,'Y',boundary='fill')[:,0].chunk({'time':1}),
dim=['XC','YC'], detrend='linear', window=True).compute()
ekehat = .5*(uhat2 + vhat2)
ekehat
Out[8]:
In [11]:
with ProgressBar():
uiso2 = xrft.isotropic_powerspectrum(grid.interp(u,'X')[0,0],
dim=['XC','YC'], detrend='linear', window=True).compute()
viso2 = xrft.isotropic_powerspectrum(grid.interp(v,'Y',boundary='fill')[0,0],
dim=['XC','YC'], detrend='linear', window=True).compute()
ekeiso = .5*(uiso2 + viso2)
ekeiso
Out[11]:
We plot $u$, $v$, $\hat{u}^2+$
In [22]:
fig, axes = plt.subplots(nrows=1, ncols=4, figsize=(20,4))
fig.set_tight_layout(True)
u[0,0].plot(ax=axes[0])
v[0,0].plot(ax=axes[1])
im = axes[2].pcolormesh(ekehat.freq_XC*1e3, ekehat.freq_YC*1e3, ekehat[0],
norm=colors.LogNorm())
axes[3].plot(ekeiso.freq_r*1e3, ekeiso)
cbar = fig.colorbar(im, ax=axes[2])
cbar.set_label(r'[m$^2$ s$^{-2}$]')
axes[3].set_xscale('log')
axes[3].set_yscale('log')
axes[2].set_xlabel(r'k [cpkm]')
axes[2].set_ylabel(r'l [cpkm]')
axes[3].set_xlabel(r'k$_r$ [cpkm]')
axes[3].set_ylabel(r'[m$^3$ s$^{-2}$]')
Out[22]:
In [31]:
with ProgressBar():
whatbhat = xrft.cross_spectrum(w.chunk({'time':1,'Zl':1}), b.chunk({'time':1,'Zl':1}),
dim=['XC','YC'], detrend='linear', window=True, density=False).compute()
whatbhat
Out[31]:
In [32]:
fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, figsize=(11,4))
fig.set_tight_layout(True)
(what*np.conjugate(bhat)).real[:,:8].mean(['time','Zl']).plot(ax=ax1)
whatbhat[:,:8].mean(['time','Zl']).plot(ax=ax2)
Out[32]:
We see that $\hat{w}\hat{b}^*$ and xrft.cross_spectrum$(w,b)$ are equivalent.
In [ ]: