{
 "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 <487316142544>\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_smooth2d_\" + 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": {},
   "outputs": [],
   "source": [
    "def downsample(highres_solution, factor):\n",
    "    assert(highres_solution.shape[0] % factor == 0)\n",
    "    assert(highres_solution.shape[1] % factor == 0)\n",
    "    ny = highres_solution.shape[0] / factor\n",
    "    nx = highres_solution.shape[1] / factor\n",
    "    \n",
    "    if (factor == 1):\n",
    "        return highres_solution\n",
    "    else:\n",
    "        return highres_solution.reshape([int(ny), int(factor), int(nx), int(factor)]).mean(3).mean(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 22336 characters]\n",
      "text/plain": [
       "<Figure size 1200x800 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.0, 2.0, 0.2726356370001644]\n"
     ]
    }
   ],
   "source": [
    "def wall_boundary_conditions(data):\n",
    "    data[0:2,:] = data[4:2:-1,:]\n",
    "    data[-2:,:] = data[-3:-5:-1,:]\n",
    "    data[:,0:2] = data[:,4:2:-1]\n",
    "    data[:,-2:] = data[:,-3:-5:-1]\n",
    "    return data\n",
    "    \n",
    "\n",
    "def gen_test_data(nx, ny, ref_nx, ref_ny, g, num_ghost_cells):\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 + yv**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",
    "    assert(nx/ref_nx == ny/ref_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",
    "    #hv_highres = 0.1*np.exp(-(xv**2/size + yv**2/size))\n",
    "    hv_highres = 0.1*0.5*(1.0 + np.cos(np.pi*r/size)) * (r < size)\n",
    "    hv[num_ghost_cells:-num_ghost_cells, num_ghost_cells:-num_ghost_cells] = downsample(hv_highres, ref_nx/nx)\n",
    "    hv_highres = None\n",
    "    gc.collect()\n",
    "    \n",
    "    #Fix boundary conditions\n",
    "    h = wall_boundary_conditions(h);\n",
    "    hu = wall_boundary_conditions(hu);\n",
    "    hv = wall_boundary_conditions(hv);\n",
    "    \n",
    "    scale = 0.7\n",
    "    max_h_estimate = 0.6\n",
    "    max_u_estimate = 0.1*np.sqrt(2.0)\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, ny=50, ref_nx=1000, ref_ny=500, 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": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_filename(simulator, nx):\n",
    "    return \"shock1d_\" + str(simulator.__name__) + \"_\" + str(nx) + \".npz\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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(simulator, courant_number, nx, ny, ref_nx, ref_ny, num_ghost_cells, g=9.81, force_rerun=False, datafilename=None):\n",
    "\n",
    "    if (datafilename and os.path.isfile(datafilename) and force_rerun == False):\n",
    "        print(\"=== WARNING ===\")\n",
    "        print(\"Existing simulation exists, skipping simulation\")\n",
    "        print(\"=== WARNING ===\")\n",
    "        return[0, 0, 0]\n",
    "    else:\n",
    "        #Generate high-res initial conditions\n",
    "        h0, hu0, hv0, dx, dy, dt = gen_test_data(nx=nx, ny=ny, ref_nx=ref_nx, ref_ny=ref_ny, g=g, num_ghost_cells=num_ghost_cells)\n",
    "        dt = dt * courant_number\n",
    "\n",
    "        #Initialize simulator\n",
    "        with Common.Timer(simulator.__name__ + \"_\" + str(nx)) as timer:\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",
    "        h = h[num_ghost_cells:-num_ghost_cells, num_ghost_cells:-num_ghost_cells]\n",
    "        hu = hu[num_ghost_cells:-num_ghost_cells, num_ghost_cells:-num_ghost_cells]\n",
    "        hv = hv[num_ghost_cells:-num_ghost_cells, num_ghost_cells:-num_ghost_cells]\n",
    "\n",
    "        sanity_check(h, \"h\")\n",
    "        sanity_check(hu, \"hu\")\n",
    "        sanity_check(hv, \"hv\")\n",
    "\n",
    "        if (datafilename):\n",
    "            np.savez_compressed(datafilename, h=h, hu=hu, hv=hv, nt=nt, t=t, elapsed=timer.secs)\n",
    "        else:\n",
    "            print(\"=== WARNING ===\")\n",
    "            print(\"No datafilename supplied - no data saved to disk!\")\n",
    "            print(\"=== WARNING ===\")\n",
    "            \n",
    "        \n",
    "    gc.collect() # Force run garbage collection to free up memory\n",
    "    \n",
    "    return [timer.secs, nt, dt]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=> KP07_50 6078.819036 ms\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 17260 characters]\n",
      "text/plain": [
       "<Figure size 1200x800 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "3354"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "datafilename = \"temp_data.npz\"\n",
    "elapsed, nt, dt = run_benchmark(KP07.KP07, \n",
    "                                courant_number=0.9, \n",
    "                                ref_nx=500, ref_ny=1000,\n",
    "                                nx=50, ny=100, \n",
    "                                num_ghost_cells=2, force_rerun=True,\n",
    "                                datafilename=datafilename)\n",
    "\n",
    "with np.load(datafilename) as data:\n",
    "    rh = downsample(data['h'], 2)\n",
    "\n",
    "    plt.figure()\n",
    "    plt.imshow(rh, extent=[0, 100, 0, 100])\n",
    "    plt.colorbar()\n",
    "    plt.show()\n",
    "    \n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warmup\n",
      "=> LxF_16 7164.669275 ms\n",
      "=== WARNING ===\n",
      "No datafilename supplied - no data saved to disk!\n",
      "=== WARNING ===\n",
      "=> FORCE_16 6598.121405 ms\n",
      "=== WARNING ===\n",
      "No datafilename supplied - no data saved to disk!\n",
      "=== WARNING ===\n",
      "=> HLL_16 3747.835875 ms\n",
      "=== WARNING ===\n",
      "No datafilename supplied - no data saved to disk!\n",
      "=== WARNING ===\n",
      "=> HLL2_16 5513.680935 ms\n",
      "=== WARNING ===\n",
      "No datafilename supplied - no data saved to disk!\n",
      "=== WARNING ===\n",
      "=> WAF_16 5508.579493 ms\n",
      "=== WARNING ===\n",
      "No datafilename supplied - no data saved to disk!\n",
      "=== WARNING ===\n",
      "=> KP07_16 18.017292 ms\n",
      "=== WARNING ===\n",
      "No datafilename supplied - no data saved to disk!\n",
      "=== WARNING ===\n",
      "=> KP07_dimsplit_16 5265.863419 ms\n",
      "=== WARNING ===\n",
      "No datafilename supplied - no data saved to disk!\n",
      "=== WARNING ===\n",
      "Simulating\n",
      "=> LxF_8 14.018774 ms\n",
      "=> LxF_16 14.011145 ms\n",
      "=> LxF_32 14.012098 ms\n",
      "=> LxF_64 15.018463 ms\n",
      "=> LxF_128 17.013311 ms\n",
      "=> LxF_256 23.023605 ms\n",
      "=> LxF_512 64.050913 ms\n",
      "=> LxF_1024 277.214766 ms\n",
      "=> LxF_2048 1894.445181 ms\n",
      "=> FORCE_8 9.007215 ms\n",
      "=> FORCE_16 8.005381 ms\n",
      "=> FORCE_32 7.007837 ms\n",
      "=> FORCE_64 8.007288 ms\n",
      "=> FORCE_128 8.013725 ms\n",
      "=> FORCE_256 12.010098 ms\n",
      "=> FORCE_512 58.043480 ms\n",
      "=> FORCE_1024 337.255001 ms\n",
      "=> FORCE_2048 2454.854488 ms\n",
      "=> HLL_8 11.008263 ms\n",
      "=> HLL_16 8.006096 ms\n",
      "=> HLL_32 8.006096 ms\n",
      "=> HLL_64 8.005619 ms\n",
      "=> HLL_128 8.006334 ms\n",
      "=> HLL_256 16.012430 ms\n",
      "=> HLL_512 70.054770 ms\n",
      "=> HLL_1024 426.322937 ms\n",
      "=> HLL_2048 3196.443558 ms\n",
      "=> HLL2_8 11.011600 ms\n",
      "=> HLL2_16 14.012337 ms\n",
      "=> HLL2_32 14.011621 ms\n",
      "=> HLL2_64 14.011383 ms\n",
      "=> HLL2_128 19.015312 ms\n",
      "=> HLL2_256 29.027224 ms\n",
      "=> HLL2_512 118.119001 ms\n",
      "=> HLL2_1024 663.526297 ms\n",
      "=> HLL2_2048 4924.751043 ms\n",
      "=> WAF_8 11.008024 ms\n",
      "=> WAF_16 11.008978 ms\n",
      "=> WAF_32 11.003256 ms\n",
      "=> WAF_64 12.010813 ms\n",
      "=> WAF_128 14.012575 ms\n",
      "=> WAF_256 33.050776 ms\n",
      "=> WAF_512 139.111280 ms\n",
      "=> WAF_1024 833.629370 ms\n",
      "=> WAF_2048 6163.692236 ms\n",
      "=> KP07_8 12.007236 ms\n",
      "=> KP07_16 14.010429 ms\n",
      "=> KP07_32 15.012741 ms\n",
      "=> KP07_64 13.002872 ms\n",
      "=> KP07_128 18.011808 ms\n",
      "=> KP07_256 51.032782 ms\n",
      "=> KP07_512 268.163919 ms\n",
      "=> KP07_1024 1841.123343 ms\n",
      "=> KP07_2048 14106.599092 ms\n",
      "=> KP07_dimsplit_8 11.006832 ms\n",
      "=> KP07_dimsplit_16 15.008926 ms\n",
      "=> KP07_dimsplit_32 11.011600 ms\n",
      "=> KP07_dimsplit_64 13.009310 ms\n",
      "=> KP07_dimsplit_128 15.004873 ms\n",
      "=> KP07_dimsplit_256 30.048847 ms\n",
      "=> KP07_dimsplit_512 109.066963 ms\n",
      "=> KP07_dimsplit_1024 611.401081 ms\n",
      "=> KP07_dimsplit_2048 4460.719347 ms\n"
     ]
    }
   ],
   "source": [
    "reference_nx = 2048\n",
    "force_rerun=True\n",
    "domain_sizes = [8, 16, 32, 64, 128, 256, 512, 1024]\n",
    "simulators = [LxF.LxF, FORCE.FORCE, HLL.HLL, HLL2.HLL2, WAF.WAF, KP07.KP07, KP07_dimsplit.KP07_dimsplit]\n",
    "simulator_ghost_cells = [1, 1, 1, 2, 2, 2, 2]\n",
    "simulator_names = ['LxF', 'FORCE', 'HLL', 'HLL2', 'WAF', 'KP07', 'KP07 dimsplit']\n",
    "simulator_courant_numbers = [0.9, 0.9, 0.9, 0.9, 0.9, 0.45, 0.9]\n",
    "\n",
    "\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",
    "print(\"Warmup\")\n",
    "for i in range(len(simulators)):\n",
    "    run_benchmark(simulator=simulators[i], \n",
    "                  courant_number=simulator_courant_numbers[i], \n",
    "                  nx=16, ny=16, ref_nx=16, ref_ny=16,\n",
    "                  num_ghost_cells=simulator_ghost_cells[i], \n",
    "                  force_rerun=force_rerun)\n",
    "    \n",
    "print(\"Simulating\")\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(simulator=simulators[i], \n",
    "                      courant_number=simulator_courant_numbers[i], \n",
    "                      nx=nx, ny=nx, ref_nx=reference_nx, ref_ny=reference_nx,\n",
    "                      num_ghost_cells=simulator_ghost_cells[i], \n",
    "                      force_rerun=force_rerun, \n",
    "                      datafilename=datafilename)\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(simulator=simulators[i], \n",
    "                  courant_number=simulator_courant_numbers[i], \n",
    "                  nx=reference_nx, ny=reference_nx, ref_nx=reference_nx, ref_ny=reference_nx,\n",
    "                  num_ghost_cells=simulator_ghost_cells[i], \n",
    "                  force_rerun=force_rerun, \n",
    "                  datafilename=datafilename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x710002a0f0>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 90240 characters]\n",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 56780 characters]\n",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 66200 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": 12,
   "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",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>LxF</th>\n",
       "      <td>2.109998</td>\n",
       "      <td>1.117344</td>\n",
       "      <td>1.770541</td>\n",
       "      <td>1.130174</td>\n",
       "      <td>1.398593</td>\n",
       "      <td>1.459872</td>\n",
       "      <td>1.585554</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>FORCE</th>\n",
       "      <td>1.803550</td>\n",
       "      <td>1.058752</td>\n",
       "      <td>1.424818</td>\n",
       "      <td>0.994047</td>\n",
       "      <td>1.264049</td>\n",
       "      <td>1.354103</td>\n",
       "      <td>1.598773</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>HLL</th>\n",
       "      <td>1.271458</td>\n",
       "      <td>1.218233</td>\n",
       "      <td>1.018041</td>\n",
       "      <td>1.023368</td>\n",
       "      <td>1.120194</td>\n",
       "      <td>1.259487</td>\n",
       "      <td>1.649370</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>HLL2</th>\n",
       "      <td>2.018988</td>\n",
       "      <td>2.466334</td>\n",
       "      <td>0.904048</td>\n",
       "      <td>0.650843</td>\n",
       "      <td>0.872441</td>\n",
       "      <td>1.026186</td>\n",
       "      <td>2.194724</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>WAF</th>\n",
       "      <td>2.016446</td>\n",
       "      <td>2.483785</td>\n",
       "      <td>1.021645</td>\n",
       "      <td>0.705272</td>\n",
       "      <td>0.891656</td>\n",
       "      <td>0.858559</td>\n",
       "      <td>2.275014</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KP07</th>\n",
       "      <td>1.560495</td>\n",
       "      <td>2.067524</td>\n",
       "      <td>1.738482</td>\n",
       "      <td>2.077667</td>\n",
       "      <td>2.065662</td>\n",
       "      <td>2.206142</td>\n",
       "      <td>2.172162</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KP07_dimsplit</th>\n",
       "      <td>1.896522</td>\n",
       "      <td>2.258217</td>\n",
       "      <td>1.073543</td>\n",
       "      <td>0.676780</td>\n",
       "      <td>0.846475</td>\n",
       "      <td>1.155170</td>\n",
       "      <td>2.138188</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                    8         16        32        64        128       256  \\\n",
       "LxF            2.109998  1.117344  1.770541  1.130174  1.398593  1.459872   \n",
       "FORCE          1.803550  1.058752  1.424818  0.994047  1.264049  1.354103   \n",
       "HLL            1.271458  1.218233  1.018041  1.023368  1.120194  1.259487   \n",
       "HLL2           2.018988  2.466334  0.904048  0.650843  0.872441  1.026186   \n",
       "WAF            2.016446  2.483785  1.021645  0.705272  0.891656  0.858559   \n",
       "KP07           1.560495  2.067524  1.738482  2.077667  2.065662  2.206142   \n",
       "KP07_dimsplit  1.896522  2.258217  1.073543  0.676780  0.846475  1.155170   \n",
       "\n",
       "                    512  \n",
       "LxF            1.585554  \n",
       "FORCE          1.598773  \n",
       "HLL            1.649370  \n",
       "HLL2           2.194724  \n",
       "WAF            2.275014  \n",
       "KP07           2.172162  \n",
       "KP07_dimsplit  2.138188  "
      ]
     },
     "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",
       "      <th>Stdev</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>LxF</th>\n",
       "      <td>2.109998</td>\n",
       "      <td>1.117344</td>\n",
       "      <td>1.510297</td>\n",
       "      <td>0.326663</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>FORCE</th>\n",
       "      <td>1.803550</td>\n",
       "      <td>0.994047</td>\n",
       "      <td>1.356870</td>\n",
       "      <td>0.265059</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>HLL</th>\n",
       "      <td>1.649370</td>\n",
       "      <td>1.018041</td>\n",
       "      <td>1.222879</td>\n",
       "      <td>0.199257</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>HLL2</th>\n",
       "      <td>2.466334</td>\n",
       "      <td>0.650843</td>\n",
       "      <td>1.447652</td>\n",
       "      <td>0.692941</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>WAF</th>\n",
       "      <td>2.483785</td>\n",
       "      <td>0.705272</td>\n",
       "      <td>1.464625</td>\n",
       "      <td>0.703906</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KP07</th>\n",
       "      <td>2.206142</td>\n",
       "      <td>1.560495</td>\n",
       "      <td>1.984019</td>\n",
       "      <td>0.222584</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>KP07_dimsplit</th>\n",
       "      <td>2.258217</td>\n",
       "      <td>0.676780</td>\n",
       "      <td>1.434985</td>\n",
       "      <td>0.599449</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                    Max       Min      Mean     Stdev\n",
       "LxF            2.109998  1.117344  1.510297  0.326663\n",
       "FORCE          1.803550  0.994047  1.356870  0.265059\n",
       "HLL            1.649370  1.018041  1.222879  0.199257\n",
       "HLL2           2.466334  0.650843  1.447652  0.692941\n",
       "WAF            2.483785  0.705272  1.464625  0.703906\n",
       "KP07           2.206142  1.560495  1.984019  0.222584\n",
       "KP07_dimsplit  2.258217  0.676780  1.434985  0.599449"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from IPython.display import display\n",
    "\n",
    "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*nx)\n",
    "        conservation[i, j] = (np.sum(h_ref_downsampled) - 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",
    "\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), 4)), columns=['Max', 'Min', 'Mean', 'Stdev'], 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",
    "        np.std(order[i,:])]\n",
    "display(order_summary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 152548 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",
    "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": [],
   "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",
    "        h = h[int(nx/2),:]\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": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 181512 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": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 111008 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": 17,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 200684 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 200784 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 190056 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 177736 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 176184 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 178876 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 180212 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",
    "\n",
    "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
}
