# coding: utf-8

# ## MITgcm Example 
# 
# xgcm is developed in close coordination with the [xmitgcm](http://xmitgcm.readthedocs.io/) package.
# The metadata in datasets constructed by xmitgcm should always be compatible with xgcm's expectations.
# xmitgcm is necessary for reading MITgcm's binary MDS file format.
# However, for this example, the MDS files have already been converted and saved as netCDF.
# 
# Below are some example of how to make calculations on mitgcm-style datasets using xgcm.
# 
# First we import xarray and xgcm:

# In[ ]:


#! pip install git+https://github.com/xgcm/xgcm.git


# In[ ]:


import xarray as xr
import numpy as np
import xgcm
from matplotlib import pyplot as plt

get_ipython().run_line_magic('matplotlib', 'inline')
plt.rcParams['figure.figsize'] = (10,6)


# Now we open the example dataset, which is stored with the xgcm github repository in the datasets folder.

# In[ ]:


#!wget http://www.ldeo.columbia.edu/~rpa/mitgcm_example_dataset.nc


# In[ ]:


ds = xr.open_dataset('mitgcm_example_dataset.nc')
ds


# In[ ]:


ds.Eta


# In[ ]:


ds.Eta[0].plot.contourf(levels=30)


# In[ ]:


surf_mask_c = ds.hFacC[0] > 0
ds.Eta[0].where(surf_mask_c).plot.contourf(levels=30)


# ### Creating the grid object
# 
# Next we create a `Grid` object from the dataset.
# We need to tell xgcm that the `X` and `Y` axes are periodic.
# (The other axes will be assumed to be non-periodic.)

# In[ ]:


get_ipython().run_line_magic('pinfo', 'xgcm.Grid')


# In[ ]:


grid = xgcm.Grid(ds, periodic=['X', 'Y'])
grid


# In[ ]:


get_ipython().run_line_magic('pinfo', 'grid.interp')


# In[ ]:


ds.THETA.dims


# In[ ]:


theta_x_interp = grid.interp(ds.THETA, 'X')
theta_x_interp


# In[ ]:


theta_z_interp = grid.interp(ds.THETA, 'Z', boundary='extend')
theta_z_interp


# In[ ]:


ds.THETA.sel(XC=200, YC=45,
             method='nearest').plot.line(y='Z', marker='o',
                                         label='original')
theta_z_interp.sel(XC=200, YC=45,
                   method='nearest').plot.line(y='Zl', marker='o',
                                               label='interpolated')
plt.ylim([-500, 0])
plt.legend(loc='lower left')


# ### Kinetic Energy
# 
# Finally, we plot the kinetic energy $1/2 (u^2 + v^2)$ by interpoloting both quantities the cell center point.

# In[ ]:


# an example of calculating kinetic energy
ke = 0.5*(grid.interp((ds.U*ds.hFacW)**2, 'X') + 
          grid.interp((ds.V*ds.hFacS)**2, 'Y'))
print(ke)
ke[0,0].plot()


# ### Vorticity and Strain
# Here we compute more dervied quantities from the velocity field.
# 
# The vertical component of the vorticity is a fundamental quantity of interest in ocean circulation theory. It is defined as
# 
# $$ \zeta = - \frac{\partial u}{\partial y} + \frac{\partial v}{\partial x} \ . $$
# 
# On the c-grid, a finite-volume representation is given by
# 
# $$ \zeta =  (- \delta_j \Delta x_c u + \delta_i \Delta y_c v ) / A_\zeta \ . $$
# 
# In xgcm, we calculate this quanity as

# In[ ]:


zeta = (-grid.diff(ds.U * ds.dxC, 'Y') +
         grid.diff(ds.V * ds.dyC, 'X'))/ds.rAz
zeta


# ...which we can see is located at the `YG, XG` horizontal position (also commonly called the vorticity point).
# 
# We plot the vertical integral of this quantity, i.e. the barotropic vorticity:

# In[ ]:


zeta_bt = (zeta * ds.drF).sum(dim='Z')
zeta_bt.plot(vmax=2e-4)


# A different way to calculate the barotropic vorticity is to take the curl of the vertically integrated velocity.
# This formulation also allows us to incorporate the $h$ factors representing partial cell thickness.

# In[ ]:


u_bt = (ds.U * ds.hFacW * ds.drF).sum(dim='Z')
v_bt = (ds.V * ds.hFacS * ds.drF).sum(dim='Z')
zeta_bt_alt = (-grid.diff(u_bt * ds.dxC, 'Y') + grid.diff(v_bt * ds.dyC, 'X'))/ds.rAz
zeta_bt_alt.plot(vmax=2e-4)


# Another interesting quantity is the horizontal strain, defined as
# 
# $$ s = \frac{\partial u}{\partial x} - \frac{\partial v}{\partial y} \ . $$
# 
# On the c-grid, a finite-volume representation is given by
# 
# $$ s =  (\delta_i \Delta y_g u - \delta_j \Delta x_g v ) / A_c \ . $$

# In[ ]:


strain = (grid.diff(ds.U * ds.dyG, 'X') - grid.diff(ds.V * ds.dxG, 'Y')) / ds.rA
strain[0,0].plot()


# ### Barotropic Transport Streamfunction
# 
# We can use the barotropic velocity to calcuate the barotropic transport streamfunction, defined via
# 
# $$ u_{bt} = - \frac{\partial \Psi}{\partial y} \ , \ \ v_{bt} = \frac{\partial \Psi}{\partial x} \ .$$
# 
# We calculate this by integrating $u_{bt}$ along the Y axis using the grid object's `cumsum` method:

# In[ ]:


psi = grid.cumsum(-u_bt * ds.dyG, 'Y', boundary='fill')
psi


# We see that xgcm automatically shifted the Y-axis position from center (YC) to left (YG) during the cumsum operation.
# 
# We convert to sverdrups and plot with a contour plot.

# In[ ]:


(psi[0] / 1e6).plot.contourf(levels=np.arange(-160, 40, 5))


# ## Tracer Budget Example

# In[ ]:


adv_flux_div = (grid.diff(ds.ADVx_TH, 'X') +
                grid.diff(ds.ADVy_TH, 'Y') +
                grid.diff(ds.ADVr_TH, 'Z', boundary='fill'))
adv_flux_div


# In[ ]:


diff_flux_div = (grid.diff(ds.DFxE_TH, 'X') +
                grid.diff(ds.DFyE_TH, 'Y') +
                grid.diff(ds.DFrE_TH + ds.DFrI_TH, 'Z', boundary='fill'))
diff_flux_div


# In[ ]:


diff_flux_div.sum(dim=['XC', 'YC']).plot.line(y='Z', marker='.', label='diffusion')
adv_flux_div.sum(dim=['XC', 'YC']).plot.line(y='Z', marker='.', label='advection')
plt.grid()
plt.legend()