{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "This notebook sets up and runs a set of benchmarks to compare\n",
    "different numerical discretizations of the SWEs\n",
    "\n",
    "Copyright (C) 2016  SINTEF ICT\n",
    "\n",
    "This program is free software: you can redistribute it and/or modify\n",
    "it under the terms of the GNU General Public License as published by\n",
    "the Free Software Foundation, either version 3 of the License, or\n",
    "(at your option) any later version.\n",
    "\n",
    "This program is distributed in the hope that it will be useful,\n",
    "but WITHOUT ANY WARRANTY; without even the implied warranty of\n",
    "MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the\n",
    "GNU General Public License for more details.\n",
    "\n",
    "You should have received a copy of the GNU General Public License\n",
    "along with this program.  If not, see <http://www.gnu.org/licenses/>.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Lets have matplotlib \"inline\"\n",
    "%matplotlib inline\n",
    "\n",
    "#Python 2.7 compatibility\n",
    "from __future__ import print_function\n",
    "\n",
    "#Import packages we need\n",
    "import numpy as np\n",
    "from matplotlib import animation, rc\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib import gridspec\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "#import mpld3\n",
    "\n",
    "import subprocess\n",
    "import os\n",
    "import gc\n",
    "import datetime\n",
    "\n",
    "import pycuda.driver as cuda\n",
    "\n",
    "try:\n",
    "    from StringIO import StringIO\n",
    "except ImportError:\n",
    "    from io import StringIO\n",
    "\n",
    "#Finally, import our simulator\n",
    "from SWESimulators import Common, LxF, FORCE, HLL, HLL2, KP07, KP07_dimsplit, WAF, PlotHelper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No existing CUDA context\n",
      "CUDA version (9, 1, 0)\n",
      "Driver version 9010\n",
      "Using GeForce 840M\n",
      " => compute capability: (5, 0)\n",
      " => memory: 2048.0 MB\n",
      "Created context <527715955808>\n"
     ]
    }
   ],
   "source": [
    "import atexit\n",
    "def exitfunc():\n",
    "    #Clean up old context if any:\n",
    "    if 'cuda_context' in globals():\n",
    "        print(\"Deleting existing context\")\n",
    "        global cuda_context\n",
    "        del cuda_context\n",
    "        gc.collect()\n",
    "    else:\n",
    "        print(\"No existing CUDA context\")\n",
    "        gc.collect()\n",
    "        \n",
    "atexit.register(exitfunc)\n",
    "exitfunc()\n",
    "    \n",
    "cuda_context = Common.CudaContext(verbose=True, blocking=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Misc plotting setup\n",
    "def setBwStyles(ax):\n",
    "    from cycler import cycler\n",
    "\n",
    "    ax.set_prop_cycle( cycler('marker', ['.', 'x', 4, '+', '*', '1', 5]) +\n",
    "                       cycler('linestyle', ['-.', '--', ':', '-.', '--', ':', '-.']) +\n",
    "                       #cycler('markersize', [5, 5, 5, 5, 5, 5]) +\n",
    "                       cycler('color', ['k', 'k', 'k', 'k', 'k', 'k', 'k']) )    \n",
    "\n",
    "#Set large figure sizes\n",
    "plt.rcParams['figure.figsize'] = [12, 8]\n",
    "plt.rcParams['figure.dpi'] = 100\n",
    "plt.rcParams['animation.html'] = 'html5'\n",
    "#plt.rcParams['legend.markerscale'] = 1.0\n",
    "#plt.rcParams['lines.markersize'] = 6\n",
    "plt.rcParams['lines.markeredgewidth'] = 1.5\n",
    "#plt.rcParams['savefig.dpi'] = 400"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_figure(fig, stem):\n",
    "    fig_filename = \"convergence_smooth1d_\" + stem + \".pdf\"\n",
    "    \n",
    "    image_text = fig_filename + \"\\nCreated \" + str(datetime.datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\"))\n",
    "    \n",
    "    #ax = fig.gca()\n",
    "    #pos1 = ax.get_position() # get the original position \n",
    "    #pos2 = [pos1.x0 + 0.3, pos1.y0 + 0.3,  pos1.width / 2.0, pos1.height / 2.0] \n",
    "    #ax.set_position(pos2) # set a new position\n",
    "    \n",
    "    #fig.gca().annotate(image_text, xy=(1.1, 0), xycoords='figure fraction', fontsize=10,\n",
    "    #                   horizontalalignment='right', verticalalignment='bottom')\n",
    "    t = plt.figtext(-0.1, .9, image_text, rotation='vertical')\n",
    "    \n",
    "    legend = fig.gca().legend_\n",
    "    if (legend != None):\n",
    "        fig.savefig(fig_filename, bbox_extra_artists=(legend,t, ), bbox_inches='tight')\n",
    "    else:\n",
    "        fig.savefig(fig_filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 11300 characters]\n",
      "text/plain": [
       "<Figure size 1200x800 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 61812 characters]\n",
      "text/plain": [
       "<Figure size 1200x800 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def downsample(highres_solution, factor):\n",
    "    if (len(highres_solution.shape) == 1):\n",
    "        highres_solution = highres_solution.reshape((1, highres_solution.size))\n",
    "\n",
    "    assert(highres_solution.shape[1] % factor == 0)\n",
    "    ny = highres_solution.shape[0] \n",
    "    nx = highres_solution.shape[1] / factor\n",
    "\n",
    "    if (factor == 1):\n",
    "        return highres_solution\n",
    "    else:\n",
    "        \n",
    "        return highres_solution.reshape([int(ny), int(1), int(nx), int(factor)]).mean(3).mean(1)\n",
    "    \n",
    "\n",
    "# Test 2D downsample\n",
    "a = np.arange(60).reshape((6, 10))\n",
    "b = downsample(a, 2)\n",
    "\n",
    "plt.figure()\n",
    "plt.subplot(1,2,1)\n",
    "plt.imshow(a, extent=[0, 1, 0, 1])\n",
    "\n",
    "plt.subplot(1,2,2)\n",
    "plt.imshow(b, extent=[0, 1, 0, 1])\n",
    "plt.show()\n",
    "\n",
    "# Test 1D downsample\n",
    "a = np.sin(np.linspace(0, 2*np.pi, 40))\n",
    "b = downsample(a, 4)\n",
    "\n",
    "plt.figure()\n",
    "plt.subplot(1,2,1)\n",
    "plt.plot(np.linspace(0, 1, a.size), a.flatten(), 'x-')\n",
    "\n",
    "plt.subplot(1,2,2)\n",
    "plt.plot(np.linspace(0, 1, b.size), b.flatten(), 'x-')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def wall_boundary_conditions(data, num_ghost_cells):\n",
    "    data[0:num_ghost_cells,:] = data[2*num_ghost_cells-1:num_ghost_cells-1:-1,:]\n",
    "    data[-num_ghost_cells:,:] = data[-num_ghost_cells-1:-2*num_ghost_cells-1:-1,:]\n",
    "    data[:,0:num_ghost_cells] = data[:,2*num_ghost_cells-1:num_ghost_cells-1:-1]\n",
    "    data[:,-num_ghost_cells:] = data[:,-num_ghost_cells-1:-2*num_ghost_cells-1:-1]\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 25296 characters]\n",
      "text/plain": [
       "<Figure size 1200x800 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.0, 50.0, 0.27710612593755635]\n"
     ]
    }
   ],
   "source": [
    "def gen_test_data(nx, ref_nx, g, num_ghost_cells):\n",
    "    assert(num_ghost_cells >= 1)\n",
    "    ref_ny = num_ghost_cells\n",
    "    ny = ref_ny\n",
    "    \n",
    "    width = 100.0\n",
    "    height = 100.0\n",
    "    ref_dx = width / float(ref_nx)\n",
    "    ref_dy = height / float(ref_ny)\n",
    "\n",
    "    x_center = ref_dx*ref_nx/2.0\n",
    "    y_center = ref_dy*ref_ny/2.0\n",
    "\n",
    "    #Create a gaussian \"dam break\" that will not form shocks\n",
    "    size = width / 5.0\n",
    "    dt = 10**10\n",
    "    \n",
    "    h  = np.zeros((ny+2*num_ghost_cells, nx+2*num_ghost_cells), dtype=np.float32); \n",
    "    hu = np.zeros((ny+2*num_ghost_cells, nx+2*num_ghost_cells), dtype=np.float32);\n",
    "    hv = np.zeros((ny+2*num_ghost_cells, nx+2*num_ghost_cells), dtype=np.float32);\n",
    "    \n",
    "    x = ref_dx*(np.arange(0, ref_nx, dtype=np.float32)+0.5) - x_center\n",
    "    y = ref_dy*(np.arange(0, ref_ny, dtype=np.float32)+0.5) - y_center\n",
    "    xv, yv = np.meshgrid(x, y, sparse=False, indexing='xy')\n",
    "    r = np.sqrt(xv**2)\n",
    "    xv = None\n",
    "    yv = None\n",
    "    gc.collect()\n",
    "    \n",
    "    #Make sure that the wanted resolution is suitable\n",
    "    assert(ref_nx >= nx)\n",
    "    assert(ref_ny >= ny)\n",
    "    \n",
    "    #Generate highres then downsample\n",
    "    #h_highres = 0.5 + 0.1*np.exp(-(xv**2/size + yv**2/size))\n",
    "    h_highres = 0.5 + 0.1*0.5*(1.0 + np.cos(np.pi*r/size)) * (r < size)\n",
    "    h[num_ghost_cells:-num_ghost_cells, num_ghost_cells:-num_ghost_cells] = downsample(h_highres, ref_nx/nx)\n",
    "    h_highres = None\n",
    "    gc.collect()\n",
    "    \n",
    "    #hu_highres = 0.1*np.exp(-(xv**2/size + yv**2/size))\n",
    "    hu_highres = 0.1*0.5*(1.0 + np.cos(np.pi*r/size)) * (r < size)\n",
    "    hu[num_ghost_cells:-num_ghost_cells, num_ghost_cells:-num_ghost_cells] = downsample(hu_highres, ref_nx/nx)\n",
    "    hu_highres = None\n",
    "    gc.collect()\n",
    "    \n",
    "    #Fix boundary conditions\n",
    "    h = wall_boundary_conditions(h, num_ghost_cells);\n",
    "    hu = wall_boundary_conditions(hu, num_ghost_cells);\n",
    "    \n",
    "    scale = 0.7\n",
    "    max_h_estimate = 0.6\n",
    "    max_u_estimate = 0.1\n",
    "    dx = width/nx\n",
    "    dy = height/ny\n",
    "    dt = scale * min(dx, dy) / (max_u_estimate + np.sqrt(g*max_h_estimate))\n",
    "    \n",
    "    return h, hu, hv, dx, dy, dt\n",
    "\n",
    "h, hu, hv, dx, dy, dt = gen_test_data(nx=100, ref_nx=1000, g=9.81, num_ghost_cells=2)\n",
    "\n",
    "plt.figure()\n",
    "ax=plt.subplot(1,3,1)\n",
    "im=plt.imshow(h, extent=[0, 100, 0, 100])\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "plt.colorbar(im, cax=cax)\n",
    "\n",
    "ax=plt.subplot(1,3,2)\n",
    "im=plt.imshow(hu, extent=[0, 100, 0, 100])\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "plt.colorbar(im, cax=cax)\n",
    "\n",
    "ax=plt.subplot(1,3,3)\n",
    "im=plt.imshow(hv, extent=[0, 100, 0, 100])\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "plt.colorbar(im, cax=cax)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "print([dx, dy, dt])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 58028 characters]\n",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#Plot initial conditions            \n",
    "fig = plt.figure()\n",
    "setBwStyles(fig.gca())\n",
    "\n",
    "for nx in [128, 1024, 16384]:\n",
    "    h0, hu0, hv0, dx, dy, dt = gen_test_data(nx, nx, 9.81, 1)\n",
    "    x = np.linspace(0.0, nx*dx, nx)\n",
    "    \n",
    "    #plt.subplot(211)\n",
    "    plt.plot(x, h0[1:-1, 1:-1].flatten(), marker=' ', label='h_' + str(nx))\n",
    "    plt.plot(x, hu0[1:-1, 1:-1].flatten(), marker=' ', label='hu_' + str(nx))\n",
    "    #plt.subplot(212)\n",
    "    #setBwStyles(fig.gca())\n",
    "    #plt.plot(x, hu0.T, marker=\" \", label='hu_' + str(nx))\n",
    "    #plt.legend(loc='upper left')\n",
    "    \n",
    "#plt.xlim([47, 53])\n",
    "#plt.ylim([0.58, 0.61])\n",
    "plt.legend(loc='upper left')\n",
    "save_figure(fig, 'initial_conditions')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_filename(simulator, nx):\n",
    "    return \"smooth1d_\" + str(simulator.__name__) + \"_\" + str(nx) + \".npz\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sanity_check(variable, variable_name):\n",
    "    maxval = np.amax(variable)\n",
    "    minval = np.amin(variable)\n",
    "    if (np.isnan(maxval) or np.isnan(minval)):\n",
    "        print(\"=== WARNING ===\")\n",
    "        print(variable_name + \" contains NaN values\")\n",
    "        print(\"=== WARNING ===\")\n",
    "\n",
    "def run_benchmark(datafilename, simulator, courant_number, nx, reference_nx, \n",
    "                  ghost_cells, g=9.81, force_rerun=False, transpose=False):\n",
    "\n",
    "    if (os.path.isfile(datafilename) and force_rerun == False):\n",
    "        print(\"=== WARNING ===\")\n",
    "        print(\"Existing simulation exists, skipping simulation\")\n",
    "        print(\"=== WARNING ===\")\n",
    "    else:\n",
    "        h0, hu0, hv0, dx, dy, dt = gen_test_data(nx, reference_nx, g, ghost_cells)\n",
    "        ny = ghost_cells\n",
    "        dt = dt * courant_number\n",
    "\n",
    "        #Initialize simulator\n",
    "        with Common.Timer(simulator.__name__ + \"_\" + str(nx)) as timer:\n",
    "            if (transpose):\n",
    "                h0 = np.ascontiguousarray(np.transpose(h0))\n",
    "                hu0, hv0 = np.ascontiguousarray(np.transpose(hv0)), np.ascontiguousarray(np.transpose(hu0))\n",
    "                dx, dy = dy, dx\n",
    "                nx, ny = ny, nx\n",
    "\n",
    "            sim = simulator(cuda_context, \n",
    "                            h0, hu0, hv0, \n",
    "                            nx, ny, \n",
    "                            dx, dy, dt, \n",
    "                            g)\n",
    "\n",
    "            t, nt = sim.step(1.0)\n",
    "            \n",
    "            h, hu, hv = sim.download()\n",
    "            \n",
    "            if (transpose):\n",
    "                h = np.transpose(h)\n",
    "                hu, hv = np.ascontiguousarray(np.transpose(hv)), np.ascontiguousarray(np.transpose(hu))\n",
    "\n",
    "            h  =  h[ghost_cells, ghost_cells:-ghost_cells]\n",
    "            hu = hu[ghost_cells, ghost_cells:-ghost_cells]\n",
    "            hv = hv[ghost_cells, ghost_cells:-ghost_cells]\n",
    "            \n",
    "            sanity_check(h, \"h\")\n",
    "            sanity_check(hu, \"hu\")\n",
    "            sanity_check(hv, \"hv\")\n",
    "\n",
    "            np.savez_compressed(datafilename, h=h, hu=hu, hv=hv)\n",
    "        \n",
    "    gc.collect() # Force run garbage collection to free up memory\n",
    "    \n",
    "    return [timer.secs, nt, dt]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=> LxF_16 7584.396839 ms\n",
      "=> FORCE_16 5857.448339 ms\n",
      "=> HLL_16 5091.929436 ms\n",
      "=> HLL2_16 5243.288517 ms\n",
      "=> KP07_16 5151.494980 ms\n",
      "=> KP07_dimsplit_16 5320.673227 ms\n",
      "=> WAF_16 4885.027409 ms\n",
      "Completed warmup in 41.114710330963135 s\n",
      "=> LxF_8 20.020962 ms\n",
      "=> LxF_16 21.036625 ms\n",
      "=> LxF_32 16.012192 ms\n",
      "=> LxF_64 20.009041 ms\n",
      "=> LxF_128 18.985271 ms\n",
      "=> LxF_256 28.013706 ms\n",
      "=> LxF_512 33.023596 ms\n",
      "=> LxF_1024 31.055212 ms\n",
      "=> LxF_2048 42.028427 ms\n",
      "=> LxF_4096 81.029415 ms\n",
      "=> LxF_16384 1678.168297 ms\n",
      "=> FORCE_8 27.021408 ms\n",
      "=> FORCE_16 27.018785 ms\n",
      "=> FORCE_32 44.031382 ms\n",
      "=> FORCE_64 48.034430 ms\n",
      "=> FORCE_128 32.021999 ms\n",
      "=> FORCE_256 35.024166 ms\n",
      "=> FORCE_512 33.024073 ms\n",
      "=> FORCE_1024 34.022808 ms\n",
      "=> FORCE_2048 55.034399 ms\n",
      "=> FORCE_4096 106.080294 ms\n",
      "=> FORCE_16384 2327.248812 ms\n",
      "=> HLL_8 20.013809 ms\n",
      "=> HLL_16 31.023026 ms\n",
      "=> HLL_32 32.022238 ms\n",
      "=> HLL_64 31.021357 ms\n",
      "=> HLL_128 29.021263 ms\n",
      "=> HLL_256 35.024643 ms\n",
      "=> HLL_512 31.026840 ms\n",
      "=> HLL_1024 33.034325 ms\n",
      "=> HLL_2048 52.037001 ms\n",
      "=> HLL_4096 142.127037 ms\n",
      "=> HLL_16384 3066.349030 ms\n",
      "=> HLL2_8 22.014856 ms\n",
      "=> HLL2_16 24.025917 ms\n",
      "=> HLL2_32 37.019730 ms\n",
      "=> HLL2_64 24.045944 ms\n",
      "=> HLL2_128 25.022507 ms\n",
      "=> HLL2_256 26.018620 ms\n",
      "=> HLL2_512 32.016039 ms\n",
      "=> HLL2_1024 39.027214 ms\n",
      "=> HLL2_2048 74.053049 ms\n",
      "=> HLL2_4096 201.141596 ms\n",
      "=> HLL2_16384 4902.546167 ms\n",
      "=> KP07_8 20.987272 ms\n",
      "=> KP07_16 33.022165 ms\n",
      "=> KP07_32 25.025129 ms\n",
      "=> KP07_64 24.017096 ms\n",
      "=> KP07_128 26.021004 ms\n",
      "=> KP07_256 28.018713 ms\n",
      "=> KP07_512 33.029556 ms\n",
      "=> KP07_1024 44.034958 ms\n",
      "=> KP07_2048 89.062214 ms\n",
      "=> KP07_4096 251.171112 ms\n",
      "=> KP07_16384 6784.766197 ms\n",
      "=> KP07_dimsplit_8 23.016691 ms\n",
      "=> KP07_dimsplit_16 30.026674 ms\n",
      "=> KP07_dimsplit_32 23.990393 ms\n",
      "=> KP07_dimsplit_64 42.029619 ms\n",
      "=> KP07_dimsplit_128 30.026436 ms\n",
      "=> KP07_dimsplit_256 28.019190 ms\n",
      "=> KP07_dimsplit_512 31.049013 ms\n",
      "=> KP07_dimsplit_1024 36.998034 ms\n",
      "=> KP07_dimsplit_2048 70.047617 ms\n",
      "=> KP07_dimsplit_4096 185.103178 ms\n",
      "=> KP07_dimsplit_16384 4402.130365 ms\n",
      "=> WAF_8 55.039167 ms\n",
      "=> WAF_16 27.032137 ms\n",
      "=> WAF_32 27.012825 ms\n",
      "=> WAF_64 43.028593 ms\n",
      "=> WAF_128 27.040482 ms\n",
      "=> WAF_256 45.030117 ms\n",
      "=> WAF_512 60.044527 ms\n",
      "=> WAF_1024 63.043594 ms\n",
      "=> WAF_2048 89.062691 ms\n",
      "=> WAF_4096 235.169888 ms\n",
      "=> WAF_16384 6115.505934 ms\n",
      "Completed simulation in 53.79148244857788 s\n"
     ]
    }
   ],
   "source": [
    "reference_nx = 16384\n",
    "domain_sizes = [8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]\n",
    "#domain_sizes = [64, 128, 256, 512]\n",
    "\n",
    "simulators = [LxF.LxF, FORCE.FORCE, HLL.HLL, HLL2.HLL2, KP07.KP07, KP07_dimsplit.KP07_dimsplit, WAF.WAF]\n",
    "simulator_ghost_cells = [1, 1, 1, 2, 2, 2, 2]\n",
    "simulator_reference = [LxF.LxF, FORCE.FORCE, HLL.HLL, HLL2.HLL2, KP07.KP07, KP07_dimsplit.KP07_dimsplit, WAF.WAF]\n",
    "simulator_courant_numbers = [0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9]\n",
    "force_rerun=True\n",
    "transpose=True\n",
    "\n",
    "#reference_nx = 4096\n",
    "#domain_sizes = [32, 64, 128, 256, 512, 1024]\n",
    "#simulators = [HLL2.HLL2, WAF.WAF]\n",
    "#simulator_ghost_cells = [2, 2]\n",
    "#simulator_reference = [HLL2.HLL2, WAF.WAF]\n",
    "#simulator_courant_numbers = [0.95, 0.95]\n",
    "#force_rerun=True\n",
    "\n",
    "sim_elapsed_time = np.zeros((len(simulators), len(domain_sizes)))\n",
    "sim_dt = np.zeros((len(simulators), len(domain_sizes)))\n",
    "sim_nt = np.zeros((len(simulators), len(domain_sizes)))\n",
    "\n",
    "#Warmup to compile kernels\n",
    "with Common.Timer(\"warmup\", verbose=False) as timer:\n",
    "    for i in range(len(simulators)):\n",
    "        run_benchmark(\"warmup\", \n",
    "                      simulators[i], \n",
    "                      simulator_courant_numbers[i], \n",
    "                      16, 16, \n",
    "                      simulator_ghost_cells[i], \n",
    "                      force_rerun=force_rerun,\n",
    "                      transpose=transpose)\n",
    "print(\"Completed warmup in {:} s\".format(timer.secs))\n",
    "    \n",
    "with Common.Timer(\"simulation\", verbose=False) as timer:\n",
    "    for i in range(len(simulators)):        \n",
    "        # Run the simulation for all the domain sizes\n",
    "        for j, nx in enumerate(domain_sizes):\n",
    "            datafilename = gen_filename(simulators[i], nx)\n",
    "            [secs, nt, dt] = run_benchmark(datafilename, \n",
    "                          simulators[i], \n",
    "                          simulator_courant_numbers[i], \n",
    "                          nx, reference_nx, \n",
    "                          simulator_ghost_cells[i], \n",
    "                          force_rerun=force_rerun,\n",
    "                          transpose=transpose)\n",
    "            sim_elapsed_time[i, j] = secs\n",
    "            sim_dt[i, j] = dt\n",
    "            sim_nt[i, j] = nt\n",
    "\n",
    "        # Run reference with a low CFL-number\n",
    "        datafilename = gen_filename(simulators[i], reference_nx)\n",
    "        run_benchmark(datafilename, \n",
    "                      simulators[i], \n",
    "                      simulator_courant_numbers[i] * 0.5, \n",
    "                      reference_nx, reference_nx, \n",
    "                      simulator_ghost_cells[i], \n",
    "                      force_rerun=force_rerun,\n",
    "                      transpose=transpose)\n",
    "print(\"Completed simulation in {:} s\".format(timer.secs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7ae5296e10>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 118084 characters]\n",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 43120 characters]\n",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 55348 characters]\n",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "setBwStyles(fig.gca())\n",
    "plt.title(\"Elapsed time vs number of cells\")\n",
    "for i, simulator in enumerate(simulators):\n",
    "    plt.plot(domain_sizes, sim_elapsed_time[i,:], label=simulator.__name__)\n",
    "plt.legend()\n",
    "    \n",
    "fig = plt.figure()\n",
    "setBwStyles(fig.gca())\n",
    "plt.title(\"Dt vs number of cells\")\n",
    "for i, simulator in enumerate(simulators):\n",
    "    plt.loglog(domain_sizes, sim_dt[i,:], label=simulator.__name__)\n",
    "plt.legend()\n",
    "    \n",
    "fig = plt.figure()\n",
    "setBwStyles(fig.gca())\n",
    "plt.title(\"Number of timesteps vs number of cells\")\n",
    "for i, simulator in enumerate(simulators):\n",
    "    plt.plot(domain_sizes, sim_nt[i,:], label=simulator.__name__)\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 135020 characters]\n",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "conservation = np.zeros((len(simulators), len(domain_sizes)))\n",
    "error = np.zeros((len(simulators), len(domain_sizes)))\n",
    "order = np.zeros((len(simulators), len(domain_sizes)-1))\n",
    "\n",
    "for i, simulator in enumerate(simulators):\n",
    "    #Read the reference solution\n",
    "    datafilename = gen_filename(simulator, reference_nx)\n",
    "    with np.load(datafilename) as data:\n",
    "        h_ref = data['h']\n",
    "    \n",
    "    #Compute conservation and error for subsequent resolutions\n",
    "    for j, nx in enumerate(domain_sizes):\n",
    "        #Read current resolution\n",
    "        datafilename = gen_filename(simulator, nx)\n",
    "        with np.load(datafilename) as data:\n",
    "            h = data['h']\n",
    "        \n",
    "        #Downsample reference to current resolution\n",
    "        h_ref_downsampled = downsample(h_ref, reference_nx/nx)\n",
    "        \n",
    "        #Compute error per cell\n",
    "        error[i, j] = np.linalg.norm((h_ref_downsampled - h).flatten(), ord=1) / nx\n",
    "        conservation[i, j] = (np.sum(h_ref) - np.sum(h))\n",
    "    \n",
    "    #Compute the numerical order\n",
    "    a = np.arange(1, len(domain_sizes))\n",
    "    order[i,:] = np.log2(error[i, :-1] / error[i, 1:])\n",
    "    \n",
    "fig = plt.figure()\n",
    "setBwStyles(fig.gca())\n",
    "\n",
    "x = np.linspace(domain_sizes[0], domain_sizes[-1], 100);\n",
    "\n",
    "scaling = np.max(error[:,0]) * domain_sizes[0] * 1.5\n",
    "plt.loglog(x, scaling/x, '-', color='gray', label='Order 1')\n",
    "\n",
    "scaling = np.min(error[:,0]) * domain_sizes[0]**2 * 0.5\n",
    "plt.loglog(x, scaling/(x*x), '-', color='gray', label='Order 2')\n",
    "\n",
    "for i, simulator in enumerate(simulators):\n",
    "    plt.loglog(domain_sizes, error[i,:], label=simulators[i].__name__)\n",
    "    \n",
    "plt.xlabel('Number of cells')\n",
    "plt.ylabel('Error')\n",
    "plt.title(\"Convergence\")\n",
    "plt.legend()\n",
    "save_figure(fig, 'convergence')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>8</th>\n",
       "      <th>16</th>\n",
       "      <th>32</th>\n",
       "      <th>64</th>\n",
       "      <th>128</th>\n",
       "      <th>256</th>\n",
       "      <th>512</th>\n",
       "      <th>1024</th>\n",
       "      <th>2048</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>LxF</th>\n",
       "      <td>2.046098</td>\n",
       "      <td>1.074325</td>\n",
       "      <td>1.834294</td>\n",
       "      <td>1.165476</td>\n",
       "      <td>1.412700</td>\n",
       "      <td>1.432408</td>\n",
       "      <td>1.524218</td>\n",
       "      <td>1.589650</td>\n",
       "      <td>-0.637310</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>FORCE</th>\n",
       "      <td>1.843813</td>\n",
       "      <td>0.964270</td>\n",
       "      <td>1.528037</td>\n",
       "      <td>1.060302</td>\n",
       "      <td>1.226554</td>\n",
       "      <td>1.206903</td>\n",
       "      <td>1.225360</td>\n",
       "      <td>1.553093</td>\n",
       "      <td>1.704214</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>HLL</th>\n",
       "      <td>0.959785</td>\n",
       "      <td>1.222461</td>\n",
       "      <td>1.098614</td>\n",
       "      <td>1.175015</td>\n",
       "      <td>1.078323</td>\n",
       "      <td>1.169317</td>\n",
       "      <td>1.385120</td>\n",
       "      <td>1.705706</td>\n",
       "      <td>0.523056</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>HLL2</th>\n",
       "      <td>2.108427</td>\n",
       "      <td>2.911967</td>\n",
       "      <td>2.415012</td>\n",
       "      <td>1.978696</td>\n",
       "      <td>0.764850</td>\n",
       "      <td>0.089884</td>\n",
       "      <td>0.008152</td>\n",
       "      <td>0.001766</td>\n",
       "      <td>0.001885</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KP07</th>\n",
       "      <td>1.469665</td>\n",
       "      <td>1.925327</td>\n",
       "      <td>1.846630</td>\n",
       "      <td>2.180099</td>\n",
       "      <td>2.103188</td>\n",
       "      <td>0.758878</td>\n",
       "      <td>0.022805</td>\n",
       "      <td>-0.018655</td>\n",
       "      <td>-0.007395</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KP07_dimsplit</th>\n",
       "      <td>2.133675</td>\n",
       "      <td>2.333060</td>\n",
       "      <td>2.385194</td>\n",
       "      <td>1.974368</td>\n",
       "      <td>1.233662</td>\n",
       "      <td>0.222392</td>\n",
       "      <td>0.038976</td>\n",
       "      <td>0.006282</td>\n",
       "      <td>0.002569</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>WAF</th>\n",
       "      <td>2.095899</td>\n",
       "      <td>2.662203</td>\n",
       "      <td>1.905463</td>\n",
       "      <td>2.039412</td>\n",
       "      <td>1.820916</td>\n",
       "      <td>0.573782</td>\n",
       "      <td>0.126937</td>\n",
       "      <td>0.000318</td>\n",
       "      <td>0.000613</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                   8         16        32        64        128       256   \\\n",
       "LxF            2.046098  1.074325  1.834294  1.165476  1.412700  1.432408   \n",
       "FORCE          1.843813  0.964270  1.528037  1.060302  1.226554  1.206903   \n",
       "HLL            0.959785  1.222461  1.098614  1.175015  1.078323  1.169317   \n",
       "HLL2           2.108427  2.911967  2.415012  1.978696  0.764850  0.089884   \n",
       "KP07           1.469665  1.925327  1.846630  2.180099  2.103188  0.758878   \n",
       "KP07_dimsplit  2.133675  2.333060  2.385194  1.974368  1.233662  0.222392   \n",
       "WAF            2.095899  2.662203  1.905463  2.039412  1.820916  0.573782   \n",
       "\n",
       "                   512       1024      2048  \n",
       "LxF            1.524218  1.589650 -0.637310  \n",
       "FORCE          1.225360  1.553093  1.704214  \n",
       "HLL            1.385120  1.705706  0.523056  \n",
       "HLL2           0.008152  0.001766  0.001885  \n",
       "KP07           0.022805 -0.018655 -0.007395  \n",
       "KP07_dimsplit  0.038976  0.006282  0.002569  \n",
       "WAF            0.126937  0.000318  0.000613  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Max</th>\n",
       "      <th>Min</th>\n",
       "      <th>Mean</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>LxF</th>\n",
       "      <td>2.046098</td>\n",
       "      <td>-0.637310</td>\n",
       "      <td>1.271318</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>FORCE</th>\n",
       "      <td>1.843813</td>\n",
       "      <td>0.964270</td>\n",
       "      <td>1.368061</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>HLL</th>\n",
       "      <td>1.705706</td>\n",
       "      <td>0.523056</td>\n",
       "      <td>1.146377</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>HLL2</th>\n",
       "      <td>2.911967</td>\n",
       "      <td>0.001766</td>\n",
       "      <td>1.142293</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KP07</th>\n",
       "      <td>2.180099</td>\n",
       "      <td>-0.018655</td>\n",
       "      <td>1.142283</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KP07_dimsplit</th>\n",
       "      <td>2.385194</td>\n",
       "      <td>0.002569</td>\n",
       "      <td>1.147798</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>WAF</th>\n",
       "      <td>2.662203</td>\n",
       "      <td>0.000318</td>\n",
       "      <td>1.247282</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                    Max       Min      Mean\n",
       "LxF            2.046098 -0.637310  1.271318\n",
       "FORCE          1.843813  0.964270  1.368061\n",
       "HLL            1.705706  0.523056  1.146377\n",
       "HLL2           2.911967  0.001766  1.142293\n",
       "KP07           2.180099 -0.018655  1.142283\n",
       "KP07_dimsplit  2.385194  0.002569  1.147798\n",
       "WAF            2.662203  0.000318  1.247282"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from IPython.display import display\n",
    "\n",
    "table = pd.DataFrame(order, columns=domain_sizes[:-1], index=[simulator.__name__ for simulator in simulators])\n",
    "display(table)\n",
    "\n",
    "order_summary = pd.DataFrame(np.empty((len(simulators), 3)), columns=['Max', 'Min', 'Mean'], index=[simulator.__name__ for simulator in simulators])\n",
    "for i, simulator in enumerate(simulators):\n",
    "    order_summary.loc[simulator.__name__] = [np.max(order[i,:]), \n",
    "        np.min(order[i,:]), \n",
    "        np.mean(order[i,:])]\n",
    "display(order_summary)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 151448 characters]\n",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "setBwStyles(fig.gca())\n",
    "\n",
    "for i, simulator in enumerate(simulators):\n",
    "    plt.plot(order[i,:], label=simulator.__name__)\n",
    "    \n",
    "plt.ylim([0, 3])\n",
    "plt.xlabel('Number of cells')\n",
    "plt.ylabel('Order')\n",
    "plt.legend()\n",
    "plt.title(\"Numerical order\")\n",
    "save_figure(fig, 'numerical_order')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_solution(simulator, nx, label, **kwargs):\n",
    "    datafilename = gen_filename(simulator, nx)\n",
    "    \n",
    "    #Read the solution\n",
    "    with np.load(datafilename) as data:\n",
    "        h = data['h']\n",
    "        \n",
    "    x = np.linspace(0.5, nx-0.5, nx)*100.0/float(nx)\n",
    "    \n",
    "    plt.plot(x, h, label=label, **kwargs)\n",
    "    \n",
    "    h = None\n",
    "    x = None\n",
    "    gc.collect() # Force run garbage collection to free up memory\n",
    "    \n",
    "def plot_comparison(nx, **kwargs):\n",
    "    plot_solution(HLL2.HLL2, reference_nx, 'Reference', marker=' ', linestyle='-')\n",
    "\n",
    "    for i, simulator in enumerate(simulators):\n",
    "        plot_solution(simulator, nx, simulator.__name__, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 189784 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "boxes = np.array([\\\n",
    "         [[28, 35], [0.499, 0.516]], \\\n",
    "         [[45, 56], [0.58, 0.60]]\n",
    "         ])\n",
    "comparison_nx = 32\n",
    "\n",
    "fig = plt.figure(figsize=(12, 12))\n",
    "gs = gridspec.GridSpec(2, 2)\n",
    "    \n",
    "ax1 = fig.add_subplot(gs[0,:])\n",
    "setBwStyles(ax1)\n",
    "for simulator in simulators:\n",
    "    plot_solution(simulator, comparison_nx, simulator.__name__, marker=' ')\n",
    "for box_x, box_y in boxes:\n",
    "    x = np.kron(box_x, np.ones((2)))\n",
    "    y = np.hstack((box_y, box_y[::-1]))\n",
    "    fig.gca().fill(x, y, fill=True, linestyle='-', color='gray', alpha=0.5)\n",
    "plt.title(\"Comparison of solutions at nx=\" + str(comparison_nx))\n",
    "plt.legend()\n",
    "\n",
    "ax2 = fig.add_subplot(gs[1,0])\n",
    "setBwStyles(ax2)\n",
    "for simulator in simulators:\n",
    "    plot_solution(simulator, comparison_nx, simulator.__name__, markevery=2)\n",
    "plt.xlim(boxes[0, 0])\n",
    "plt.ylim(boxes[0, 1])\n",
    "\n",
    "ax3 = fig.add_subplot(gs[1,1])\n",
    "setBwStyles(ax3)\n",
    "setBwStyles(ax2)\n",
    "for simulator in simulators:\n",
    "    plot_solution(simulator, comparison_nx, simulator.__name__, markevery=2)\n",
    "plt.xlim(boxes[1, 0])\n",
    "plt.ylim(boxes[1, 1])\n",
    "    \n",
    "save_figure(fig, 'solution')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 104328 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(12, 12))\n",
    "gs = gridspec.GridSpec(2, 2)\n",
    "    \n",
    "ax1 = fig.add_subplot(gs[0,:])\n",
    "setBwStyles(ax1)\n",
    "for simulator in simulators:\n",
    "    plot_solution(simulator, reference_nx, simulator.__name__, marker=' ')\n",
    "for box_x, box_y in boxes:\n",
    "    x = np.kron(box_x, np.ones((2)))\n",
    "    y = np.hstack((box_y, box_y[::-1]))\n",
    "    fig.gca().fill(x, y, fill=True, linestyle='-', color='gray', alpha=0.5)\n",
    "plt.title(\"Comparison of reference solutions\")\n",
    "plt.legend()\n",
    "\n",
    "ax2 = fig.add_subplot(gs[1,0])\n",
    "setBwStyles(ax2)\n",
    "for simulator in simulators:\n",
    "    plot_solution(simulator, reference_nx, simulator.__name__, marker=' ')\n",
    "plt.xlim(boxes[0, 0])\n",
    "plt.ylim(boxes[0, 1])\n",
    "\n",
    "ax3 = fig.add_subplot(gs[1,1])\n",
    "setBwStyles(ax3)\n",
    "setBwStyles(ax2)\n",
    "for simulator in simulators:\n",
    "    plot_solution(simulator, reference_nx, simulator.__name__, marker=' ')\n",
    "plt.xlim(boxes[1, 0])\n",
    "plt.ylim(boxes[1, 1])\n",
    "    \n",
    "save_figure(fig, 'reference_solutions')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 213968 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 224112 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 214892 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 189816 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 195344 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 193696 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 190508 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for simulator in simulators:\n",
    "    fig = plt.figure(figsize=(12, 12))\n",
    "    \n",
    "    gs = gridspec.GridSpec(2, 2)\n",
    "    \n",
    "    ax1 = fig.add_subplot(gs[0,:])\n",
    "    setBwStyles(ax1)\n",
    "    plot_solution(simulator, reference_nx, 'Reference', linestyle='-', marker=' ')\n",
    "    for j, nx in enumerate(domain_sizes):\n",
    "        plot_solution(simulator, nx, 'nx=' + str(nx), marker=' ')\n",
    "    for box_x, box_y in boxes:\n",
    "        x = np.kron(box_x, np.ones((2)))\n",
    "        y = np.hstack((box_y, box_y[::-1]))\n",
    "        fig.gca().fill(x, y, fill=True, linestyle='-', color='gray', alpha=0.5)\n",
    "    plt.title(str(simulator.__name__))\n",
    "    plt.legend()\n",
    "    \n",
    "    ax2 = fig.add_subplot(gs[1,0])\n",
    "    setBwStyles(ax2)\n",
    "    plot_solution(simulator, reference_nx, 'Reference', linestyle='-', marker=' ')\n",
    "    for j, nx in enumerate(domain_sizes):\n",
    "        plot_solution(simulator, nx, 'nx=' + str(nx), marker=' ')  \n",
    "    plt.xlim(boxes[0, 0])\n",
    "    plt.ylim(boxes[0, 1])\n",
    "    \n",
    "    ax3 = fig.add_subplot(gs[1,1])\n",
    "    setBwStyles(ax3)\n",
    "    plot_solution(simulator, reference_nx, 'Reference', linestyle='-', marker=' ')\n",
    "    for j, nx in enumerate(domain_sizes):\n",
    "        plot_solution(simulator, nx, 'nx=' + str(nx), marker=' ')   \n",
    "    plt.xlim(boxes[1, 0])\n",
    "    plt.ylim(boxes[1, 1])\n",
    "    \n",
    "    save_figure(fig, simulator.__name__)"
   ]
  },
  {
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
