course2/03_multiprocessing_multithreading_python.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Notebook 3 - working with processes / threads \n",
    "-----------------------------------------------------------------------\n",
    "\n",
    "## Table of Content <a id=\"toc\"></a>\n",
    "\n",
    "1. **[Multiprocessing (and refactoring)](#8)**  \n",
    "   <br>\n",
    "   \n",
    "2. **[Numba and parallelization](#9)**  \n",
    "    2.1. [automatic parallelization](#2.1)  \n",
    "    2.2. [explicit parallelization (prange)](#2.2)  \n",
    "    2.3. [controlling the number of threads used](#2.3)  \n",
    "    <br>\n",
    "    \n",
    "**Supplementary material:**  \n",
    "   * Annex 1 - [Parallelization of pairwise distance computation with multiprocess](#annexa)  \n",
    "   * Annex 2 - [Parallelization of pairwise distance computation with numba](#annexb)  \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 1 Multiprocessing (and refactoring) <a id='8'></a>\n",
    "\n",
    "We can take advantage of the multiple cores available on our computers by using the **`multiprocessing`** module. \n",
    "\n",
    "In this approach, separate __processes__ are used, __not threads__. \n",
    "\n",
    "The use of threads is generally blocked by Python because of the \"*Global Interpreter Lock*\". This was a necessary design feature as a trade-off for the enormous flexibility in memory management that Python makes possible. This means that there is no shared memory when using multiprocessing, and thus the individual tasks must be independent.\n",
    "\n",
    "`multiprocessing` generally works well with lists, where one maps a function to each element of the list and these operations are computed as separated processes, on separate cores per element of the list. \n",
    "\n",
    "Indeed, any kind of parallelization technique is really only worth it if the task you want to do is actually *parallelizable*. It is sometimes hard to judge what is and is not easily parallelizable, and can often require that you refactor your code quite a bit.\n",
    "\n",
    "A rule of thumb for whether parallelization is possible is to evaluate whether the task can be divided into subtasks which: \n",
    "1. **Do not depend on each others results.**\n",
    "2. Are very similar.\n",
    "3. Use independent parts of the data.\n",
    "\n",
    "Point 1. is the most important, the others are helpful but not entirely necessary.\n",
    "\n",
    "<br>\n",
    "\n",
    "**Example:** consider our function to compute an integral from the previous lesson:\n",
    "\n",
    "* Here is the code in **pure python**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_native(x):\n",
    "    return x ** 2 - x\n",
    "\n",
    "def integrate_f_native(a, b, N):\n",
    "    s = 0\n",
    "    dx = (b - a) / N\n",
    "    for i in range(N):\n",
    "        s += f_native(a + i * dx)\n",
    "    return s * dx\n",
    "\n",
    "print(integrate_f_native(0, 2, 100))\n",
    "%timeit -n 10 -r 3 integrate_f_native(0, 2, 1000000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ideally, we would like to reduce this something that looks like:\n",
    "\n",
    "```python\n",
    "for i in range(len(data)):\n",
    "    result[i] = function(data[i])\n",
    "\n",
    "```\n",
    "\n",
    "Equivalent to:\n",
    "```python\n",
    "map(function, data)\n",
    "```\n",
    "\n",
    "So, we apply a `function` to each element (`data[i]`) of `data`.  \n",
    "The game is thus to re-write it slightly so it fits this template.\n",
    "\n",
    "Let's work outside the function and focus on the main loop:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = 0\n",
    "b = 2\n",
    "N = 1000000\n",
    "dx = (b - a) / N\n",
    "\n",
    "data = list(range(N))\n",
    "\n",
    "def f2(i):\n",
    "    x = a + i*dx\n",
    "    return x ** 2 - x\n",
    "\n",
    "result = map(f2, data)\n",
    "\n",
    "# Equivalent to\n",
    "# for i in range(N):\n",
    "#     result[i] = f_native(data[i])\n",
    "\n",
    "final_result = sum(result) * dx\n",
    "print(final_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, everything is ready for us to use `multiprocessing`.  \n",
    "The simplest usage is to open up a pool of processes using the `with` keyword:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiprocessing as mp\n",
    "\n",
    "with mp.Pool(2) as pool:\n",
    "    result = pool.map(f2, data)\n",
    "    \n",
    "final_result2 = sum(result) * dx\n",
    "print(final_result2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ok so we get the same result when splitting the task on 2 processes, but does it perform faster?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit -n 10 -r 3 list(map(f2, data))\n",
    "\n",
    "with mp.Pool(2) as pool:\n",
    "    %timeit -n 10 -r 3 pool.map(f2, data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Mhmm, what? the multiprocessing version is actually slower...\n",
    "\n",
    "This is due to the fact that opening, closing, and communicating data to and from processes are costly operation.\n",
    "In other words, **multiprocessing a task has some overhead**, and it therefore tends to work better with a few long tasks than with a lot of very small ones (note: each parallelization techniques have different overhead and react differently to this).\n",
    "\n",
    "For instance, let's try with a few long tasks:\n",
    "* Our \"long\" task will be to run the `integrate_f_native()` function between 0 and 1, with 400 thousand points.  \n",
    "  This takes around 0.06 seconds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit -n 10 -r 3 integrate_f_native(0, 1, 4 * 10**5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Around 0.06 seconds per task\n",
    "def task(i):\n",
    "    return integrate_f_native(0, i, 4 * 10**5)\n",
    "\n",
    "# 100 tasks to perform.\n",
    "data = list(range(1, 101))\n",
    "\n",
    "# Serial execution: ~6 seconds.\n",
    "%time _ = list(map(task, data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with mp.Pool(2) as pool:\n",
    "    %time  pool.map(task, data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With such long tasks, the overhead is lower than the gained time.  \n",
    "Indeed, on the basis of 0.06 seconds per task, we would expect 100 tasks on 2 processes to take ~ 3 seconds, so we have ~0.45 seconds of overhead here.\n",
    "\n",
    "<br>\n",
    "\n",
    "Let's now **vary the number of processes** to see whether adding more processors (CPUs) allows to further speed-up the computing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for number_of_processors in [1, 2, 4, 8]:\n",
    "    print(number_of_processors)\n",
    "    with mp.Pool(number_of_processors) as pool:\n",
    "        %time pool.map(task, data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Note: as we increase the number of processes, the overhead increases and after some value this\n",
    "> actually hurts the overall performance.\n",
    ">\n",
    "> * **Multiprocessing works better when the individual tasks are longer.**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "<div class=\"alert alert-block alert-success\">\n",
    "\n",
    "## Exercise 3.1\n",
    "\n",
    "See the dedicated `exercises_course2.ipynb` notebook.\n",
    "\n",
    "<div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "<br>\n",
    "<br>\n",
    "\n",
    "[Back to ToC](#toc)\n",
    "\n",
    "## 2. Numba and parallelization  <a id='9'></a>\n",
    "\n",
    "It is possible to provide a `numba` function to `mp.pool`, but `numba` already provides what's necessary to parallelize code.\n",
    "* By setting **`parallel=True`** when calling `@jit` (in no-python mode), numba will attempt\n",
    "  to automatically parallelize your code.\n",
    "* In particular, by default, it works on the array operations.\n",
    "\n",
    "<br>\n",
    "\n",
    "### 2.1 Automatic parallelization <a id='2.1'></a>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from numba import njit  # njit is a shortcut for jit(nopython=True).\n",
    "\n",
    "def integrate_f(a, b, N):\n",
    "    dx = (b - a) / N\n",
    "    X = np.arange(a,b,dx)\n",
    "    return (X**2 - X).sum() * dx\n",
    "\n",
    "# The code does not change, so no need to re-write it; just give the function to njit.\n",
    "integrate_f_numba = njit(integrate_f)\n",
    "integrate_f_numba_parallel = njit(integrate_f , parallel=True)\n",
    "\n",
    "# Check that we get similar results (additionally, this lets numba do the compilation now).\n",
    "print(\"native         :\", integrate_f(0,2,100))\n",
    "print(\"numba          :\", integrate_f_numba(0,2,100))\n",
    "print(\"numba parallel :\", integrate_f_numba_parallel(0,2,100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "Let's time the different versions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 10**7\n",
    "print(\"native         :\")\n",
    "%timeit integrate_f(0, 2, N)\n",
    "\n",
    "print(\"numba          :\")\n",
    "%timeit integrate_f_numba(0, 2, N)\n",
    "\n",
    "print(\"numba parallel :\")\n",
    "%timeit integrate_f_numba_parallel(0, 2, N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So the basic numba seems less efficient than numpy, but the parallel version is showing quite a speedup!  \n",
    "Here, numba was able to parallelize the `np.arange`, all the array operations, and the `sum()`, so actually almost all the code.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "[Back to ToC](#toc)\n",
    "\n",
    "### 2.2 Explicit parallelization (prange) <a id='2.2'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def integrate_f2(a, b, N):\n",
    "    dx = (b - a) / N\n",
    "    s = 0\n",
    "    for i in range(N):\n",
    "        x = a + i*dx\n",
    "        s += x**2 - x\n",
    "    return s * dx\n",
    "\n",
    "integrate_f2_numba = njit(integrate_f2)\n",
    "integrate_f2_numba_parallel = njit(integrate_f2,parallel=True)\n",
    "\n",
    "\n",
    "# Check that we get similar results (additionally this let's numba do the compilation now).\n",
    "print(\"native         :\", integrate_f2(0,2,100))\n",
    "print(\"numba          :\", integrate_f2_numba(0,2,100))\n",
    "print(\"numba parallel :\", integrate_f2_numba_parallel(0,2,100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import prange \n",
    "\n",
    "@njit(parallel=True)\n",
    "def integrate_f2_numba_parallel(a, b, N):\n",
    "    dx = (b - a) / N\n",
    "    s = 0\n",
    "    \n",
    "    # prange() tells numba which is the loop to parallelize. \n",
    "    for i in prange(N):\n",
    "        x = a + i*dx\n",
    "        s += x**2 - x\n",
    "    return s * dx\n",
    "\n",
    "integrate_f2_numba_parallel(0, 2, 100) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "Let's time it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 10**7\n",
    "print(\"numba from native         : \", end=\"\")\n",
    "%timeit -n 10 -r 3 integrate_f2_numba(0, 2, N)\n",
    "\n",
    "print(\"numba parallel with prange: \", end=\"\")\n",
    "%timeit -n 10 -r 3 integrate_f2_numba_parallel(0, 2, N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's use a bit more data to compare the 2 parallel versions (auto and manual):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 10**8\n",
    "print(\"numba parallel auto       : \", end=\"\")\n",
    "%timeit -n 10 -r 3 integrate_f_numba_parallel(0, 2, N)\n",
    "\n",
    "print(\"numba parallel with prange: \", end=\"\")\n",
    "%timeit -n 10 -r 3 integrate_f2_numba_parallel(0, 2, N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The two results are similar. The one you end up using will depend on the structure of your problem and the shape of your code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "[Back to ToC](#toc)\n",
    "\n",
    "### 2.3 Controlling the number of threads used <a id=\"2.3\"></a>\n",
    "\n",
    "Up until now, we have let numba use its default number of threads."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numba\n",
    "numba.config.NUMBA_DEFAULT_NUM_THREADS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To control the number of threads, just use the `set_num_threads` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import set_num_threads\n",
    "\n",
    "N = 10**8\n",
    "\n",
    "# Max number of threads.\n",
    "print(\"Default number of threads: \", end=\"\")\n",
    "%timeit -n 10 -r 3 integrate_f2_numba_parallel(0, 2, N)\n",
    "\n",
    "for num_thread in range(2, 8):\n",
    "    print(\"Thread number set to -> \", num_thread, \": \", sep=\"\", end=\"\")\n",
    "    set_num_threads(num_thread)\n",
    "    %timeit -n 10 -r 3 integrate_f2_numba_parallel(0, 2, N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To go further, we recommend you have a look at [numba documentation on parallelization](https://numba.pydata.org/numba-doc/latest/user/parallel.html) which explains what can, and what cannot be parallelized, and how to diagnose the automatic parallelization process."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "<br>\n",
    "<br>\n",
    "<br>\n",
    "\n",
    "[Back to ToC](#toc)\n",
    "\n",
    "# Additional material\n",
    "------------------------------\n",
    "\n",
    "## Annex 1 - parallelization of pairwise distance computation with multiprocess <a id='annexa'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "\n",
    "def pairwise_distance_numpy(X):\n",
    "\n",
    "    num_vectors = X.shape[0]\n",
    "    num_measurements = X.shape[1] \n",
    "    D = np.empty((num_vectors, num_vectors), dtype=np.float64)\n",
    "    \n",
    "    for i in range(num_vectors):\n",
    "        for j in range(num_vectors):\n",
    "            d = np.square(np.subtract(X[i], X[j]))\n",
    "            D[i, j] = np.sqrt(np.sum(d))\n",
    "            \n",
    "    return(D)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Right now, this function operates onto a whole array. but ideally, we would like to reduce this something that looks like:\n",
    "\n",
    "```python\n",
    "for i in range(len(data)):\n",
    "    result[i] = function(data[i])\n",
    "```\n",
    "\n",
    "Equivalent to:\n",
    "\n",
    "```python\n",
    "map(function,data)\n",
    "```\n",
    "\n",
    "So, we apply a `function` to each element (`data[i]`) of `data`.\n",
    "\n",
    "**Question:** how can we go from the `pairwise_distance_numpy` function to this? what would be `function`? `data`? \n",
    "\n",
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ... don't scroll - spoilers ahead ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "So, my proposition to solve this (not the only one possible, maybe not even the best) is that :\n",
    " 1. the `function` is computing distance between 2 vectors\n",
    " 2. the `data[i]` is a couple of vector\n",
    " 3. consequently, `data` is a list of couples of vectors.\n",
    "\n",
    "\n",
    "I will even go one (small) step further, and rather than keeping the whole vectors in data, I will just keep the vector indexes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate 200 vectors with 100 measurements each \n",
    "data = np.random.uniform(size=(200, 100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pairwise_list_I(X):\n",
    "    \"\"\"Create a list of the pairs of vector index we have to\n",
    "    compute distances for (ie. all possible pair of indexes).\n",
    "    \"\"\"\n",
    "    list_of_tuples = list()\n",
    "    \n",
    "    num_vectors = X.shape[0]\n",
    "    num_measurements = X.shape[1] \n",
    "    \n",
    "    for i in range(num_vectors):\n",
    "        for j in range(num_vectors):\n",
    "            list_of_tuples.append((i,j))\n",
    "            \n",
    "    return list_of_tuples\n",
    "\n",
    "def pairwise_distance_from_indexes(indexes ):\n",
    "    \"\"\"Takes a tuple containing a pair of indexes,\n",
    "    and computes the distance between the 2.\n",
    "    \"\"\"\n",
    "    assert(len(indexes) == 2)\n",
    "    X1 = data[indexes[0]]\n",
    "    X2 = data[indexes[1]]\n",
    "    \n",
    "    return np.sqrt(np.sum(np.square(X1 - X2)))\n",
    "\n",
    "\n",
    "list_of_tuples_I = pairwise_list_I(data)\n",
    "\n",
    "%timeit -n 1 -r 3  result = list(map(pairwise_distance_from_indexes,  list_of_tuples_I))\n",
    "\n",
    "import multiprocessing as mp\n",
    "\n",
    "with mp.Pool(2) as pool :\n",
    "    \n",
    "    %timeit -n 1 -r 3  result2 = pool.map(pairwise_distance_from_indexes, list_of_tuples_I)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some speedup, but nothing tremendous.\n",
    "\n",
    "Let's see if that holds up :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for NP in range(1,7):\n",
    "    print(NP)\n",
    "    with mp.Pool(NP) as pool :\n",
    "        %time result = pool.map(pairwise_distance_from_indexes, list_of_tuples_I)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Of course, we want to compare this with the original version of the function.\n",
    "%timeit pairwise_distance_numpy(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So there is some gain. \n",
    "Nothing tremendous, but sill a 1.5x speedup, and it beats the alternative of having all these core idle.\n",
    "\n",
    "<br>\n",
    "\n",
    "Going one step further, we know multiprocessing works better when the task are somewhat large. So, instead of say that a task is \"compute a single distance\", and having NxN tasks, we could have the task be \"compute a full row of the distance matrix\", and then we only have N tasks.\n",
    "\n",
    "So for the task, we compute the distance between one vector and all the others. You will see this is a very good idea, even in a non-multiprocessing framework, because this plays into some of numpy's strength.\n",
    "\n",
    "First, let's implement our \"task\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Computing reference results for testing validity.\n",
    "toy_data = np.random.uniform(size=(10, 100))\n",
    "\n",
    "res = pairwise_distance_numpy(toy_data)\n",
    "# We want our task to compute something like this:\n",
    "res[0]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "We have seen that numpy makes operation between 2 vector easy. but actually, operation between a matrix and a vector works as well so, matrix - vector will perform the subtraction on each row independently.  \n",
    "Then, if we make the sum also on each row independently we can get the distances we want!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_distance_row( i ):\n",
    "    # Here, we assume that there exists a DATA_GLOB\n",
    "    # variable in global memory with my data in it.\n",
    "    \n",
    "    squared_diff = (DATA_GLOB - DATA_GLOB[i])**2  # Squared differences between the matrix and a single vector.\n",
    "    sums = np.sum( squared_diff , axis=1)         # axis=1 --> to get 1 sum per row.\n",
    "    return np.sqrt( sums )                        # Compute square root of all these sums.\n",
    "\n",
    "# Of course we want to test this:\n",
    "DATA_GLOB = toy_data\n",
    "res_new = compute_distance_row(0)\n",
    "print(res_new)\n",
    "print(res_new == res[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All good so far. How does it perform?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Redefine the data used by the function as the bigger dataset.\n",
    "DATA_GLOB = data\n",
    "\n",
    "# Map this onto the list of possible indices : from 0 to N\n",
    "%timeit -n 1 -r 3  result = list(map(compute_distance_row, range(DATA_GLOB.shape[0])))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So you can see how this actually performs even better even with a single process.\n",
    "\n",
    "Actually, let's make the data larger."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "big_data = np.random.uniform(size=(500,100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_GLOB = big_data\n",
    "%timeit -n 5 -r 3  result = list(map(compute_distance_row, range(DATA_GLOB.shape[0])))\n",
    "for NP in range(1,7):\n",
    "    print(NP)\n",
    "    with mp.Pool(NP) as pool :\n",
    "        %time result = pool.map(compute_distance_row, range(DATA_GLOB.shape[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Same as before: with larger individual takes we seem to get better speedup in general (~x2 speedup for 4 processes)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "<br>\n",
    "\n",
    "[Back to ToC](#toc)\n",
    "\n",
    "## Annex 2 - parallelization of pairwise distance computation with numba <a id='annexb'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import njit, prange  # njit -> no-python jit\n",
    "\n",
    "\n",
    "@njit(parallel=True)\n",
    "def pairwise_distance_numba_prange(X):\n",
    "\n",
    "    num_vectors = X.shape[0]\n",
    "    num_measurements = X.shape[1] \n",
    "    D = np.empty((num_vectors, num_vectors), dtype=np.float64)\n",
    "    \n",
    "    for i in prange(num_vectors): # Note the usage of prange\n",
    "        for j in range(num_vectors):\n",
    "            d = 0.\n",
    "            for k in range(num_measurements):\n",
    "                d += np.square(np.subtract(X[i][k], X[j][k]))\n",
    "            D[i, j] = np.sqrt(d)\n",
    "    return(D)\n",
    "\n",
    "# Create toy data to launch the function once and compile it.\n",
    "toy_data = np.random.uniform(size=(10,10))\n",
    "toy_result = pairwise_distance_numba_prange(toy_data) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"numba parallel=True\")\n",
    "%timeit -n 5 -r 3 result = pairwise_distance_numba_prange(big_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Environment (conda_py311)",
   "language": "python",
   "name": "conda_py311"
  },
  "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}