<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Realistic example using outputs from MITgcm
This example requires the understanding of xgcm.grid and xmitgcm.open_mdsdataset.
Step1:  One year of daily-averaged output from MITgcm.
Step2:  Discrete Fourier Transform
We chunk the data along the time and Z axes to allow parallelized computation and detrend and window the data before taking the DFT along the horizontal axes.
Step3:  Power spectrum
We compute the surface eddy kinetic energy spectrum.
Step4:  Isotropic wavenumber spectrum
We now isotropize the spectrum
Step5:  We plot $u$, $v$, $\hat{u}^2+$
Step6:  Cross Spectrum
We calculate the cross correlation between vertical velocity ($w$) and buoyancy ($b$)
<ASSISTANT_TASK:>
Python Code:
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
ddir = '/swot/SUM05/takaya/MITgcm/channel/runs/'
Explanation: Realistic example using outputs from MITgcm
This example requires the understanding of xgcm.grid and xmitgcm.open_mdsdataset.
End of explanation
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
grid = Grid(ds, periodic=['X'])
u = ds.UVEL     #zonal velocity
v = ds.VVEL     #meridional velocity
w = ds.WVEL     #vertical velocity
phi = ds.PHIHYD #hydrostatic pressure
Explanation: One year of daily-averaged output from MITgcm.
End of explanation
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
Explanation: Discrete Fourier Transform
We chunk the data along the time and Z axes to allow parallelized computation and detrend and window the data before taking the DFT along the horizontal axes.
End of explanation
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
Explanation: Power spectrum
We compute the surface eddy kinetic energy spectrum.
End of explanation
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
Explanation: Isotropic wavenumber spectrum
We now isotropize the spectrum:
End of explanation
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}$]')
Explanation: We plot $u$, $v$, $\hat{u}^2+$
End of explanation
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
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)
Explanation: Cross Spectrum
We calculate the cross correlation between vertical velocity ($w$) and buoyancy ($b$):
End of explanation
<END_TASK>