{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Python notebook: Analysis of Tutorial Reentrant Channel Output\n",
    "\n",
    "We require standard python packages ``numpy`` and ``matplotlib``,\n",
    "and use package ``MITgcmutils`` (as described in chapter 11 of the MITgcm users guide)\n",
    "to load model grid and diagnostic output data (in binary format). Users should feel free to substitute\n",
    "``xmitgcm`` (and ``xarray``) for this purpose if they are comfortable with these packages.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.patches import Rectangle # used to shade land cells\n",
    "from MITgcmutils import diagnostics      # used to read in statistical diagnostics output\n",
    "from MITgcmutils import mds\n",
    "import cmocean                           # uses cmocean colormaps https://matplotlib.org/cmocean/\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, load up grid information. Note there is no straightforward way to load the `pkg/layers` temperature coordinate, so we copy the data from `data.layers` configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = '../run/'                         # point to run directory\n",
    "XC = mds.rdmds(path + 'XC')\n",
    "YC = mds.rdmds(path + 'YC')\n",
    "XG = mds.rdmds(path + 'XG')\n",
    "YG = mds.rdmds(path + 'YG')\n",
    "RC = mds.rdmds(path + 'RC')\n",
    "RF = mds.rdmds(path + 'RF')\n",
    "DXC = mds.rdmds(path + 'DXC')\n",
    "DYC = mds.rdmds(path + 'DYC')\n",
    "DXG = mds.rdmds(path + 'DXG')\n",
    "DYG = mds.rdmds(path + 'DYG')\n",
    "DRF = mds.rdmds(path + 'DRF')\n",
    "DRC = mds.rdmds(path + 'DRC')\n",
    "RAC = mds.rdmds(path + 'RAC')\n",
    "hFacC = mds.rdmds(path + 'hFacC')\n",
    "hFacW = mds.rdmds(path + 'hFacW')\n",
    "hFacS = mds.rdmds(path + 'hFacS')\n",
    "\n",
    "# create Nx+1,Ny+1 sized x,y vectors that include the eastern and northern\n",
    "# domain edge locations respectively\n",
    "Xp1 = np.append(XG[0,:], XG[0,-1] + DXG[0,-1])\n",
    "Yp1 = np.append(YG[:,0], YG[-1,0] + DYG[-1,0])\n",
    "\n",
    "# as specified in data.layers configuration\n",
    "layer_bounds = np.array([-2.00, -1.75, -1.50, -1.25,\n",
    "                         -1.00, -0.75, -0.50, -0.25,\n",
    "                          0.00,  0.25,  0.50,  0.75,\n",
    "                          1.00,  1.25,  1.50,  1.75,\n",
    "                          2.00,  2.25,  2.50,  2.75,\n",
    "                          3.00,  3.25,  3.50,  3.75,\n",
    "                          4.00,  4.25,  4.50,   5.0,\n",
    "                          5.5,    6.0,   6.5,   7.0,\n",
    "                          7.5,    8.0,   8.5,   9.0,\n",
    "                          9.5,   10.0,])\n",
    "\n",
    "# extract overall MITgcm domain dimensions\n",
    "Ny = XC.shape[0]\n",
    "Nx = XC.shape[1]\n",
    "Nr = RC.size\n",
    "Nr_layers = layer_bounds.size - 1\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "The hyperbolic tangent method of Stewart et al. (2017) was used to generate an optimal vertical grid: https://github.com/kialstewart/vertical_grid_for_ocean_models\n",
    "The vertical spacing is plotted below; near the surface, the vertical grid is quite refined,\n",
    "with a uppermost cell thickness of 5.5 m, whereas the deepest cell is 150 m thick.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32648 characters]\n",
      "text/plain": [
       "<Figure size 1008x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# plot vertical grid spacing\n",
    "\n",
    "plt.figure(figsize=(14,5))\n",
    "plt.subplot(1,2,1)\n",
    "plt.plot(DRF[:,0,0], RC[:,0,0], '-o')\n",
    "plt.grid()\n",
    "plt.title('Vertical grid spacing')\n",
    "plt.xlabel('Grid spacing (m)')\n",
    "plt.ylabel('Depth (m)')\n",
    "plt.subplot(1,2,2)\n",
    "plt.plot(RF[:,0,0], '-o')\n",
    "plt.grid()\n",
    "plt.title('Vertical layer face depths')\n",
    "plt.xlabel('Level')\n",
    "plt.ylabel('Depth (m)'); \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The bathymetry consists of a reentrant channel in the x-coordinate, including a deep ridge running N-S\n",
    "with a sloping notch cut out in the middle. A wall (\"land cells\") exists at the southernmost row of \n",
    "the domain (row of wall cells not plotted in figure below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 173956 characters]\n",
      "text/plain": [
       "<Figure size 720x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#plot bathymetry as 3D surface plot, as shown in fig 4.10\n",
    "\n",
    "bathy = np.fromfile(path + 'bathy.50km.bin', dtype='>f4')\n",
    "bathy.shape = (Ny,Nx)\n",
    "plt.figure(figsize=(10,8)) \n",
    "ax = plt.axes(projection='3d')\n",
    "ax.plot_surface(XC[1:,:]/1000, YC[1:,:]/1000, bathy[1:,:], cmap='copper', edgecolor='none');\n",
    "ax.view_init(40, -50)\n",
    "ax.set_zlim(-4000,0)\n",
    "ax.set_zlabel('Z (m)', fontsize=16)\n",
    "ax.set_title('Tutorial SO Reentrant Channel Bathymetry', fontsize=18)\n",
    "plt.xlabel('x (km)',fontsize=16); plt.ylabel('y (km)', fontsize=16)\n",
    "plt.xlim(0,1000)\n",
    "plt.ylim(0,2000);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the surface layer, temperature is restored to a linear profile:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 19312 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# SST linear restoring temp\n",
    "\n",
    "sst_relax = np.fromfile(path + 'SST_relax.50km.bin', dtype='>f4')\n",
    "sst_relax.shape = (Ny,Nx)\n",
    "plt.plot(YC[:,0]/1000, sst_relax[:,0], linewidth=3)\n",
    "plt.grid()\n",
    "plt.ylabel('Temperature ($^{\\circ}$C)')\n",
    "plt.xlabel('Y (km)')\n",
    "plt.title('SST Restoring Temperature', fontsize=18);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The vertical profile of initial conditions in temperature at the northern and southern ends of the domain is shown.\n",
    "Note that during the run, temperature is restored back to these conditions along the northen boundary.\n",
    "\n",
    "The vertical profile is specified as follows (where $T_{max}$ is the restoring SST(y)\n",
    "and $T_{min}$ = -2 $^{\\circ}$C), and $H$ is the maximum ocean depth):\n",
    "\n",
    "\n",
    "$$ T(z) = (T_{max} - T_{min}) \\frac{exp(z/h) - exp(-H/h)}{1 - exp(-H/h)} + T_{min} $$ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 25440 characters]\n",
      "text/plain": [
       "<Figure size 720x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Temperature initial conditions, as shown in  fig 4.11\n",
    "\n",
    "T_3D = np.fromfile(path + 'temperature.50km.bin', dtype='>f4')\n",
    "T_3D.shape = (Nr,Ny,Nx)\n",
    "plt.figure(figsize=(10,4))\n",
    "plt.plot(T_3D[:,1,1], RC[:,0,0], linewidth=3, label='Southern end')\n",
    "plt.plot(T_3D[:,-1,1], RC[:,0,0], linewidth=3, label='Northern end')\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.xlabel('Temperature ($^{\\circ}$C)')\n",
    "plt.ylabel('Depth (m)')\n",
    "plt.title('Temperature: Initial Conditions', fontsize=18);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The \"mask\" used for 3D temperature restoring (`pkg/rbcs`) is plotted below (showing YZ plane). Note that we are not\n",
    "using `pkg/rbcs` in the surface layer (`T_msk[0,:,:]=0`), instead the SST is restored through\n",
    "options set in file `data`. Mask values range between 0 and 1, with 0=no restoring, 1=full restoring with timescale prescibed in `data.rbcs`, with fractional strength in between."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 16484 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# restoring mask in yz plane, showing restoring along northern boundary\n",
    "\n",
    "T_msk = np.fromfile(path + 'T_relax_mask.50km.bin',dtype='>f4')\n",
    "T_msk.shape = (Nr,Ny,Nx)\n",
    "# just shade where restoring occurs (non-zero mask), white figure background = no restoring\n",
    "plt.pcolormesh(Yp1/1000, RF, np.ma.array(T_msk[:,:,0], mask=(T_msk[:,:,0]==0)), cmap='Reds', vmin=0)\n",
    "plt.title('Temperature Restoring Mask YZ plane', fontsize=18)\n",
    "plt.colorbar()\n",
    "plt.xlabel('Y (km)')\n",
    "plt.ylabel('Depth (m)');\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Applied surface wind stress in the x-direction\n",
    "\n",
    "$$ \\tau_x(y) = \\tau_{0}\\sin \\left( \\frac{y}{2 L_y} \\pi \\right) $$\n",
    "\n",
    "is as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 22048 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# wind stress profile\n",
    "\n",
    "taux = np.fromfile(path + 'zonal_wind.50km.bin',dtype='>f4')\n",
    "taux.shape = (Ny,Nx)\n",
    "plt.plot(taux[:,0], YC[:,0]/1000, linewidth=3)\n",
    "plt.grid()\n",
    "plt.xlabel('N/m$^2$')\n",
    "plt.ylabel('Y (km)')\n",
    "plt.title(r'Wind stress $\\tau_x$',fontsize=18);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load diagnostic data\n",
    "\n",
    "Here we have separate directories labelled ``run`` and ``run_noGM`` and we load up \n",
    "sets of statistical diagnostics and  diagnostics of (annual mean) year 30 data (corresponding to iteration number 933120) from each.\n",
    "\n",
    "Diagnostics output (as specified in ``data.diagnostics``) consists of:\n",
    "\n",
    "  state[0,1,2,3,4] = ``THETA``, ``VVEL``, ``UVEL``, ``WVEL``, ``CONVADJ`` <br>\n",
    "  diag2D[0,1,2] = ``TRELAX``, ``MXDEPTH``, ``ETAN`` <br>\n",
    "  gm_diags[0,1] = ``GM_PsiX``, ``GM_PsiY`` <br>\n",
    "  laydiag[0,1,2] = ``LaVH1TH``, ``LaHs1TH``, ``LaVa1TH`` <br>\n",
    "  \n",
    "Statistical diagnostics are loaded into a dictionary-like object which can be accessed\n",
    "by key ``THETA`` or ``TRELAX``: \n",
    " - statdiags_2D is dimensioned (iter, [ave,std,min,max,vol])\n",
    " - statdiags_bylev is dimensioned (iter, kLev, [ave,std,min,max,vol]).\n",
    "\n",
    "See output file ``available_diagnostics.log`` for a description of these diagnostics, including units.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose iteration number corresponding to year 30 (annual average)\n",
    "\n",
    "iter = 933120\n",
    "\n",
    "state_wGM = mds.rdmds(path + 'Diags/state', iter)\n",
    "diag2D_wGM = mds.rdmds(path + 'Diags/2D_diags', iter)\n",
    "gm_diags = mds.rdmds(path + 'Diags/GM_diags', iter)\n",
    "laydiag_wGM = mds.rdmds(path + 'Diags/layDiag', iter)\n",
    "stdiags_bylev_wGM,stdiags_2D_wGM,iters = diagnostics.readstats(path + 'dynStDiag.0000000000.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = '../run_noGM/'     # switch path to load output data from run without GM\n",
    "state_woGM = mds.rdmds(path + 'Diags/state', iter)\n",
    "diag2D_woGM = mds.rdmds(path + 'Diags/2D_diags', iter)\n",
    "laydiag_woGM = mds.rdmds(path + 'Diags/layDiag', iter)\n",
    "stdiags_bylev_woGM,stdiags_2D_woGM,iters = diagnostics.readstats(path + 'dynStDiag.0000000000.txt')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Statistical Diagnostics Time Series\n",
    "\n",
    "First we'll examine the time series of net surface heat flux into the ocean\n",
    "(here, diagnostic ``TRELAX``, from temperature relaxation),\n",
    "showing both w/GM and w/o GM\n",
    "runs on the same plot. As an equilibrium is approached, these series level off to a constant non-zero value because\n",
    "of net heat loss at the sponge layer northern boundary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28408 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#statistical diagnostics - TRELAX\n",
    "\n",
    "# note we have to do a conversion of iteration number into time units, \n",
    "# using timestep deltaT=1000 sec (and convert seconds -> years)\n",
    "\n",
    "deltaT = 1000\n",
    "plt.plot(np.array(iters['TRELAX'][:]) * deltaT/86400/360, stdiags_2D_wGM['TRELAX'][:,0],\n",
    "                  'b--', linewidth=4, label='Mean')\n",
    "plt.grid()\n",
    "plt.plot(np.array(iters['TRELAX'][:]) * deltaT/86400/360, stdiags_2D_wGM['TRELAX'][:,1], \n",
    "                  'm--', linewidth=4, label='Std Dev')\n",
    "plt.legend()\n",
    "plt.plot(np.array(iters['TRELAX'][:]) * deltaT/86400/360, stdiags_2D_woGM['TRELAX'][:,0],\n",
    "                  'b', linewidth=4)\n",
    "plt.plot(np.array(iters['TRELAX'][:]) * deltaT/86400/360, stdiags_2D_woGM['TRELAX'][:,1], \n",
    "                  'm', linewidth=4)\n",
    "plt.xlim(0,30)\n",
    "plt.xlabel('Time (yrs)')\n",
    "plt.ylabel('W/m$^2$')\n",
    "plt.title('Net Heat Flux into Ocean; dashed=w/GM', fontsize=18);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similary, we plot a time series of (global mean) ``THETA`` at the surface, mid-depth, and abyss,\n",
    "which appear to show similar equilibration at 30 years except for\n",
    "some continued warming at depth in the w/o GM run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 47764 characters]\n",
      "text/plain": [
       "<Figure size 576x504 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#statistical diagnostics - THETA\n",
    "\n",
    "# choose depths for time series plots of THETA \n",
    "kloc1 = 0    # MITgcm k=1 surface\n",
    "kloc2 = 16   # MITgcm k=17 270m\n",
    "kloc3 = 39   # MITgcm k=40 2577m depth\n",
    "\n",
    "fig, axarr = plt.subplots(3, 1, sharex=True, figsize=(8,7))\n",
    "\n",
    "axarr[0].plot(np.array(iters['THETA'][:]) * deltaT/86400/360, stdiags_bylev_wGM['THETA'][:,kloc1,0],\n",
    "                       'c--', linewidth=4)\n",
    "axarr[0].plot(np.array(iters['THETA'][:]) * deltaT/86400/360, stdiags_bylev_woGM['THETA'][:,kloc1,0],\n",
    "                       'c', linewidth=4)\n",
    "axarr[0].grid()\n",
    "axarr[0].set_ylabel('$^{\\circ}$C')\n",
    "axarr[0].set_title('Mean Temperature, %d m Depth; dashed=w/GM' %np.abs(RC[kloc1]), fontsize=16)\n",
    "axarr[1].plot(np.array(iters['THETA'][:]) * deltaT/86400/360, stdiags_bylev_wGM['THETA'][:,kloc2,0],\n",
    "                       'g--', linewidth=4)\n",
    "axarr[1].plot(np.array(iters['THETA'][:]) * deltaT/86400/360, stdiags_bylev_woGM['THETA'][:,kloc2,0],\n",
    "                       'g', linewidth=4);\n",
    "axarr[1].set_title('Mean Temperature, %d m Depth; dashed=w/GM' %np.abs(RC[kloc2]), fontsize=16)\n",
    "axarr[1].grid()\n",
    "axarr[1].set_ylabel('$^{\\circ}$C')\n",
    "axarr[2].plot(np.array(iters['THETA'][:]) * deltaT/86400/360, stdiags_bylev_wGM['THETA'][:,kloc3,0],\n",
    "                       'r--', linewidth=4)\n",
    "axarr[2].plot(np.array(iters['THETA'][:]) * deltaT/86400/360, stdiags_bylev_woGM['THETA'][:,kloc3,0],\n",
    "                       'r', linewidth=4)\n",
    "axarr[2].grid()\n",
    "axarr[2].set_ylabel('$^{\\circ}$C')\n",
    "axarr[2].set_title('Mean Temperature, %d m Depth; dashed=w/GM' %np.abs(RC[kloc3]), fontsize=16);\n",
    "axarr[2].set_xlim(0,30)\n",
    "axarr[2].set_xlabel('Time (yrs)');\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data.diagnostics output, year 30 annual mean\n",
    "\n",
    "To look at temperature structure in the yz plane,\n",
    "we compute and plot zonal mean temperature and zonal mean mixed layer depth. We use the grid variable\n",
    "``hFacC``, which contains the partial cell fraction, to compute a weighted mean.\n",
    "\n",
    "To plot, we use ``contourf`` to shade zonal mean temperature; unless we take extra steps,\n",
    "the plot will include a thin strip of white\n",
    "(or background color) space between land points and valid data at ``YC``, ``RC`` points. To avoid this,\n",
    "we create a \"padded\" copy (``zmTpad``) of the zonal mean temperature (``zmT``) data with an extra row and column,\n",
    "and do the following:\n",
    "\n",
    "  - duplicate the ``zmT`` data adjacent to the southern wall onto these \"land\" points\n",
    "  - copy the northernmost ``zmT`` data onto the extra column (at the north)\n",
    "  - copy the bottommost ``zmT`` data onto the extra row (at the bottom)\n",
    "  - we plot a gray rectangle over the southern wall points to properly show the land location\n",
    "\n",
    "Finally, when we call ``contourf``, we need to augment the both the ``YC`` and ``RC`` vectors\n",
    "with an additional gridpoint location.\n",
    "\n",
    "Note that we did NOT \"pad\" the top zmT row which is located at ``RC[0]``= -2.7 m\n",
    "(whereas our plot extends to the surface).\n",
    "Given that our vertical axis ranges over 2500 m depth, this 2.7 m worth of white space\n",
    "is not noticeable -- but it is obvious if one limits the plot\n",
    "to the top 100 m, for example, in which case repeating the padding procedure at the\n",
    "surface (i.e. augmenting a surface value\n",
    "point at 0 m depth) might be worthwhile.\n",
    "\n",
    "This may seem like a lot of extra effort when ``pcolormesh`` will avoid the white space problem, whether you pass\n",
    "``pcolormesh`` with ``Yp1``, ``RF`` sized dimensions or ``YC``, ``RC``\n",
    "dimensions with ``shading='auto'``. Indeed, it is. However\n",
    "``pcolormesh`` will generally give you a plot that looks pixelated whereas a ``contourf`` plot will render smooth\n",
    "gradations, depending on your choice of contour interval, often permitting better visualization of results.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 55948 characters]\n",
      "text/plain": [
       "<Figure size 1008x360 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# compute and compare zonal mean temperature and mixed layer depth\n",
    "\n",
    "# compute zonal mean T, weight by partial cell factor\n",
    "# using Cartesian grid (uniform horiz. spacing), all grid cells(k) have same volume\n",
    "# using a more complex grid, one might need to weight by cell area RAC\n",
    "# when computing area-weighted quantities\n",
    "\n",
    "zmT_woGM = np.ma.average(state_woGM[0,:,:,:], axis=2, weights=hFacC)\n",
    "zmT_wGM = np.ma.average(state_wGM[0,:,:,:], axis=2, weights=hFacC)\n",
    "\n",
    "# to pad zmT with temp values on land points, to avoid white space in contourf plot e.g.\n",
    "# between y=50km and first YC point with data, 75km\n",
    "\n",
    "zmTpad_woGM = np.zeros((RF.size, Ny+1))\n",
    "zmTpad_woGM[:-1,:-1] = zmT_woGM        # augment zmT shape (+1,+1)\n",
    "zmTpad_woGM[:-1,0] = zmT_woGM[:,1]     # copy zmT values onto southern land column\n",
    "zmTpad_woGM[:-1,-1] = zmT_woGM[:,-1]\n",
    "zmTpad_woGM[-1,:] = zmTpad_woGM[-2,:]  # copy onto added northern col. and bottom row\n",
    "zmTpad_wGM = np.zeros((RF.size, Ny+1))\n",
    "zmTpad_wGM[:-1,:-1] = zmT_wGM          # repeat padding procedure with GM run data\n",
    "zmTpad_wGM[:-1,0] = zmT_wGM[:,1]\n",
    "zmTpad_wGM[:-1,-1] = zmT_wGM[:,-1]\n",
    "zmTpad_wGM[-1,:] = zmTpad_wGM[-2,:]\n",
    "\n",
    "# make the plots\n",
    "# the plot in the users manual was done using MATLAB parula colormap\n",
    "# (which Mathworks arguably claim they have copyrighted, albeit posted here:\n",
    "# https://stackoverflow.com/questions/34859628/has-someone-made-the-parula-colormap-in-matplotlib )\n",
    "# instead we will use cmocean_v2 thermal, which out-of-the-box is a bit dark, so we lighten a bit.\n",
    "\n",
    "light_thermal = cmocean.tools.lighten(cmocean.cm.thermal, 0.90)\n",
    "fig, axarr = plt.subplots(1, 2, sharey=True, figsize=(14,5))\n",
    "\n",
    "im = axarr[0].contourf(np.append(YC[:,0], YC[-1,0] + DYC[-1,0])/1000,\n",
    "                       np.append(RC[:,0,0], RC[-1,0,0] + DRC[-1,0,0]),\n",
    "                       zmTpad_woGM, np.arange(-2,10,.2), cmap=light_thermal)\n",
    "plt.colorbar(im, ax =axarr[0])\n",
    "im = axarr[1].contourf(np.append(YC[:,0], YC[-1,0] + DYC[-1,0])/1000,\n",
    "                       np.append(RC[:,0,0], RC[-1,0,0] + DRC[-1,0,0]),\n",
    "                       zmTpad_wGM, np.arange(-2,10,.2), cmap=light_thermal)\n",
    "plt.colorbar(im, ax = axarr[1])\n",
    "axarr[0].set_xlim(0,2000)\n",
    "axarr[1].set_xlim(0,2000)\n",
    "axarr[0].set_ylim(-2500,0)\n",
    "axarr[0].add_patch(Rectangle((0, RF[-1,0,0]), 50, -RF[-1,0,0], facecolor='lightgray'))\n",
    "axarr[1].add_patch(Rectangle((0, RF[-1,0,0]), 50, -RF[-1,0,0], facecolor='lightgray'))\n",
    "axarr[0].set_xlabel('y-coordinate (km)')\n",
    "axarr[1].set_xlabel('y-coordinate (km)')\n",
    "axarr[0].set_ylabel('Depth (m)')\n",
    "axarr[1].set_ylabel('Depth (m)')\n",
    "\n",
    "# compute zonal mean ML depth and add to the plots\n",
    "zmML_woGM = np.ma.average(diag2D_woGM[1,:,:], axis=1, weights=hFacC[0,:,:])\n",
    "zmML_wGM = np.ma.average(diag2D_wGM[1,:,:], axis=1, weights=hFacC[0,:,:])\n",
    "axarr[0].plot(YC[1:,0]/1000, -zmML_woGM[1:], linewidth=4, color='aquamarine')\n",
    "axarr[1].plot(YC[1:,0]/1000, -zmML_wGM[1:], linewidth=4, color='aquamarine')\n",
    "\n",
    "axarr[0].set_title('w/o GM', fontsize=16)\n",
    "axarr[1].set_title('w/GM', fontsize=16)\n",
    "fig.suptitle('Zonal Mean Temp. ($^{\\circ}$C) and ML Depth',fontsize=18);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we plot the convective adjustment index (``CONVADJ``) fairly close to the surface\n",
    "(here, we choose 92 m depth).\n",
    "Typically, it is useful to know where\n",
    "convective mixing is occuring (and how often; 0=never convecting, 1= always convecting), and in this\n",
    "setup we see a big difference in mixed layer depths w/ and w/o GM. This is also readily\n",
    "apparent in the mixed layer depth plotted above.\n",
    "\n",
    "For this plot ``pcolormesh`` is perfectly acceptable. For clarity we again show the southern wall as a gray\n",
    "rectangle patch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 27732 characters]\n",
      "text/plain": [
       "<Figure size 1008x720 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# compare convective adjustment index\n",
    "\n",
    "#plot at MITgcm k=10 which is 92 m depth\n",
    "klev = 9\n",
    "\n",
    "# make plots; note here we use pcolormesh passed with size XC,YC\n",
    "# and let matplotlib figure out the XG,YG locations\n",
    "# uses cmocean_v2 colormap 'deep'\n",
    "\n",
    "fig, axarr = plt.subplots(1, 2, sharey=True, figsize=(14,10))\n",
    "\n",
    "im = axarr[0].pcolormesh(XC/1000, YC/1000, state_woGM[4,klev,:,:],\n",
    "                         cmap=cmocean.cm.deep, shading='auto')\n",
    "plt.colorbar(im, ax=axarr[0])\n",
    "axarr[0].add_patch(Rectangle((0, 0), 1000, 50, facecolor='lightgray'))\n",
    "axarr[0].set_xlabel('x-coordinate (km)')\n",
    "axarr[0].set_ylabel('y-coordinate (km)')\n",
    "axarr[0].set_aspect('equal')\n",
    "im = axarr[1].pcolormesh(XC/1000, YC/1000, state_wGM[4,klev,:,:],\n",
    "                         cmap=cmocean.cm.deep, shading='auto')\n",
    "plt.colorbar(im, ax =axarr[1])\n",
    "axarr[1].add_patch(Rectangle((0, 0), 1000, 50, facecolor='lightgray'));\n",
    "axarr[1].set_aspect('equal')\n",
    "axarr[1].set_xlabel('x-coordinate (km)')\n",
    "axarr[1].set_ylabel('y-coordinate (km)')\n",
    "\n",
    "axarr[0].set_title('w/o GM', fontsize=16)\n",
    "axarr[1].set_title('w/GM', fontsize=16)\n",
    "fig.suptitle('Convective Adj. Index (depth %d m)' %np.abs(RC[klev]), fontsize=18);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's examine the barotropic streamfunction to assess the strength of the zonal current. Similar to the\n",
    "computation in tutorial Baroclinic Ocean Gyre, first compute a depth-integrated\n",
    "zonal transport $U_{bt}$ (i.e., $u$ multiplied by the vertical cell depth ``DRF``).\n",
    "Note we use ``np.pad`` option ``'wrap'`` with specification ``[(0,0),(0,0),(0,1)]`` which will duplicate the upper end\n",
    "of the **third** axis (i.e., the x-axis) with values taken from the lower end. It is necessary to wrap the\n",
    "u-velocity entering the domain from the west to the domain eastern edge, in order to plot a contour of\n",
    "streamfunction over the full domain. We also weight the u-velocity by ``hFacW``,\n",
    "the partial cell factor at the western\n",
    "edge of each gridcell; this is the portion of the gridcell edge open to flow. The streamfunction $\\Psi_{bt}$\n",
    "is then formulated as follows:\n",
    "\n",
    "$$U_{bt} = - \\frac{\\partial \\Psi_{bt}}{\\partial y}, \\phantom{WW} V_{bt} = \\frac{\\partial \\Psi_{bt}}{\\partial x}$$\n",
    "\n",
    "which we compute by multiplying $U_{bt}$ by the length of the western gridcell edge (``DYG``) and doing\n",
    "a cumulative sum in the north-south direction. When $\\Psi_{bt}$ is plotted we add a row of zeros (i.e., no flow)\n",
    "at the southern boundary.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 136544 characters]\n",
      "text/plain": [
       "<Figure size 1008x720 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# barotropic streamfunction computation and comparison\n",
    "\n",
    "ubt_woGM = ((np.pad(state_woGM[2,:,:,:], [(0,0),(0,0),(0,1)], 'wrap') * \n",
    "             np.pad(hFacW, [(0,0),(0,0),(0,1)], 'wrap')) * DRF).sum(0)       # depth-integrated uvel\n",
    "psi_woGM = (np.pad(DYG,[(0,0),(0,1)], 'wrap') * (-ubt_woGM)).cumsum(0) / 1E6 # cumsum in N-S direction\n",
    "\n",
    "ubt_wGM = ((np.pad(state_wGM[2,:,:,:],[(0,0),(0,0),(0,1)], 'wrap')  * \n",
    "            np.pad(hFacW, [(0,0),(0,0),(0,1)], 'wrap')) * DRF).sum(0)        # repeat calculation w/GM data\n",
    "psi_wGM = (np.pad(DYG,[(0,0),(0,1)], 'wrap') * (-ubt_wGM)).cumsum(0) / 1E6\n",
    "\n",
    "# make plots\n",
    "# use cmocean_v2 colormap 'speed_r'\n",
    "\n",
    "fig, axarr = plt.subplots(1, 2, sharey=True, figsize=(14,10))\n",
    "\n",
    "im=axarr[0].contourf(Xp1/1000, Yp1/1000, np.concatenate((np.zeros((Nx+1,1)).T, psi_woGM)),\n",
    "                     levels=np.arange(-240,40,20), cmap=cmocean.cm.speed_r)\n",
    "plt.colorbar(im, ax=axarr[0])\n",
    "axarr[0].contour(Xp1/1000, Yp1/1000, np.concatenate((np.zeros((Nx+1,1)).T, psi_woGM)),\n",
    "                 levels=np.arange(-240,40,20), colors='k')\n",
    "axarr[0].add_patch(Rectangle((0, 0), 1000, 50, facecolor='lightgray'))\n",
    "axarr[0].set_xlabel('x-coordinate (km)')\n",
    "axarr[0].set_ylabel('y-coordinate (km)')\n",
    "axarr[0].set_aspect('equal')\n",
    "\n",
    "im = axarr[1].contourf(Xp1/1000, Yp1/1000, np.concatenate((np.zeros((Nx+1,1)).T, psi_wGM)),\n",
    "                       levels=np.arange(-240,40,20), cmap=cmocean.cm.speed_r)\n",
    "plt.colorbar(im, ax=axarr[1])\n",
    "axarr[1].contour(Xp1/1000, Yp1/1000, np.concatenate((np.zeros((Nx+1,1)).T, psi_wGM)),\n",
    "                 levels=np.arange(-240,40,20), colors='k')\n",
    "axarr[1].add_patch(Rectangle((0, 0), 1000, 50, facecolor='lightgray'))\n",
    "axarr[1].set_xlabel('x-coordinate (km)')\n",
    "axarr[1].set_ylabel('y-coordinate (km)')\n",
    "axarr[1].set_aspect('equal')\n",
    "\n",
    "axarr[0].set_title('w/o GM', fontsize=16)\n",
    "axarr[1].set_title('w/GM', fontsize=16)\n",
    "fig.suptitle('Barotropic Streamfunction', fontsize=18);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Shifting focus to the overturning circulation in the vertical, we compute and plot the \n",
    "meridional overturning circulation (MOC) streamfunction as:\n",
    "\n",
    "$$\\Psi_{moc} = \\int_{z}^0 \\int_\\mathrm{west}^\\mathrm{east} v \\,\\mathrm{d} x \\,\\mathrm{d} z $$\n",
    "\n",
    "First, we compute meridional transport: v-velocity multiplied\n",
    "by the southern edge gridcell length (``DXG``) and gridcell thickness (``DRF``),\n",
    "weighted by the partial cell factor on this southern edge (``hFacS``), and take its zonal sum (variable ``vzi``).\n",
    "A row of zeros is added at the bottom; we compute the streamfunction by doing a cumulative sum of ``vzi``\n",
    "in reverse k-index order. In general, integrating from the bottom upward is preferable, and in fact\n",
    "avoids issues if a non-linear free surface option is employed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 99744 characters]\n",
      "text/plain": [
       "<Figure size 1080x360 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# compute and compare Eulerian MOC\n",
    "\n",
    "# first, take zonal sum of v*dx*dz, adding a row of zeros at the bottom (RF spans 0:-3982 m dimension 50)\n",
    "vzi_woGM = np.concatenate((np.sum(state_woGM[1,:,:,:] * DXG * DRF * hFacS, 2), np.zeros((Ny,1)).T))\n",
    "moc_woGM = -np.cumsum(vzi_woGM[::-1,:],0)[::-1] / 1E6\n",
    "# repeat w/GM run\n",
    "vzi_wGM = np.concatenate((np.sum(state_wGM[1,:,:,:] * DXG * DRF * hFacS, 2), np.zeros((Ny,1)).T))\n",
    "moc_wGM = -np.cumsum(vzi_wGM[::-1,:],0)[::-1] / 1E6\n",
    "\n",
    "# make plots\n",
    "fig, axarr = plt.subplots(1, 2, sharey=True, figsize=(15,5))\n",
    "\n",
    "im = axarr[0].contourf(YG[:,0]/1000, RF[:,0,0], moc_woGM, np.arange(-3,3.1,.1), cmap='bwr')\n",
    "plt.colorbar(im, ax=axarr[0])\n",
    "axarr[0].contour(YG[:,0]/1000, RF[:,0,0], moc_woGM, np.arange(-3,3.5,.5), colors='k', linewidths=.5)\n",
    "im = axarr[1].contourf(YG[:,0]/1000, RF[:,0,0], moc_wGM, np.arange(-3,3.1,.1), cmap='bwr')\n",
    "plt.colorbar(im, ax=axarr[1])\n",
    "axarr[1].contour(YG[:,0]/1000, RF[:,0,0], moc_wGM, np.arange(-3,3.5,.5), colors='k', linewidths=.5)\n",
    "\n",
    "axarr[0].set_xlabel('y-coordinate (km)')\n",
    "axarr[1].set_xlabel('y-coordinate (km)')\n",
    "axarr[0].set_ylabel('Depth (m)')\n",
    "axarr[1].set_ylabel('Depth (m)')\n",
    "axarr[0].add_patch(Rectangle((0, RF[-1,0,0]), 50, -RF[-1,0,0], facecolor='lightgray'))\n",
    "axarr[1].add_patch(Rectangle((0, RF[-1,0,0]), 50, -RF[-1,0,0], facecolor='lightgray'))\n",
    "\n",
    "axarr[0].set_title('w/o GM', fontsize=16)\n",
    "axarr[1].set_title('w/GM', fontsize=16)\n",
    "fig.suptitle('Eulerian MOC (Sv)', fontsize=18);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above figures plot the Eulerian MOC, not the residual MOC, the latter being defined as the Eulerian component\n",
    "plus a contribution from the eddy induced circulation, here formulated as a bolus velocity.\n",
    "The output diagnostics from ``pkg/gmredi``  include bolus transport streamfunction (below, variable ``psi_bol``, i.e., diagnostic ``GM_PsiY``). Bolus v-velocity is computed as:\n",
    "\n",
    "$$ v_\\mathrm{bolus} = \\frac{- \\partial \\: \\mathrm{GM}_{\\Psi y}}{\\partial z} $$\n",
    "\n",
    "Conveniently, ``GM_PsiY`` is output at v-velocity points in the horizontal plane, but at the RF\n",
    "points in the vertical (i.e. vertical cell separation depths, not cell centers);\n",
    "see sections 2.11.4-2.11.5 in the MITgcm users guide. Thus, computing a\n",
    "derivative of ``GM_PsiY`` in the vertical (differencing and dividing by ``DRF``) locates\n",
    "the resulting bolus velocity at the vertical cell center (i.e. ``RC`` depths),\n",
    "coincident with the ``VVEL`` locations, so they can be added together without any spatial\n",
    "averaging. Note that ``GM_PsiY`` lacks a value at RF[-1] so we\n",
    "first augment variable ``psi_bol`` with a plane of zeros at the bottom.\n",
    "In the calculation of  ``vzi``, the zonally summed,  Eulerian+bolus meridional transport,\n",
    "it is not necessary to multiply ``bolV`` by ``hFacS`` because ``GM_PsiY`` is a transport diagnostic\n",
    "and has already been weighted by ``hFacS``.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 65360 characters]\n",
      "text/plain": [
       "<Figure size 468x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Compute and plot residual MOC (first computing bolus velocity from GM diagnostics)\n",
    "\n",
    "psi_bol = np.zeros((RF.size,Ny,Nx))\n",
    "psi_bol[:-1,:,:] = gm_diags[1,:,:,:] # pad bolus streamfn with zeros at bottom\n",
    "bolV = (psi_bol[1:,:,:] - psi_bol[:-1,:,:,]) / DRF\n",
    "# compute moc as before, except adding bolV to v-velocity\n",
    "vzi = np.concatenate((np.sum((state_wGM[1,:,:,:] * hFacS + bolV) * DXG * DRF,2), np.zeros((Ny,1)).T))\n",
    "moc_res = -np.cumsum(vzi[::-1,:],0)[::-1] / 1E6\n",
    "\n",
    "#make plot\n",
    "plt.figure(figsize=(6.5,5))\n",
    "plt.contourf(YG[:,0]/1000, RF[:,0,0], moc_res, np.arange(-3,3.2,.1), cmap='bwr')\n",
    "plt.colorbar()\n",
    "plt.colorbar(im, ax=axarr[0])\n",
    "plt.contour(YG[:,0]/1000, RF[:,0,0], moc_res,np.arange(-3,3.5,.5), colors='k', linewidths=.5)\n",
    "ax = plt.gca()\n",
    "ax.add_patch(Rectangle((0, RF[-1,0,0]), 50, -RF[-1,0,0], facecolor='lightgray'))\n",
    "\n",
    "plt.xlabel('y-coordinate (km)')\n",
    "plt.ylabel('Depth (m)')\n",
    "plt.title('Eulerian + Bolus MOC (Sv)', fontsize=18);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Package Layer Diagnostics\n",
    "\n",
    "### In temperature (density space):\n",
    "\n",
    "Computing and plotting the MOC in density space is straightforward; ``pkg/layers`` has\n",
    "already done the heavy lifting for you. While computing the MOC in z-coordinates is a fairly standard diagnostic\n",
    "for illustrating circulation, examining the overturning in density space is illuminating from a watermass\n",
    "transformation perspective.\n",
    "Given ``LaVH1TH``, the meridional mass transport in layers with units m$^2$s$^{-1}$\n",
    "(using a temperature vertical coordinate to define layers, see Fig. 4.12\n",
    "in the MITgcm users guide; also note this includes the bolus contribution), we can follow a procedure analogous\n",
    "to the above (depth-coordinate) MOC computations. Specifically, we take ``LaVH1TH`` (computed\n",
    "at v-velocity locations), do a cumlative sum in the vertical, multiply by \n",
    "the southern edge gridcell length (``DXG``) to get units of m$^3$s$^{-1}$, and then finally\n",
    "take the zonal sum (variable ``vti``).\n",
    "One important caveat: in depth space,\n",
    "k=[0:Nr] proceeds from ocean surface to bottom, whereas\n",
    "for layers diagnostics k=[0:Nr_layers] proceeds from the ocean bottom to the surface\n",
    "(this is because vector ``layer_bounds``, which is otherwise analogous to grid variable ``RF``\n",
    "for the vertical temperature coordinate,\n",
    "increases from coldest to warmest model temperature). Inserting a row of zeros at the top of ``vti``\n",
    "specifyies a zero starting value at the ocean bottom (as we did previously).\n",
    "\n",
    "Because ``pkg/layers`` is essentially doing calculations in a vertical coordinate that is different from\n",
    "the true model (grid) coordinate, some error is introduced in calculations and the value of streamfunction at\n",
    "the surface (from a bottom-up ``cumsum``) is not precisely zero. This residual could show up as \"noise\" contours in\n",
    "the upper left quadrant of our output plots. To eliminate this, we zero out small positive and \n",
    "negative values of ``vti``, here using a threshold of 0.005 Sv, which produces a clean contourf plot.\n",
    "\n",
    "In interpreting the output plot, the upper contour reflects the surface temperature (i.e., warmest at the north end, coldest at the southern wall). However, given that SST is not zonally uniform,\n",
    "we have added green dashed lines of maximum and minimum SST to delineate the surface layer.\n",
    "As such, contours between the SST max and min lines do not reflect overturning in the vertical,\n",
    "but are indicative of watermass transformation (i.e., surface heat exchange) in the surface layer.\n",
    "\n",
    "Finally, as noted above, the diagnostic ``LaVH1TH`` includes bolus transport from ``pkg/gmredi``\n",
    "(if used) so our plot below shows the residual MOC;\n",
    "in other words, the figure below w/GM is the density coordinate representation of the Eulerian + Bolus\n",
    "depth MOC shown above. In density space, any circulation that occurs strictly within a density \"class\", for example\n",
    "between ``layer_bounds[1]`` and ``layer_bounds[2]`` (temperature -1.75 to -1.5 $^{\\circ}$C)\n",
    "is not represented in the plot below; see additional discussion of these plots in section 4.3.5.1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 107496 characters]\n",
      "text/plain": [
       "<Figure size 1080x360 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Compute and plot residual MOC from package layer diagnostics, in density space\n",
    "\n",
    "vti_woGM = -np.concatenate((np.zeros((Ny,1)).T, (laydiag_woGM[0,:,:,:].cumsum(0)*DXG).sum(2))) / 1E6\n",
    "# due to small error in calculation, at the ocean top the value is not precisely zero\n",
    "# force small residual -> 0 to make a cleaner plot\n",
    "vti_woGM[abs(vti_woGM) < .005] = 0\n",
    "vti_wGM = -np.concatenate((np.zeros((Ny,1)).T, (laydiag_wGM[0,:,:,:].cumsum(0)*DXG).sum(2))) / 1E6\n",
    "vti_wGM[abs(vti_wGM) < .005] = 0\n",
    " \n",
    "# make plots\n",
    "fig, axarr = plt.subplots(1, 2, sharey=True, figsize=(15,5))\n",
    "\n",
    "im = axarr[0].contourf(YG[:,0]/1000, layer_bounds, vti_woGM, np.arange(-3,3,.01), cmap='bwr')\n",
    "plt.colorbar(im, ax=axarr[0])\n",
    "axarr[0].contour(YG[:,0]/1000, layer_bounds, vti_woGM, np.arange(-3,3,.5), linewidths=0.5, colors='k')\n",
    "# trick matplotlib to get it to draw top 0 contour\n",
    "axarr[0].contour(YG[:,0]/1000, layer_bounds, vti_woGM,[-0.00001], linewidths=0.5, colors='k') \n",
    "axarr[0].plot(YC[1:,0]/1000, state_woGM[0,0,1:,:].max(1), 'g--') # plot maximum SST(y) \n",
    "axarr[0].plot(YC[1:,0]/1000, state_woGM[0,0,1:,:].min(1), 'g--') # and minimum SST(y) to delineate the surface layer\n",
    "# in an eddy-resolving run, you need to dump data more frequently to obtain\n",
    "# reasonable estimates for SST max, min (e.g. daily, not monthly, dumps)\n",
    "\n",
    "im = axarr[1].contourf(YG[:,0]/1000, layer_bounds, vti_wGM, np.arange(-3,3,.01), cmap='bwr')\n",
    "plt.colorbar(im, ax=axarr[1])\n",
    "axarr[1].contour(YG[:,0]/1000, layer_bounds,vti_wGM, np.arange(-3,3,.5), linewidths=0.5, colors='k')\n",
    "axarr[1].contour(YG[:,0]/1000, layer_bounds, vti_wGM, [-0.00001], linewidths=0.5, colors='k') \n",
    "axarr[1].plot(YC[1:,0]/1000, state_wGM[0,0,1:,:].max(1), 'g--')\n",
    "axarr[1].plot(YC[1:,0]/1000, state_wGM[0,0,1:,:].min(1), 'g--')\n",
    "\n",
    "axarr[0].set_xlabel('y-coordinate (km)')\n",
    "axarr[0].set_ylabel('Temperature ($^{\\circ}$C)')\n",
    "axarr[1].set_xlabel('y-coordinate (km)')\n",
    "axarr[1].set_ylabel('Temperature ($^{\\circ}$C)')\n",
    "axarr[0].set_title('w/o GM', fontsize=16)\n",
    "axarr[1].set_title('w/GM', fontsize=16)\n",
    "fig.suptitle('Residual MOC (Sv)', fontsize=18);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Converted back to depth (z) space:\n",
    "\n",
    "While the merits of visualizing the overturning in density space is apparent, these plots can be\n",
    "a challenge to interpret in terms of the true spatial circulation. Moreover, it is often useful\n",
    "to compute and compare the depth-coordinate rendering of the above plots with our earlier\n",
    "z-coordinate MOC plots; the former will eliminate any adiabatic circulation as compared to the latter.\n",
    "\n",
    "For this calculation, we will first calculate the z-location of layer depths (``csum_th``)\n",
    "through a surface-to-bottom cumulative sum of diagnostic ``LaHs1TH`` (note, this is the layer thickness\n",
    "on the southern face of each cell, corresponding spatially with diagnostic ``LAVH1TH``).\n",
    "It is again necessary to first augment the result with a row\n",
    "of zeros, corresponding to the surface at z=0. Then we take a zonal mean. To calculate\n",
    "the layer transports (``csum_tr``), the computation is the same as above (using ``LAVH1TH``)\n",
    "except for reversing the\n",
    "k-index on the ``cumsum`` given for this calculation we proceed from the surface downward.\n",
    "\n",
    "To assess the extent to which interior circulation is adiabatic, also plotted are\n",
    "zonal mean temperature contours, using variable ``zmT`` computed earlier in this notebook. Not surprisingly,\n",
    "we see a big difference between the run without GM (which includes large parameterized horizontal mixing)\n",
    "versus the run with GM.\n",
    "\n",
    "Note that using a zonal mean to compute layer depths in the YZ plane may be problematic in a solution\n",
    "where layer depths vary significantly in a zonal band, e.g., in a setup with  more complex topography,\n",
    "if using a non-linear\n",
    "equation of state, and/or which includes marginal seas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 92184 characters]\n",
      "text/plain": [
       "<Figure size 1080x360 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#residual MOC from layers, converted back into depth coordinates\n",
    "\n",
    "# calculate layer depths through cumsum of layer thickness, surface to bottom\n",
    "# need to reverse k-index on cumsum given k[0] is coldest temperature (ocean bottom)\n",
    "csum_th_woGM = -np.concatenate((np.zeros((40,1)).T, laydiag_woGM[1,::-1,:,:].cumsum(0).mean(2)))\n",
    "# then cumsum the layer transports from ocean surface downward, multiply by DXG, and integrate zonally\n",
    "csum_tr_woGM = np.concatenate((np.zeros((40,1)).T, (laydiag_woGM[0,::-1,:,:].cumsum(0)*DXG).sum(2))) / 1E6\n",
    "# then this can be plotted specifying 2D arrays for x-axis (Y-coor) and y-axis (layer z-locations)\n",
    "\n",
    "# repeat calculation for run w/GM\n",
    "csum_th_wGM = -np.concatenate((np.zeros((40,1)).T, laydiag_wGM[1,::-1,:,:].cumsum(0).mean(2)))\n",
    "csum_tr_wGM = np.concatenate((np.zeros((40,1)).T, (laydiag_wGM[0,::-1,:,:].cumsum(0)*DXG).sum(2))) / 1E6\n",
    "\n",
    "Y = np.tile(YG[:,0]/1000, (38, 1))  # this is basically the Y output of a meshgrid(layer_bounds,YG[:,0])\n",
    "\n",
    "# make plots\n",
    "fig, axarr = plt.subplots(1, 2, sharey=True, figsize=(15,5))\n",
    "\n",
    "axarr[0].add_patch(Rectangle((0, -4000), 2000, 520.5, facecolor='lightgray'))\n",
    "im = axarr[0].contourf(Y, csum_th_woGM, csum_tr_woGM,np.arange(-3,3,.01), cmap='bwr', zorder=2)\n",
    "plt.colorbar(im, ax=axarr[0])\n",
    "# use zonal mean temperature variable zmT computed above\n",
    "axarr[0].contour(YC[:,0]/1000,RC[:,0,0], zmT_woGM, np.arange(-2,10,1), colors='k')\n",
    "axarr[0].add_patch(Rectangle((0, -4000), 50, 4000, facecolor='lightgray'))\n",
    "\n",
    "axarr[1].add_patch(Rectangle((0, -4000), 2000, 520.5, facecolor='lightgray'))\n",
    "im = axarr[1].contourf(Y, csum_th_wGM, csum_tr_wGM, np.arange(-3,3,.01), cmap='bwr', zorder=2)\n",
    "plt.colorbar(im, ax=axarr[1])\n",
    "axarr[1].contour(YC[:,0]/1000, RC[:,0,0], zmT_wGM, np.arange(-2,10,1), colors='k')\n",
    "axarr[1].add_patch(Rectangle((0, -4000), 50, 4000, facecolor='lightgray'))\n",
    "\n",
    "axarr[0].set_xlabel('y-coordinate (km)')\n",
    "axarr[0].set_ylabel('Depth (m)')\n",
    "axarr[1].set_xlabel('y-coordinate (km)')\n",
    "axarr[1].set_ylabel('Depth (m)')\n",
    "axarr[0].set_title('w/o GM', fontsize=16)\n",
    "axarr[1].set_title('w/GM', fontsize=16)\n",
    "fig.suptitle('Residual MOC converted to depth space (Sv)', fontsize=18);\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
