Parallel Computing with Python public.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Parallel Computing with Python\n",
    "================================\n",
    "\n",
    "**[Rodrigo Nemmen](http://rodrigonemmen.com), IAG USP**\n",
    "\n",
    "This IPython notebook illustrates a few simple ways of doing parallel computing.\n",
    "\n",
    "Practical examples included:\n",
    "\n",
    "1. Parallel function mapping to a list of arguments (multiprocessing module)\n",
    "2. Parallel execution of array function (scatter/gather) + parallel execution of scripts\n",
    "4. Easy parallel Monte Carlo (parallel magics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Mapping a model to a grid of parameters \n",
    "\n",
    "<!---\n",
    "Inspired on \"useful parallel\".\n",
    "-->\n",
    "\n",
    "Uses the *multiprocess* module. *Note: We use `multiprocess` instead of `multiprocessing` because even though the latter comes by default with python it is known to be problematic with Jupyter.*\n",
    "\n",
    "Idea: you have a function $f(\\mathbf{x},\\mathbf{y})$ of two parameters (e.g., $f$ may represent your model) stored in the arrays $(\\mathbf{x},\\mathbf{y})$. Given the arrays $\\mathbf{x}$ and $\\mathbf{y}$, you want to compute the values of $f(\\mathbf{x},\\mathbf{y})$. Let's assume for simplicity that there is no dependence on the neighbours. This is an embarassingly parallel problem.\n",
    "\n",
    "<!---\n",
    "### TODO\n",
    "\n",
    "* Random sampling of parameter space if desired\n",
    "-->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiprocess as mp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Time wasting function that depends on two parameters. Here, I generate 1E5 random numbers based on the normal distribution and then sum them. The two parameters are $\\mu$ and $\\sigma$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(z):\n",
    "    x=np.random.normal(z[0], z[1], 100000)\n",
    "    return x.sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Arrays of input parameters. You could easily modify this to take as input a matrix, not two arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "n=3000\n",
    "X=np.linspace(-1,1,n) # mean\n",
    "Y=np.linspace(0.1,1,n) # std. dev."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# creates list of arguments [Xi, Yi]\n",
    "pargs=[]\t# this is a list of lists!\n",
    "for i in range(X.size):\n",
    "\tpargs.append([X[i],Y[i]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Parallel execution. Check out all the cores being used with a tool like [htop](http://hisham.hm/htop/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "ncores=mp.cpu_count() # number of cores\n",
    "pool = mp.Pool(processes=ncores) # initializes parallel engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.45 s, sys: 28.6 ms, total: 2.48 s\n",
      "Wall time: 29.5 s\n"
     ]
    }
   ],
   "source": [
    "%%time \n",
    "t=pool.map(f, pargs)\t# parallel function map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "pool.close()\t# close the parallel engine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Serial execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 636 µs, sys: 3 µs, total: 639 µs\n",
      "Wall time: 645 µs\n"
     ]
    }
   ],
   "source": [
    "%time t=map(f, pargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to convert the list to an array use `y=np.array(t)`. Notes: (1) there is a `map` method for `ipyparallel`. (2) the example I picked is not the best one because numpy is highly optimized hence faster than doing in parallel where communication between processes will kill the performance in this case. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Parallel execution of array function\n",
    "\n",
    "Uses [`ipyparallel`](https://ipyparallel.readthedocs.io/en/latest/index.html). Consider a function $f(x)$ which takes an array $x$ containing the grid of input parameters. We want to split the function calls (\"split the array\") to the different cores in our machine:\n",
    "\n",
    "\n",
    "![test](array_proc1.gif)\n",
    "![test](array_proc2.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start parallel engine\n",
    "\n",
    "Make sure you have ipyparallel installed: \n",
    "\n",
    "    conda install ipyparallel\n",
    "\n",
    "Now, start the parallel engines. Open a terminal and use the command:\n",
    "\n",
    "    ipcluster start"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Function to be parallelized"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our time-waster function $f(x)$ that can be applied to an array of integers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test if n is prime\n",
    "def isprime(n):\n",
    "    for i in range(3, n):\n",
    "        if n % i == 0:\n",
    "            return False\n",
    "    return True\n",
    "\n",
    "# tests each element of an array if it is prime\n",
    "def f(x):\n",
    "    return list(map(isprime,x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generates big array (10k elements) of random integers between 0 and 100000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "x=np.random.randint(low=0, high=100000, size=10000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tests"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Serial execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 8.17 s, sys: 32.4 ms, total: 8.2 s\n",
      "Wall time: 9.74 s\n"
     ]
    }
   ],
   "source": [
    "%time y=f(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now explain how ipyparallel works **(here I show a slide)**. See documentation at the end of the notebook for details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting 8 engines with <class 'ipyparallel.cluster.launcher.LocalEngineSetLauncher'>\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Cluster(cluster_id='1683573801-aucc', profile='default', controller=<running>, engine_sets=['1683573802'])>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ipyparallel as ipp\n",
    "\n",
    "cluster = ipp.Cluster(n=ncores)\n",
    "cluster.start_cluster_sync()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are going to use the direct view, which means that commands always run on all nodes. This as opposed to a balanced view, which asynchronously executes code on nodes which are idle. In addition, we are going to turn blocking on. This means that jobs will block further execution until all nodes have finished."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "rc = cluster.connect_client_sync()\n",
    "rc.wait_for_engines()\n",
    "\n",
    "direct = rc[:] # use all engines\n",
    "direct.block=True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Splits the input array $x$ between the cores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "direct.scatter('x',x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Verify that the array was indeed divided equally"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1250, 1250, 1250, 1250, 1250, 1250, 1250, 1250]"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "direct['x.size']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([94305, 72839, 17104, ..., 29346, 73755, 29269]),\n",
       " array([31625, 37515, 37053, ..., 76381, 32938, 13199]),\n",
       " array([44846, 30440, 38205, ..., 83728,  5019, 84130]),\n",
       " array([29578, 88280, 80813, ..., 32620, 52857, 27595])]"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "direct['x']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try to apply the function in each different core"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2:execute]\n",
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)\n",
      "Cell \u001b[0;32mIn[1], line 1\u001b[0m\n",
      "\u001b[0;32m----> 1\u001b[0m y\u001b[38;5;241m=\u001b[39m\u001b[43mf\u001b[49m(x)\n",
      "\n",
      "\u001b[0;31mNameError\u001b[0m: name 'f' is not defined\n",
      "[3:execute]\n",
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)\n",
      "Cell \u001b[0;32mIn[1], line 1\u001b[0m\n",
      "\u001b[0;32m----> 1\u001b[0m y\u001b[38;5;241m=\u001b[39m\u001b[43mf\u001b[49m(x)\n",
      "\n",
      "\u001b[0;31mNameError\u001b[0m: name 'f' is not defined\n",
      "[7:execute]\n",
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)\n",
      "Cell \u001b[0;32mIn[1], line 1\u001b[0m\n",
      "\u001b[0;32m----> 1\u001b[0m y\u001b[38;5;241m=\u001b[39m\u001b[43mf\u001b[49m(x)\n",
      "\n",
      "\u001b[0;31mNameError\u001b[0m: name 'f' is not defined\n",
      "[0:execute]\n",
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)\n",
      "Cell \u001b[0;32mIn[1], line 1\u001b[0m\n",
      "\u001b[0;32m----> 1\u001b[0m y\u001b[38;5;241m=\u001b[39m\u001b[43mf\u001b[49m(x)\n",
      "\n",
      "\u001b[0;31mNameError\u001b[0m: name 'f' is not defined\n",
      "[6:execute] NameError: name 'f' is not defined\n",
      "[4:execute] NameError: name 'f' is not defined\n",
      "[1:execute] NameError: name 'f' is not defined\n",
      "[5:execute] NameError: name 'f' is not defined\n"
     ]
    },
    {
     "ename": "AlreadyDisplayedError",
     "evalue": "8 errors",
     "output_type": "error",
     "traceback": [
      "8 errors"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "y=f(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Why the errors above? Because each core does not see the local engine. They work as separate machines and you have to load all variables and modules separately in each engine. Fortunately, that's easy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting myscript.py\n"
     ]
    }
   ],
   "source": [
    "%%file myscript.py\n",
    "# test if n is prime\n",
    "def isprime(n):\n",
    "    for i in range(3, n):\n",
    "        if n % i == 0:\n",
    "            return False\n",
    "    return True\n",
    "\n",
    "# tests each element of an array if it is prime\n",
    "def f(x):\n",
    "    return list(map(isprime,x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Execute code which defines the methods on the different engines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<AsyncResult(execute): finished>"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "direct.run(\"myscript.py\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now compute the \"model grid\" correctly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c2be13917ecd44949901a3e7586b3b6d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "%px:   0%|          | 0/8 [00:00<?, ?tasks/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 107 ms, sys: 47.7 ms, total: 155 ms\n",
      "Wall time: 2.62 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "%px y=f(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively to the command above, you could use\n",
    "\n",
    "    direct.apply(f,x)\n",
    "or\n",
    "\n",
    "    direct.execute('y=f(x)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have the separate arrays $y$ containing the results on each engine. How to get it back to the local engine?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[0:5]: \u001b[0m1250"
      ]
     },
     "metadata": {
      "after": null,
      "completed": null,
      "data": {},
      "engine_id": 0,
      "engine_uuid": "03851eed-c2ebd02e989a5c72d465dadc",
      "error": null,
      "execute_input": "import numpy\nnumpy.size(y)\n",
      "execute_result": {
       "data": {
        "text/plain": "1250"
       },
       "execution_count": 5,
       "metadata": {}
      },
      "follow": null,
      "msg_id": null,
      "outputs": [],
      "received": null,
      "started": null,
      "status": null,
      "stderr": "",
      "stdout": "",
      "submitted": "2023-05-08T19:27:20.035048Z"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[1:5]: \u001b[0m1250"
      ]
     },
     "metadata": {
      "after": null,
      "completed": null,
      "data": {},
      "engine_id": 1,
      "engine_uuid": "c7a2e15d-21cf60e6f29db63558bfe3e2",
      "error": null,
      "execute_input": "import numpy\nnumpy.size(y)\n",
      "execute_result": {
       "data": {
        "text/plain": "1250"
       },
       "execution_count": 5,
       "metadata": {}
      },
      "follow": null,
      "msg_id": null,
      "outputs": [],
      "received": null,
      "started": null,
      "status": null,
      "stderr": "",
      "stdout": "",
      "submitted": "2023-05-08T19:27:20.036435Z"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[3:5]: \u001b[0m1250"
      ]
     },
     "metadata": {
      "after": [],
      "completed": "2023-05-08T19:27:20.055305Z",
      "data": {},
      "engine_id": 3,
      "engine_uuid": "29382c1f-b7d5519040f13c1070f4d9bd",
      "error": null,
      "execute_input": "import numpy\nnumpy.size(y)\n",
      "execute_result": {
       "data": {
        "text/plain": "1250"
       },
       "execution_count": 5,
       "metadata": {}
      },
      "follow": [],
      "is_broadcast": false,
      "is_coalescing": false,
      "msg_id": "b308550d-e694c8100b28fc416d1f91ca_76213_36",
      "outputs": [],
      "received": "2023-05-08T19:27:20.080771Z",
      "started": "2023-05-08T19:27:20.046216Z",
      "status": "ok",
      "stderr": "",
      "stdout": "",
      "submitted": "2023-05-08T19:27:20.037440Z"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[2:5]: \u001b[0m1250"
      ]
     },
     "metadata": {
      "after": [],
      "completed": "2023-05-08T19:27:20.055140Z",
      "data": {},
      "engine_id": 2,
      "engine_uuid": "d4e940ef-ae032e638ca505de07944254",
      "error": null,
      "execute_input": "import numpy\nnumpy.size(y)\n",
      "execute_result": {
       "data": {
        "text/plain": "1250"
       },
       "execution_count": 5,
       "metadata": {}
      },
      "follow": [],
      "is_broadcast": false,
      "is_coalescing": false,
      "msg_id": "b308550d-e694c8100b28fc416d1f91ca_76213_35",
      "outputs": [],
      "received": "2023-05-08T19:27:20.073630Z",
      "started": "2023-05-08T19:27:20.045609Z",
      "status": "ok",
      "stderr": "",
      "stdout": "",
      "submitted": "2023-05-08T19:27:20.037191Z"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[4:5]: \u001b[0m1250"
      ]
     },
     "metadata": {
      "after": [],
      "completed": "2023-05-08T19:27:20.055832Z",
      "data": {},
      "engine_id": 4,
      "engine_uuid": "93a29946-7d7d65a6544917d87ef2b75c",
      "error": null,
      "execute_input": "import numpy\nnumpy.size(y)\n",
      "execute_result": {
       "data": {
        "text/plain": "1250"
       },
       "execution_count": 5,
       "metadata": {}
      },
      "follow": [],
      "is_broadcast": false,
      "is_coalescing": false,
      "msg_id": "b308550d-e694c8100b28fc416d1f91ca_76213_37",
      "outputs": [],
      "received": "2023-05-08T19:27:20.084359Z",
      "started": "2023-05-08T19:27:20.046604Z",
      "status": "ok",
      "stderr": "",
      "stdout": "",
      "submitted": "2023-05-08T19:27:20.037629Z"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[7:5]: \u001b[0m1250"
      ]
     },
     "metadata": {
      "after": [],
      "completed": "2023-05-08T19:27:20.058890Z",
      "data": {},
      "engine_id": 7,
      "engine_uuid": "92b5c49b-6c9402ca407e5ccd51c1f9aa",
      "error": null,
      "execute_input": "import numpy\nnumpy.size(y)\n",
      "execute_result": {
       "data": {
        "text/plain": "1250"
       },
       "execution_count": 5,
       "metadata": {}
      },
      "follow": [],
      "is_broadcast": false,
      "is_coalescing": false,
      "msg_id": "b308550d-e694c8100b28fc416d1f91ca_76213_40",
      "outputs": [],
      "received": "2023-05-08T19:27:20.088441Z",
      "started": "2023-05-08T19:27:20.050157Z",
      "status": "ok",
      "stderr": "",
      "stdout": "",
      "submitted": "2023-05-08T19:27:20.038129Z"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[5:5]: \u001b[0m1250"
      ]
     },
     "metadata": {
      "after": [],
      "completed": "2023-05-08T19:27:20.057607Z",
      "data": {},
      "engine_id": 5,
      "engine_uuid": "a022b31c-bc75a3ef0e062bc5f1e8834c",
      "error": null,
      "execute_input": "import numpy\nnumpy.size(y)\n",
      "execute_result": {
       "data": {
        "text/plain": "1250"
       },
       "execution_count": 5,
       "metadata": {}
      },
      "follow": [],
      "is_broadcast": false,
      "is_coalescing": false,
      "msg_id": "b308550d-e694c8100b28fc416d1f91ca_76213_38",
      "outputs": [],
      "received": "2023-05-08T19:27:20.087164Z",
      "started": "2023-05-08T19:27:20.048313Z",
      "status": "ok",
      "stderr": "",
      "stdout": "",
      "submitted": "2023-05-08T19:27:20.037798Z"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[6:5]: \u001b[0m1250"
      ]
     },
     "metadata": {
      "after": [],
      "completed": "2023-05-08T19:27:20.057075Z",
      "data": {},
      "engine_id": 6,
      "engine_uuid": "4c47efda-d372d2bd56a569708b4c039c",
      "error": null,
      "execute_input": "import numpy\nnumpy.size(y)\n",
      "execute_result": {
       "data": {
        "text/plain": "1250"
       },
       "execution_count": 5,
       "metadata": {}
      },
      "follow": [],
      "is_broadcast": false,
      "is_coalescing": false,
      "msg_id": "b308550d-e694c8100b28fc416d1f91ca_76213_39",
      "outputs": [],
      "received": "2023-05-08T19:27:20.086061Z",
      "started": "2023-05-08T19:27:20.048842Z",
      "status": "ok",
      "stderr": "",
      "stdout": "",
      "submitted": "2023-05-08T19:27:20.037965Z"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%px \n",
    "import numpy\n",
    "numpy.size(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "y=direct.gather('y')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have the array magically reassembled back in the local engine.  :)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Easy parallel Monte Carlo\n",
    "\n",
    "Suppose you need to do 100k Monte Carlo simulations. Wouldn't it be great if you could easily split them among your (hopefully many) cores?\n",
    "\n",
    "In this example, I will perform 100k realizations of a 300x300 array of random floats."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "import astropy.io.ascii as ascii"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "data=ascii.read('/Users/nemmen/Dropbox/science/projects/finished/jetpower/data/allrb.dat')\n",
    "\n",
    "xdata=data['pb']\n",
    "ydata=data['pj']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "# number of desired bootstraps\n",
    "nboot=200000\n",
    "\n",
    "# number of bootstraps that will go to each core\n",
    "n=int(nboot/np.size(rc.ids))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Passes variables to the engines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "direct.push(dict(n=n,xdata=xdata,ydata=ydata))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "import numpy as np \n",
    "import nmmn.lsd, scipy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now everything below is executed in parallel! (IPython magic)\n",
    "\n",
    "<!---\n",
    "Have a look also at the %autopx command.\n",
    "-->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "92c70762ce0642a7b20f30c9485afcc5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "%px:   0%|          | 0/8 [00:00<?, ?tasks/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 118 ms, sys: 30.4 ms, total: 148 ms\n",
      "Wall time: 8.21 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "%%px\n",
    "r=np.zeros(n)\n",
    "\n",
    "for i in range(n):\n",
    "    [xsim,ysim]=nmmn.lsd.bootstrap([xdata,ydata]) # data bootstrapping\n",
    "    r[i]=scipy.stats.pearsonr(xsim,ysim)[0]\t # just compute simple Pearson coefficient"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assemble result from cores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "r=direct.gather('r')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For comparison, how long does it take to do the same simulation in serial mode?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 44.8 s, sys: 2.42 s, total: 47.2 s\n",
      "Wall time: 56.6 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "for i in range(nboot):\n",
    "    [xsim,ysim]=nmmn.lsd.bootstrap([xdata,ydata])\n",
    "    t=scipy.stats.pearsonr(xsim,ysim)[0]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}