Convergence exploration.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convergence exploration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports\n",
    "\n",
    "We use *bokeh* for the convergence plot, *seaborn* for its colour palette, *pandas* to load data series from the `convergence.txt` files, as well as some more regular libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets\n",
    "from IPython.display import display\n",
    "\n",
    "from bokeh.io import output_notebook, push_notebook, show\n",
    "from bokeh.plotting import figure\n",
    "from bokeh.colors import RGB\n",
    "from bokeh.models import ColumnDataSource\n",
    "from bokeh.models.callbacks import CustomJS\n",
    "from bokeh.models.tools import HoverTool, TapTool\n",
    "\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os.path\n",
    "from glob import glob\n",
    "from netCDF4 import Dataset\n",
    "\n",
    "from pyGVtools import m6toolbox\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "output_notebook()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Experiment configuration\n",
    "\n",
    "We set up a path to the directory containing our data, and all the parameters swept across in our study. The experiments all exist in separate directories name `expt-cut<c>_min<m>_rest<r>`, where `<c>`, `<m>` and `<r>` are the parameter values. Being a bit more clever, we could do this automatically, but sometimes it's nice to be explicit about the data we want!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "basepath = os.path.expanduser('~/phd/data/adaptive-convergence/')\n",
    "paths = glob(os.path.join(basepath, 'expt-cut*'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = {\n",
    "    'cutoff': ['1e-3', '1e-2'],\n",
    "    'smooth': ['0', '0.01', '0.05', '0.1', '0.25', '0.5'],\n",
    "    'restoring': ['0', '1e2', '1e3', '1e4', '1e5', '1e6']\n",
    "}\n",
    "\n",
    "# we want to display more understandable names in the interactive widget\n",
    "# so this dictionary just translates those names to the short names\n",
    "# I used in naming the experiments\n",
    "param_trans = {\n",
    "    'cutoff': 'cut', 'smooth': 'min', 'restoring': 'rest'\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pre-load data\n",
    "\n",
    "To make plotting quicker, we pre-load all the data we'll need for section plots. The section slices were already prepared from full diagnostic output files with `ncks` as a post-processing script to save on storage space on my laptop. This way, we just have to read the full file in. For smaller files, it might be fine to load them on the fly, or cache the data only when it's been loaded for plotting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slices = {}\n",
    "for path in paths:\n",
    "    d = Dataset(os.path.join(path, 'prog-slice.nc'), 'r')\n",
    "    yq = d.variables['yq'][:]\n",
    "    e = d.variables['e'][0,:,:,0]\n",
    "    s = d.variables['adapt_coord_v'][0,:-1,:-1,0]\n",
    "    d.close()\n",
    "\n",
    "    # remove some erroneous missing values\n",
    "    s = np.maximum(s, 1e-30)\n",
    "    \n",
    "    slices[os.path.basename(path)] = (yq, e, s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameter filter\n",
    "\n",
    "For each varied parameter (`cut`, `min` and `rest`), we set up an ipywidget allowing us to select only the values we want to see in the convergence plot. The following routine performs this filtering, pushing the result back to the Bokeh plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initially, every parameter value is selected -- it's good to have this as\n",
    "# a baseline so the updates work\n",
    "selections = {}\n",
    "for param, values in params.items():\n",
    "    selections[param_trans[param]] = tuple(values)\n",
    "\n",
    "def update(change):\n",
    "    # the `change` structure contains details about the widget we interacted with\n",
    "    # here, we're interested in:\n",
    "    # - the name of the parameter we changed: 'cutoff', 'smooth' or 'restoring',\n",
    "    #   which we translate back to its \"raw\" parameter name\n",
    "    # - a tuple of selected values for this parameter\n",
    "    selections[param_trans[change['owner'].description]] = change['new']\n",
    "\n",
    "    # Bokeh produces each line as a separate \"renderer\", if it's not selected, we\n",
    "    # just make it invisible, otherwise make it visible if it is selected\n",
    "    for expt, renderer in lines.items():\n",
    "        # start off by assuming the line should be visible\n",
    "        enabled = True\n",
    "        \n",
    "        # iterate over the parameter name and its selected values\n",
    "        for param, selected in selections.items():\n",
    "            # is the current line within the selections for the current parameter?\n",
    "            for value in selected:\n",
    "                if param + value + '_' in expt:\n",
    "                    # yes; check the next parameter\n",
    "                    break\n",
    "            else:\n",
    "                # we hit this \"else\" block if we *don't* break out of the loop\n",
    "                # explicitly... that means the current line isn't selected for\n",
    "                # this parameter, so it should be invisible\n",
    "                enabled = False\n",
    "                break\n",
    "                \n",
    "        # finally, set the visibility\n",
    "        renderer.visible = enabled\n",
    "        \n",
    "    # this tells Bokeh to update the notebook with the changes we've made to\n",
    "    # renderer visibility\n",
    "    push_notebook()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section plot routine\n",
    "\n",
    "This uses the `m6toolbox` (just a package for plotting stuff on the arbitrary grids of MOM6) to plot a transect of the along-coordinate slope term. It's indended to be driven by an ipywidget so we can quickly select an experiment to plot.\n",
    "\n",
    "- To select the actual experiment to plot with its string name `expt`.\n",
    "- To change the colour scale (for example, to compare different sections while maintaining the same scale), use the `vmin` parameter.\n",
    "- To switch between a full-depth plot of the section, or just the upper 500m, use `full`.\n",
    "- If a section looks interesting, we can save it by providing a filename in the `save` parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_section(expt, vmin=-12, full=False, save=None):\n",
    "    plt.figure(figsize=(8,10))\n",
    "    \n",
    "    data = slices[expt]\n",
    "    norm = mpl.colors.LogNorm(vmin=10**vmin, vmax=1.0)\n",
    "    x,z,q = m6toolbox.section2quadmesh(*data)\n",
    "    \n",
    "    plt.pcolormesh(x, z, q, norm=norm)\n",
    "    plt.plot(x, z.T, color='white', linewidth=0.6)\n",
    "    if not full:\n",
    "        plt.xlim(-60, -40)\n",
    "        plt.ylim(-500, 0)\n",
    "    plt.xlabel('latitude')\n",
    "    plt.ylabel('height (m)')\n",
    "    plt.title(expt)\n",
    "    plt.colorbar()\n",
    "    if save is not None:\n",
    "        plt.savefig(save)\n",
    "    else:\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convergence plots\n",
    "\n",
    "Here we plot the whole suite of convergence tests, but allow them to be filtered by the various parameters. We can switch between the L1 and L2 norms of our convergence measure, or take a look at the total amount of limiting that occurred during the convergence test."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide_input": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# create a new Bokeh figure with given axis labels\n",
    "p = figure(x_axis_label='iteration', y_axis_label='normalised global error')\n",
    "# create a unique colour for each experiment\n",
    "palette = sns.hls_palette(len(paths))\n",
    "# here's the dictionary that saves the line renderers for each experiment\n",
    "lines = {}\n",
    "\n",
    "# read and plot a line for each experiment (combine with its colour)\n",
    "for path, col in zip(paths, palette):\n",
    "    # we append '_' to the experiment name so we can distinguish between e.g.\n",
    "    # min0 and min0.5 in the selection test -- we would see whether 'min0_' or\n",
    "    # 'min0.5_' appear in the experiment name\n",
    "    expt = os.path.basename(path) + '_'\n",
    "    # use pandas to read the whitespace-delimited text file of convergence metrics\n",
    "    # into a data table\n",
    "    data = pd.read_table(os.path.join(path, 'convergence.txt'),\n",
    "                         header=None, delim_whitespace=True, names=('l1', 'l2', 'li'))\n",
    "    \n",
    "    # this is a little bit of Bokeh magic: we load a data table as the data source for the line\n",
    "    # where we gave the three columns names 'l1', 'l2' and 'li'\n",
    "    # to choose which column is plotted, we only need to tell Bokeh the name of the column,\n",
    "    # rather than providing the full data!\n",
    "    l = p.line(x='index', y='l2', source=data, line_width=2, \n",
    "               color=RGB(*[int(x*255) for x in col]))\n",
    "    # save the line renderer in our dictionary\n",
    "    lines[expt] = l\n",
    "    \n",
    "    # add a hover tool that'll show us the experiment name, iteration and value\n",
    "    p.add_tools(HoverTool(toggleable=False, renderers=[l], tooltips=[\n",
    "        ('Expt', expt[:-1]),\n",
    "        ('Iteration', '$x{int}'),\n",
    "        ('Error', '$y')\n",
    "    ]))\n",
    "    \n",
    "# hide legend and render it under the hovertools so we can inspect the entire line\n",
    "# if we just hid it, it would still technically exist and prevent us from hovering over\n",
    "# parts of the plot it would have covered (which could be a lot because we might be\n",
    "# displaying many lines!)\n",
    "p.legend.visible = False\n",
    "p.legend.level = 'underlay'\n",
    "\n",
    "# actually show the figure in the notebook\n",
    "# by setting `notebook_handle=True`, we're allowing the use of `push_notebook()`\n",
    "# to update the figure\n",
    "show(p, notebook_handle=True)\n",
    "\n",
    "# first, count the maximum number of values any parameter takes\n",
    "# this is so we can set the height of the selection widgets\n",
    "rows = 0\n",
    "for v in params.values():\n",
    "    rows = max(rows, len(v))\n",
    "    \n",
    "# now, for each parameter, we create a `SelectMultiple` widget to\n",
    "# allow selecting different values\n",
    "widgets = []\n",
    "for param, values in params.items():\n",
    "    widget = ipywidgets.SelectMultiple(options=values, value=values,\n",
    "                                       rows=rows, description=param, disabled=False)\n",
    "    # when the selection is changed, call our `update()` function\n",
    "    widget.observe(update, names='value')\n",
    "    # save the widget so we can actually display it with our plot\n",
    "    widgets.append(widget)\n",
    "\n",
    "# this function just changes which column of the data source is plotted\n",
    "# we use 'l' + the second character of the option name ('1', '2' or 'i')\n",
    "# to choose the column -- these are the column names that were set up initially\n",
    "def set_norm(change):\n",
    "    for line in lines.values():\n",
    "        # choose series\n",
    "        l.glyph.update(y='l' + change['new'][1])\n",
    "        # force refresh by making the line invisible, then visible again\n",
    "        l.visible = False\n",
    "        l.visible = True\n",
    "        \n",
    "    # force recalculation of the y range of the plot -- I'm not sure this\n",
    "    # is strictly necessary\n",
    "    p.y_range = None\n",
    "\n",
    "    push_notebook()\n",
    "    \n",
    "# define the buttons for toggling between convergence metrics (defaulting to the\n",
    "# L2 norm), and call the `set_norm()` function when the value changes\n",
    "norm_buttons = ipywidgets.ToggleButtons(options=['L1 norm', 'L2 norm', 'Limiting'],\n",
    "                                        value='L2 norm',\n",
    "                                       description='Convergence norm')\n",
    "norm_buttons.observe(set_norm, names='value')\n",
    "\n",
    "# finally, display the buttons and the selection widgets so we can actually\n",
    "# interact with them\n",
    "display(norm_buttons)\n",
    "ipywidgets.HBox(widgets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section plots\n",
    "\n",
    "Here we can choose to plot a section of the spatial map of the convergence metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide_input": false
   },
   "outputs": [],
   "source": [
    "w = ipywidgets.Select(options=[os.path.basename(path) for path in paths])\n",
    "ipywidgets.interactive(lambda expt, vmin, full=False: plot_section(expt, vmin, full, save=None), expt=w, vmin=(-12,-2))"
   ]
  }
 ],
 "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}