{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Starting cluster\n",
    "\n",
    "## Prerequisites\n",
    "First, you need to install MPI, on windows use MS-MPI:\n",
    "https://msdn.microsoft.com/en-us/library/bb524831(v=vs.85).aspx\n",
    "\n",
    "\n",
    "## With a profile (not working)\n",
    "In theory, you should be able to create a profile using\n",
    "```\n",
    "ipython profile create --parallel --profile=myprofile\n",
    "```\n",
    "and then set\n",
    "```\n",
    "c.IPClusterEngines.engine_launcher_class = 'MPIEngineSetLauncher'\n",
    "```\n",
    "in ```<IPYTHON-DIR>/profile_myprofile/ipcluster_config.py```. This should then enable you to start a cluster using\n",
    "```\n",
    "ipcluster start --profile=myprofile\n",
    "```\n",
    "or alternatively through the Clusters tab in Jupyter\n",
    "\n",
    "\n",
    "## Without a profile (not working)\n",
    "An alternative is to run\n",
    "```\n",
    "ipcluster start --engines=MPI\n",
    "```\n",
    "\n",
    "\n",
    "## Manual start (working)\n",
    "This, however, does *not* work for me on Windows. What does work is the following:\n",
    "\n",
    "Start a controller using\n",
    "```\n",
    "ipcontroller --ip='*'\n",
    "```\n",
    "and then start several engines using mpiexec:\n",
    "```\n",
    "mpiexec -n 4 ipengine --mpi\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from GPUSimulators import IPythonMagic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Console logger using level INFO\n",
      "File logger using level DEBUG to mpi.log\n",
      "Python version 3.6.9 |Anaconda, Inc.| (default, Jul 30 2019, 14:00:49) [MSC v.1915 64 bit (AMD64)]\n"
     ]
    }
   ],
   "source": [
    "%setup_logging --out mpi.log my_logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Starting IPController\n",
      "Starting IPEngines\n",
      "Waiting for cluster...\n",
      "Waiting for cluster...\n",
      "Waiting for cluster...\n",
      "Waiting for cluster...\n",
      "Waiting for cluster...\n",
      "Waiting for cluster...\n",
      "Waiting for cluster...\n",
      "Waiting for cluster...\n",
      "Done\n"
     ]
    }
   ],
   "source": [
    "%setup_mpi mpi_context --num_engines 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from mpi4py import MPI\n",
    "import time\n",
    "import json\n",
    "\n",
    "from GPUSimulators import IPythonMagic, MPISimulator, Common"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[stderr:0] \n",
      "Console logger using level INFO\n",
      "File logger using level DEBUG to mpi_0.log\n",
      "Python version 3.6.9 |Anaconda, Inc.| (default, Jul 30 2019, 14:00:49) [MSC v.1915 64 bit (AMD64)]\n",
      "[stderr:1] \n",
      "Console logger using level INFO\n",
      "File logger using level DEBUG to mpi_1.log\n",
      "Python version 3.6.9 |Anaconda, Inc.| (default, Jul 30 2019, 14:00:49) [MSC v.1915 64 bit (AMD64)]\n",
      "[stderr:2] \n",
      "Console logger using level INFO\n",
      "File logger using level DEBUG to mpi_2.log\n",
      "Python version 3.6.9 |Anaconda, Inc.| (default, Jul 30 2019, 14:00:49) [MSC v.1915 64 bit (AMD64)]\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "%setup_logging --out \"'mpi_' + str(MPI.COMM_WORLD.rank) + '.log'\" my_logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[stderr:0] \n",
      "Registering my_context in user workspace\n",
      "PyCUDA version 2018.1.1\n",
      "CUDA version (10, 0, 0)\n",
      "Driver version 10000\n",
      "Using device 0/1 'GeForce 840M' (0000:03:00.0) GPU\n",
      "Created context handle <650919533984>\n",
      "Using CUDA cache dir C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\GPUSimulators\\cuda_cache\n",
      "Autotuning enabled. It may take several minutes to run the code the first time: have patience\n",
      "[stderr:1] \n",
      "Registering my_context in user workspace\n",
      "PyCUDA version 2018.1.1\n",
      "CUDA version (10, 0, 0)\n",
      "Driver version 10000\n",
      "Using device 0/1 'GeForce 840M' (0000:03:00.0) GPU\n",
      "Created context handle <669212383056>\n",
      "Using CUDA cache dir C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\GPUSimulators\\cuda_cache\n",
      "Autotuning enabled. It may take several minutes to run the code the first time: have patience\n",
      "[stderr:2] \n",
      "Registering my_context in user workspace\n",
      "PyCUDA version 2018.1.1\n",
      "CUDA version (10, 0, 0)\n",
      "Driver version 10000\n",
      "Using device 0/1 'GeForce 840M' (0000:03:00.0) GPU\n",
      "Created context handle <553691644496>\n",
      "Using CUDA cache dir C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\GPUSimulators\\cuda_cache\n",
      "Autotuning enabled. It may take several minutes to run the code the first time: have patience\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "%cuda_context_handler my_context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "def plotExtent(extent, *args, **kwargs):\n",
    "    x0, x1, y0, y1 = extent\n",
    "    plt.plot([x0, x1, x1, x0, x0], [y0, y0, y1, y1, y0], *args, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] 0\n",
      "[stdout:1] 1\n",
      "[stdout:2] 2\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "\n",
    "import importlib\n",
    "importlib.reload(MPISimulator)\n",
    "\n",
    "grid = MPISimulator.MPIGrid(MPI.COMM_WORLD)\n",
    "print(grid.getLocalRank())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] 0\n",
      "[stdout:1] 1\n",
      "[stdout:2] 2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[stderr:0] \n",
      "Constructed in 3.840158462524414 seconds\n",
      "Initialized C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_0_0012.nc\n",
      "Opening C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_0_0012.nc\n",
      "Keyword arguments: {'mode': 'w', 'clobber': False}\n",
      "JSON: Converting boundary_conditions to string ([north=Type.Periodic, south=Type.Periodic, east=Type.Periodic, west=Type.Periodic])\n",
      "JSON: Converting context to string (CudaContext id 650919533984)\n",
      "JSON: Converting grid to string (<GPUSimulators.MPISimulator.MPIGrid object at 0x000000979166F4A8>)\n",
      "MPISimulator [128x128]: 0% [####==========================] 100%. Total: 33s, elapsed: 5s, remaining: 28s\n",
      "MPISimulator [128x128]: 0% [##########################====] 100%. Total: 11s, elapsed: 10s, remaining: 1s\n",
      "MPISimulator [128x128]: 0% [############==================] 100%. Total: 12s, elapsed: 5s, remaining: 7s\n",
      "MPISimulator [128x128]: 0% [########################======] 100%. Total: 12s, elapsed: 10s, remaining: 2s\n",
      "MPISimulator [128x128]: 0% [################==============] 100%. Total: 9s, elapsed: 5s, remaining: 4s\n",
      "MPISimulator [128x128]: 0% [###################===========] 100%. Total: 8s, elapsed: 5s, remaining: 3s\n",
      "MPISimulator [128x128]: 0% [##################============] 100%. Total: 8s, elapsed: 5s, remaining: 3s\n",
      "MPISimulator [128x128]: 0% [##################============] 100%. Total: 8s, elapsed: 5s, remaining: 3s\n",
      "MPISimulator [128x128]: 0% [######################========] 100%. Total: 7s, elapsed: 5s, remaining: 1s\n",
      "MPISimulator [128x128]: 0% [#######################=======] 100%. Total: 6s, elapsed: 5s, remaining: 1s\n",
      "MPISimulator [128x128]: 0% [#####################=========] 100%. Total: 6s, elapsed: 5s, remaining: 1s\n",
      "Closing C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_0_0012.nc\n",
      "[stderr:1] \n",
      "Constructed in 3.6445236206054688 seconds\n",
      "Initialized C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_1_0012.nc\n",
      "Opening C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_1_0012.nc\n",
      "Keyword arguments: {'mode': 'w', 'clobber': False}\n",
      "JSON: Converting boundary_conditions to string ([north=Type.Periodic, south=Type.Periodic, east=Type.Periodic, west=Type.Periodic])\n",
      "JSON: Converting context to string (CudaContext id 669212383056)\n",
      "JSON: Converting grid to string (<GPUSimulators.MPISimulator.MPIGrid object at 0x0000009BD17DF438>)\n",
      "MPISimulator [128x128]: 0% [####==========================] 100%. Total: 35s, elapsed: 5s, remaining: 30s\n",
      "MPISimulator [128x128]: 0% [##########################====] 100%. Total: 11s, elapsed: 10s, remaining: 1s\n",
      "MPISimulator [128x128]: 0% [############==================] 100%. Total: 12s, elapsed: 5s, remaining: 7s\n",
      "MPISimulator [128x128]: 0% [########################======] 100%. Total: 12s, elapsed: 10s, remaining: 2s\n",
      "MPISimulator [128x128]: 0% [################==============] 100%. Total: 9s, elapsed: 5s, remaining: 4s\n",
      "MPISimulator [128x128]: 0% [###################===========] 100%. Total: 8s, elapsed: 5s, remaining: 3s\n",
      "MPISimulator [128x128]: 0% [##################============] 100%. Total: 8s, elapsed: 5s, remaining: 3s\n",
      "MPISimulator [128x128]: 0% [##################============] 100%. Total: 8s, elapsed: 5s, remaining: 3s\n",
      "MPISimulator [128x128]: 0% [#####################=========] 100%. Total: 6s, elapsed: 5s, remaining: 1s\n",
      "MPISimulator [128x128]: 0% [#######################=======] 100%. Total: 6s, elapsed: 5s, remaining: 1s\n",
      "MPISimulator [128x128]: 0% [#####################=========] 100%. Total: 7s, elapsed: 5s, remaining: 1s\n",
      "Closing C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_1_0012.nc\n",
      "[stderr:2] \n",
      "Constructed in 3.6400203704833984 seconds\n",
      "Initialized C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_2_0004.nc\n",
      "Opening C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_2_0004.nc\n",
      "Keyword arguments: {'mode': 'w', 'clobber': False}\n",
      "JSON: Converting boundary_conditions to string ([north=Type.Periodic, south=Type.Periodic, east=Type.Periodic, west=Type.Periodic])\n",
      "JSON: Converting context to string (CudaContext id 553691644496)\n",
      "JSON: Converting grid to string (<GPUSimulators.MPISimulator.MPIGrid object at 0x00000080EE2ED3C8>)\n",
      "MPISimulator [128x128]: 0% [####==========================] 100%. Total: 35s, elapsed: 5s, remaining: 30s\n",
      "MPISimulator [128x128]: 0% [##########################====] 100%. Total: 11s, elapsed: 10s, remaining: 1s\n",
      "MPISimulator [128x128]: 0% [############==================] 100%. Total: 12s, elapsed: 5s, remaining: 7s\n",
      "MPISimulator [128x128]: 0% [########################======] 100%. Total: 12s, elapsed: 10s, remaining: 2s\n",
      "MPISimulator [128x128]: 0% [################==============] 100%. Total: 9s, elapsed: 5s, remaining: 4s\n",
      "MPISimulator [128x128]: 0% [###################===========] 100%. Total: 8s, elapsed: 5s, remaining: 3s\n",
      "MPISimulator [128x128]: 0% [##################============] 100%. Total: 8s, elapsed: 5s, remaining: 3s\n",
      "MPISimulator [128x128]: 0% [##################============] 100%. Total: 8s, elapsed: 5s, remaining: 3s\n",
      "MPISimulator [128x128]: 0% [######################========] 100%. Total: 7s, elapsed: 5s, remaining: 1s\n",
      "MPISimulator [128x128]: 0% [#######################=======] 100%. Total: 6s, elapsed: 5s, remaining: 1s\n",
      "MPISimulator [128x128]: 0% [#####################=========] 100%. Total: 7s, elapsed: 5s, remaining: 1s\n",
      "Closing C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_2_0004.nc\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "\n",
    "from GPUSimulators.helpers import InitialConditions\n",
    "from GPUSimulators.Simulator import BoundaryCondition\n",
    "\n",
    "my_context.autotuner = None\n",
    "\n",
    "nx = 128\n",
    "ny = 128\n",
    "gamma = 1.4\n",
    "save_times = np.linspace(0, 5.0, 10)\n",
    "outfile = \"mpi_out_\" + str(MPI.COMM_WORLD.rank) + \".nc\"\n",
    "save_var_names = ['rho', 'rho_u', 'rho_v', 'E']\n",
    "\n",
    "\n",
    "grid = MPISimulator.MPIGrid(MPI.COMM_WORLD)\n",
    "print(grid.getLocalRank())\n",
    "\n",
    "#arguments = InitialConditions.genShockBubble(nx, ny, gamma, grid=grid)\n",
    "arguments = InitialConditions.genKelvinHelmholtz(nx, ny, gamma, grid=grid)\n",
    "#arguments = InitialConditions.genRayleighTaylor(nx, ny, gamma, grid=grid)\n",
    "\n",
    "arguments['context'] = my_context\n",
    "arguments['theta'] = 1.2\n",
    "arguments['grid'] = grid\n",
    "\n",
    "from GPUSimulators import EE2D_KP07_dimsplit\n",
    "\n",
    "def genSim(grid, **kwargs):\n",
    "    local_sim = EE2D_KP07_dimsplit.EE2D_KP07_dimsplit(**kwargs)\n",
    "    sim = MPISimulator.MPISimulator(local_sim, grid)\n",
    "    return sim\n",
    "\n",
    "\n",
    "outfile = Common.runSimulation(genSim, arguments, outfile, save_times, save_var_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] \n",
      "C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_0_0012.nc created Mon Aug 12 13:52:30 2019 contains 10 timesteps\n",
      "Simulator arguments: \n",
      " {'rho': '[...]', 'rho_u': '[...]', 'rho_v': '[...]', 'E': '[...]', 'nx': 128, 'ny': 128, 'dx': 0.005208333333333333, 'dy': 0.0078125, 'g': 0.0, 'gamma': 1.4, 'boundary_conditions': '[north=Type.Periodic, south=Type.Periodic, east=Type.Periodic, west=Type.Periodic]', 'context': 'CudaContext id 650919533984', 'theta': 1.2, 'grid': '<GPUSimulators.MPISimulator.MPIGrid object at 0x000000979166F4A8>'}\n",
      "[stdout:1] \n",
      "C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_1_0012.nc created Mon Aug 12 13:52:29 2019 contains 10 timesteps\n",
      "Simulator arguments: \n",
      " {'rho': '[...]', 'rho_u': '[...]', 'rho_v': '[...]', 'E': '[...]', 'nx': 128, 'ny': 128, 'dx': 0.005208333333333333, 'dy': 0.0078125, 'g': 0.0, 'gamma': 1.4, 'boundary_conditions': '[north=Type.Periodic, south=Type.Periodic, east=Type.Periodic, west=Type.Periodic]', 'context': 'CudaContext id 669212383056', 'theta': 1.2, 'grid': '<GPUSimulators.MPISimulator.MPIGrid object at 0x0000009BD17DF438>'}\n",
      "[stdout:2] \n",
      "C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_2_0004.nc created Mon Aug 12 13:52:29 2019 contains 10 timesteps\n",
      "Simulator arguments: \n",
      " {'rho': '[...]', 'rho_u': '[...]', 'rho_v': '[...]', 'E': '[...]', 'nx': 128, 'ny': 128, 'dx': 0.005208333333333333, 'dy': 0.0078125, 'g': 0.0, 'gamma': 1.4, 'boundary_conditions': '[north=Type.Periodic, south=Type.Periodic, east=Type.Periodic, west=Type.Periodic]', 'context': 'CudaContext id 553691644496', 'theta': 1.2, 'grid': '<GPUSimulators.MPISimulator.MPIGrid object at 0x00000080EE2ED3C8>'}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[stderr:0] \n",
      "Initialized C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_0_0012.nc\n",
      "Opening C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_0_0012.nc\n",
      "Arguments: ('r',)\n",
      "Closing C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_0_0012.nc\n",
      "[stderr:1] \n",
      "Initialized C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_1_0012.nc\n",
      "Opening C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_1_0012.nc\n",
      "Arguments: ('r',)\n",
      "Closing C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_1_0012.nc\n",
      "[stderr:2] \n",
      "Initialized C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_2_0004.nc\n",
      "Opening C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_2_0004.nc\n",
      "Arguments: ('r',)\n",
      "Closing C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_2_0004.nc\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[output:0]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 17884 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 49444 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 62520 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 80200 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 92348 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 92724 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 103992 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 115964 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 123660 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 136088 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%px\n",
    "\n",
    "with Common.DataDumper(outfile, 'r') as indata:\n",
    "    t = indata.ncfile.variables['time'][:]\n",
    "    x = indata.ncfile.variables['x'][:]\n",
    "    y = indata.ncfile.variables['y'][:]\n",
    "    \n",
    "    x_data = grid.gather(x)\n",
    "    y_data = grid.gather(y)\n",
    "    \n",
    "    created = indata.ncfile.created\n",
    "    sim_args = json.loads(indata.ncfile.sim_args)\n",
    "    for key in sim_args:\n",
    "        if isinstance(sim_args[key], list):\n",
    "            sim_args[key] = \"[...]\"\n",
    "    num_steps = len(t)\n",
    "    print(\"{:s} created {:s} contains {:d} timesteps\".format(outfile, created, num_steps))\n",
    "    print(\"Simulator arguments: \\n\", sim_args)\n",
    "    \n",
    "    \n",
    "    for i in range(num_steps):\n",
    "        rho = indata.ncfile.variables['rho'][i]\n",
    "        rho_data = grid.gather(rho)\n",
    "        \n",
    "        #Plot on rank 0\n",
    "        if (grid.comm.rank == 0):\n",
    "            plt.figure(figsize=(16, 12))\n",
    "            for k in range(grid.comm.size):\n",
    "                extent = [x_data[k].min(), x_data[k].max(), y_data[k].min(), y_data[k].max()]\n",
    "                plt.imshow(rho_data[k], extent=extent, origin='lower', vmin=1.0, vmax=2.0)\n",
    "                plotExtent(extent, 'k:', alpha=0.1)\n",
    "            plt.title(\"t=\" + str(t[i]))\n",
    "            plt.colorbar(shrink=0.33)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[stderr:0] \n",
      "Constructed in 0.020514488220214844 seconds\n",
      "Initialized C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_0_0013.nc\n",
      "Opening C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_0_0013.nc\n",
      "Keyword arguments: {'mode': 'w', 'clobber': False}\n",
      "JSON: Converting context to string (CudaContext id 650919533984)\n",
      "JSON: Converting boundary_conditions to string ([north=Type.Reflective, south=Type.Reflective, east=Type.Reflective, west=Type.Reflective])\n",
      "JSON: Converting grid to string (<GPUSimulators.MPISimulator.MPIGrid object at 0x000000978BCC20F0>)\n",
      "Closing C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_0_0013.nc\n",
      "[stderr:1] \n",
      "Constructed in 0.05653858184814453 seconds\n",
      "Initialized C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_1_0013.nc\n",
      "Opening C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_1_0013.nc\n",
      "Keyword arguments: {'mode': 'w', 'clobber': False}\n",
      "JSON: Converting context to string (CudaContext id 669212383056)\n",
      "JSON: Converting boundary_conditions to string ([north=Type.Reflective, south=Type.Reflective, east=Type.Reflective, west=Type.Reflective])\n",
      "JSON: Converting grid to string (<GPUSimulators.MPISimulator.MPIGrid object at 0x0000009BD5E20358>)\n",
      "Closing C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_1_0013.nc\n",
      "[stderr:2] \n",
      "Constructed in 0.02651834487915039 seconds\n",
      "Initialized C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_2_0005.nc\n",
      "Opening C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_2_0005.nc\n",
      "Keyword arguments: {'mode': 'w', 'clobber': False}\n",
      "JSON: Converting context to string (CudaContext id 553691644496)\n",
      "JSON: Converting boundary_conditions to string ([north=Type.Reflective, south=Type.Reflective, east=Type.Reflective, west=Type.Reflective])\n",
      "JSON: Converting grid to string (<GPUSimulators.MPISimulator.MPIGrid object at 0x00000080F39E8320>)\n",
      "Closing C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_2_0005.nc\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "\n",
    "from GPUSimulators.helpers import InitialConditions\n",
    "from GPUSimulators.Simulator import BoundaryCondition\n",
    "\n",
    "nx = 128\n",
    "ny = 128\n",
    "g = 9.81\n",
    "dt = 0.05\n",
    "width = 50\n",
    "height = 100\n",
    "save_times = np.linspace(0, 80, 10)\n",
    "outfile = \"mpi_out_\" + str(MPI.COMM_WORLD.rank) + \".nc\"\n",
    "save_var_names = ['h', 'hu', 'hv']\n",
    "\n",
    "\n",
    "if (MPI.COMM_WORLD.rank == 0):\n",
    "    h0, hu0, hv0, dx, dy = InitialConditions.bump(nx, ny, width, height, g, x_center=0.75, y_center=0.55)\n",
    "else:\n",
    "    h0, hu0, hv0, dx, dy = InitialConditions.bump(nx, ny, width, height, g, h_ref=0.5, h_amp=0.0, u_amp=0.0, v_amp=0.0)\n",
    "\n",
    "bc = BoundaryCondition({\n",
    "    'north': BoundaryCondition.Type.Reflective,\n",
    "    'south': BoundaryCondition.Type.Reflective,\n",
    "    'east': BoundaryCondition.Type.Reflective,\n",
    "    'west': BoundaryCondition.Type.Reflective\n",
    "})\n",
    "grid = MPISimulator.MPIGrid(MPI.COMM_WORLD)\n",
    "\n",
    "arguments = {\n",
    "    'context': my_context,\n",
    "    'h0': h0, 'hu0': hu0, 'hv0': hv0,\n",
    "    'nx': nx, 'ny': ny,\n",
    "    'dx': dx, 'dy': dy, \n",
    "    'g': g,\n",
    "    'boundary_conditions': bc,\n",
    "    'grid': grid\n",
    "} \n",
    "\n",
    "from GPUSimulators import HLL2\n",
    "\n",
    "def genSim(grid, **kwargs):\n",
    "    local_sim = HLL2.HLL2(**kwargs)\n",
    "    sim = MPISimulator.MPISimulator(local_sim, grid)\n",
    "    return sim\n",
    "\n",
    "outfile = Common.runSimulation(genSim, arguments, outfile, save_times, save_var_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] \n",
      "C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_0_0013.nc created Mon Aug 12 13:54:08 2019 contains 10 timesteps\n",
      "Simulator arguments: \n",
      " {'context': 'CudaContext id 650919533984', 'h0': '[...]', 'hu0': '[...]', 'hv0': '[...]', 'nx': 128, 'ny': 128, 'dx': 0.390625, 'dy': 0.78125, 'g': 9.81, 'boundary_conditions': '[north=Type.Reflective, south=Type.Reflective, east=Type.Reflective, west=Type.Reflective]', 'grid': '<GPUSimulators.MPISimulator.MPIGrid object at 0x000000978BCC20F0>'}\n",
      "[stdout:1] \n",
      "C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_1_0013.nc created Mon Aug 12 13:54:08 2019 contains 10 timesteps\n",
      "Simulator arguments: \n",
      " {'context': 'CudaContext id 669212383056', 'h0': '[...]', 'hu0': '[...]', 'hv0': '[...]', 'nx': 128, 'ny': 128, 'dx': 0.390625, 'dy': 0.78125, 'g': 9.81, 'boundary_conditions': '[north=Type.Reflective, south=Type.Reflective, east=Type.Reflective, west=Type.Reflective]', 'grid': '<GPUSimulators.MPISimulator.MPIGrid object at 0x0000009BD5E20358>'}\n",
      "[stdout:2] \n",
      "C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_2_0005.nc created Mon Aug 12 13:54:08 2019 contains 10 timesteps\n",
      "Simulator arguments: \n",
      " {'context': 'CudaContext id 553691644496', 'h0': '[...]', 'hu0': '[...]', 'hv0': '[...]', 'nx': 128, 'ny': 128, 'dx': 0.390625, 'dy': 0.78125, 'g': 9.81, 'boundary_conditions': '[north=Type.Reflective, south=Type.Reflective, east=Type.Reflective, west=Type.Reflective]', 'grid': '<GPUSimulators.MPISimulator.MPIGrid object at 0x00000080F39E8320>'}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[stderr:0] \n",
      "Initialized C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_0_0013.nc\n",
      "Opening C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_0_0013.nc\n",
      "Arguments: ('r',)\n",
      "Closing C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_0_0013.nc\n",
      "[stderr:1] \n",
      "Initialized C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_1_0013.nc\n",
      "Opening C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_1_0013.nc\n",
      "Arguments: ('r',)\n",
      "Closing C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_1_0013.nc\n",
      "[stderr:2] \n",
      "Initialized C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_2_0005.nc\n",
      "Opening C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_2_0005.nc\n",
      "Arguments: ('r',)\n",
      "Closing C:\\Users\\anbro\\Documents\\projects\\GPUSimulators\\mpi_out_2_0005.nc\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[output:0]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 25408 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 38332 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 43908 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 48924 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 52260 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 54668 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 55620 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 57924 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 55960 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 58052 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "engine": 0,
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%px\n",
    "\n",
    "with Common.DataDumper(outfile, 'r') as indata:\n",
    "    t = indata.ncfile.variables['time'][:]\n",
    "    x = indata.ncfile.variables['x'][:]\n",
    "    y = indata.ncfile.variables['y'][:]\n",
    "    \n",
    "    x_data = grid.gather(x)\n",
    "    y_data = grid.gather(y)\n",
    "    \n",
    "    created = indata.ncfile.created\n",
    "    sim_args = json.loads(indata.ncfile.sim_args)\n",
    "    for key in sim_args:\n",
    "        if isinstance(sim_args[key], list):\n",
    "            sim_args[key] = \"[...]\"\n",
    "    num_steps = len(t)\n",
    "    print(\"{:s} created {:s} contains {:d} timesteps\".format(outfile, created, num_steps))\n",
    "    print(\"Simulator arguments: \\n\", sim_args)\n",
    "    \n",
    "    \n",
    "    for i in range(num_steps):\n",
    "        h = indata.ncfile.variables['h'][i]\n",
    "        #hu = indata.ncfile.variables['hu'][i]\n",
    "        #hv = indata.ncfile.variables['hv'][i]\n",
    "        h_data = grid.gather(h)\n",
    "        \n",
    "        #Plot on rank 0\n",
    "        if (grid.comm.rank == 0):\n",
    "            plt.figure(figsize=(16, 12))\n",
    "            for k in range(grid.comm.size):\n",
    "                extent = [x_data[k].min(), x_data[k].max(), y_data[k].min(), y_data[k].max()]\n",
    "                plt.imshow(h_data[k], extent=extent, origin='lower', vmin=0.49, vmax=0.52)\n",
    "                plotExtent(extent, 'k:', alpha=0.1)\n",
    "            plt.title(\"t=\" + str(t[i]))\n",
    "            plt.colorbar(shrink=0.33)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:root] *",
   "language": "python",
   "name": "conda-root-py"
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
