{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Produces an interactive heatmap with x and y as lattitude and longitude, the colour represents the value of the selected tracer variable, and two sliders control the depth and timestep.\n",
    "\n",
    "Modified version of this script:\n",
    "http://stackoverflow.com/questions/6697259/interactive-matplotlib-plot-with-two-sliders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from pylab import *\n",
    "from matplotlib.widgets import Slider, Button, RadioButtons\n",
    "\n",
    "import netCDF4 as nc\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "grid_t = nc.Dataset('/data/jpetrie/MEOPAR/NEMO-3.6-code/NEMOGCM/CONFIG/mygyre/EXP00/GYRE_5d_00010101_00011230_grid_T.nc')\n",
    "#/data/jpetrie/MEOPAR/SalishSea/results/0bc30428-1611-11e6-9e0c-0025909a8461/SS5x5_1h_20041019_20041023_ptrc_T.nc')\n",
    "#/data/eolson/MEOPAR/SS36runs/runNoBioPISCES/SOG2D_1h_20041019_20041019_grid_T.nc')\n",
    "\n",
    "\n",
    "lons = grid_t.variables['nav_lon']\n",
    "lats = grid_t.variables['nav_lat']\n",
    "keys = grid_t.variables.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Looks for all variables with a 4 dimensional input. \n",
    "# Should probably check dimensions match as well, but doesn't currently\n",
    "plottable_vars = list() \n",
    "for key in keys:\n",
    "    if np.size(grid_t.variables[key].shape) == 4:\n",
    "        plottable_vars.append(key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Global variable used to share information between the update functions about what tracer is currently being plotted\n",
    "current_var = plottable_vars[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Function used to update the heatmap when depth/time sliders are moved\n",
    "def update(val):\n",
    "    time_index = stime.val\n",
    "    depth_index = sdepth.val\n",
    "    # set_array() requires a 1D array. ravel() is used to get around this\n",
    "    # The [:-1,:-1] is to counteract the weird indexing that pcolormesh does\n",
    "    #   See: http://stackoverflow.com/questions/29009743/using-set-array-with-pyplot-pcolormesh-ruins-figure\n",
    "    mesh.set_array(grid_t.variables[current_var][round(time_index),round(depth_index),:,:][:-1,:-1].ravel())\n",
    "    draw()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def updateVar(variable_name):\n",
    "    global current_var\n",
    "    time_index = stime.val\n",
    "    depth_index = sdepth.val\n",
    "    # set_array() requires a 1D array. ravel() is used to get around this\n",
    "    # The [:-1,:-1] is to counteract the weird indexing that pcolormesh does\n",
    "    #   See: http://stackoverflow.com/questions/29009743/using-set-array-with-pyplot-pcolormesh-ruins-figure\n",
    "    mesh.set_array(grid_t.variables[variable_name][round(time_index),round(depth_index),:,:][:-1,:-1].ravel())\n",
    "    current_var = variable_name\n",
    "    draw()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7fb31ad633c8>"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "subplots_adjust(left=0.25, bottom=0.25)\n",
    "\n",
    "mesh = ax.pcolormesh(lats, lons,grid_t.variables[plottable_vars[0]][0,0,:,:])\n",
    "fig.colorbar(mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Dimensions for each of the sliders\n",
    "axfreq = axes([0.25, 0.1,0.65, 0.03])\n",
    "axamp  = axes([0.25, 0.15, 0.65, 0.03])\n",
    "vax = axes([0.025, 0.5, 0.15, 0.15])\n",
    "\n",
    "\n",
    "stime = Slider(axfreq, 'Time', 1, grid_t.dimensions['time_counter'].size -1)\n",
    "sdepth = Slider(axamp, 'Depth', 1,grid_t.dimensions['deptht'].size - 1)\n",
    "radio = RadioButtons(vax, plottable_vars, active=0) # creates a list of checkboxes with the different tracer options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stime.on_changed(update)\n",
    "sdepth.on_changed(update)\n",
    "radio.on_clicked(updateVar)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.5.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}