{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MITgcm Example \n",
    "\n",
    "xgcm is developed in close coordination with the [xmitgcm](http://xmitgcm.readthedocs.io/) package.\n",
    "The metadata in datasets constructed by xmitgcm should always be compatible with xgcm's expectations.\n",
    "xmitgcm is necessary for reading MITgcm's binary MDS file format.\n",
    "However, for this example, the MDS files have already been converted and saved as netCDF.\n",
    "\n",
    "Below are some example of how to make calculations on mitgcm-style datasets using xgcm.\n",
    "\n",
    "First we import xarray and xgcm:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#! pip install git+https://github.com/xgcm/xgcm.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "import xgcm\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "%matplotlib inline\n",
    "plt.rcParams['figure.figsize'] = (10,6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we open the example dataset, which is stored with the xgcm github repository in the datasets folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!wget http://www.ldeo.columbia.edu/~rpa/mitgcm_example_dataset.nc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir(xgcm.grid.Grid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ds = xr.open_dataset('mitgcm_example_dataset.nc')\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.Eta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.Eta[0].plot.contourf(levels=30)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "surf_mask_c = ds.hFacC[0] > 0\n",
    "ds.Eta[0].where(surf_mask_c).plot.contourf(levels=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating the grid object\n",
    "\n",
    "Next we create a `Grid` object from the dataset.\n",
    "We need to tell xgcm that the `X` and `Y` axes are periodic.\n",
    "(The other axes will be assumed to be non-periodic.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xgcm.Grid?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.XC.attrs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid = xgcm.Grid(ds, periodic=['X', 'Y'])\n",
    "grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid.interp?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.THETA.dims"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta_x_interp = grid.interp(ds.THETA, 'X')\n",
    "theta_x_interp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.THETA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta_z_interp = grid.interp(ds.THETA, 'Z', boundary='extend')\n",
    "theta_z_interp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.THETA[0,5].plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta_x_interp[0,5,].plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.THETA.sel(XC=200, YC=45,\n",
    "             method='nearest').plot.line(y='Z', marker='o',\n",
    "                                         label='original')\n",
    "theta_z_interp.sel(XC=200, YC=45,\n",
    "                   method='nearest').plot.line(y='Zl', marker='o',\n",
    "                                               label='interpolated')\n",
    "# plt.ylim([-500, 0])\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Kinetic Energy\n",
    "\n",
    "Finally, we plot the kinetic energy $1/2 (u^2 + v^2)$ by interpoloting both quantities the cell center point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# an example of calculating kinetic energy\n",
    "ke = 0.5*(grid.interp((ds.U*ds.hFacW)**2, 'X') + \n",
    "          grid.interp((ds.V*ds.hFacS)**2, 'Y'))\n",
    "print(ke)\n",
    "ke[0,0].plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vorticity and Strain\n",
    "Here we compute more dervied quantities from the velocity field.\n",
    "\n",
    "The vertical component of the vorticity is a fundamental quantity of interest in ocean circulation theory. It is defined as\n",
    "\n",
    "$$ \\zeta = - \\frac{\\partial u}{\\partial y} + \\frac{\\partial v}{\\partial x} \\ . $$\n",
    "\n",
    "On the c-grid, a finite-volume representation is given by\n",
    "\n",
    "$$ \\zeta =  (- \\delta_j \\Delta x_c u + \\delta_i \\Delta y_c v ) / A_\\zeta \\ . $$\n",
    "\n",
    "In xgcm, we calculate this quanity as"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "zeta = (-grid.diff(ds.U * ds.dxC, 'Y') +\n",
    "         grid.diff(ds.V * ds.dyC, 'X'))/ds.rAz\n",
    "zeta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...which we can see is located at the `YG, XG` horizontal position (also commonly called the vorticity point).\n",
    "\n",
    "We plot the vertical integral of this quantity, i.e. the barotropic vorticity:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "zeta_bt = (zeta * ds.drF).sum(dim='Z')\n",
    "zeta_bt.plot(vmax=2e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A different way to calculate the barotropic vorticity is to take the curl of the vertically integrated velocity.\n",
    "This formulation also allows us to incorporate the $h$ factors representing partial cell thickness."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_bt = (ds.U * ds.hFacW * ds.drF).sum(dim='Z')\n",
    "v_bt = (ds.V * ds.hFacS * ds.drF).sum(dim='Z')\n",
    "zeta_bt_alt = (-grid.diff(u_bt * ds.dxC, 'Y') + grid.diff(v_bt * ds.dyC, 'X'))/ds.rAz\n",
    "zeta_bt_alt.plot(vmax=2e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another interesting quantity is the horizontal strain, defined as\n",
    "\n",
    "$$ s = \\frac{\\partial u}{\\partial x} - \\frac{\\partial v}{\\partial y} \\ . $$\n",
    "\n",
    "On the c-grid, a finite-volume representation is given by\n",
    "\n",
    "$$ s =  (\\delta_i \\Delta y_g u - \\delta_j \\Delta x_g v ) / A_c \\ . $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "strain = (grid.diff(ds.U * ds.dyG, 'X') - grid.diff(ds.V * ds.dxG, 'Y')) / ds.rA\n",
    "strain[0,0].plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Barotropic Transport Streamfunction\n",
    "\n",
    "We can use the barotropic velocity to calcuate the barotropic transport streamfunction, defined via\n",
    "\n",
    "$$ u_{bt} = - \\frac{\\partial \\Psi}{\\partial y} \\ , \\ \\ v_{bt} = \\frac{\\partial \\Psi}{\\partial x} \\ .$$\n",
    "\n",
    "We calculate this by integrating $u_{bt}$ along the Y axis using the grid object's `cumsum` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi = grid.cumsum(-u_bt * ds.dyG, 'Y', boundary='fill')\n",
    "psi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that xgcm automatically shifted the Y-axis position from center (YC) to left (YG) during the cumsum operation.\n",
    "\n",
    "We convert to sverdrups and plot with a contour plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(psi[0] / 1e6).plot.contourf(levels=np.arange(-160, 40, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tracer Budget Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adv_flux_div = (grid.diff(ds.ADVx_TH, 'X') +\n",
    "                grid.diff(ds.ADVy_TH, 'Y') +\n",
    "                grid.diff(ds.ADVr_TH, 'Z', boundary='fill'))\n",
    "adv_flux_div"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diff_flux_div = (grid.diff(ds.DFxE_TH, 'X') +\n",
    "                grid.diff(ds.DFyE_TH, 'Y') +\n",
    "                grid.diff(ds.DFrE_TH + ds.DFrI_TH, 'Z', boundary='fill'))\n",
    "diff_flux_div"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diff_flux_div.sum(dim=['XC', 'YC']).plot.line(y='Z', marker='.', label='diffusion')\n",
    "adv_flux_div.sum(dim=['XC', 'YC']).plot.line(y='Z', marker='.', label='advection')\n",
    "plt.grid()\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "gist_info": {
   "gist_id": null,
   "gist_url": null
  },
  "kernelspec": {
   "display_name": "Py3 pangeo",
   "language": "python",
   "name": "pangeo"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}