tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***https://github.com/jeremiedbb/tutorial-euroscipy-2019.git***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Requirements**\n",
    "- python > 3.5\n",
    "- jupyter\n",
    "\n",
    "\n",
    "- numpy\n",
    "- scipy\n",
    "\n",
    "\n",
    "- cython\n",
    "- numba\n",
    "\n",
    "\n",
    "- snakeviz\n",
    "\n",
    "\n",
    "- pandas\n",
    "- matplotlib\n",
    "\n",
    "\n",
    "- a C compiler supporting OpenMP\n",
    "  (for macOS users: ``conda install gcc``)\n",
    "  \n",
    "**For conda users**, you can get a ready-to go environment by cloning the repo and running ``conda env create``.\n",
    "  \n",
    "<br></br>\n",
    "<br></br>\n",
    "<br></br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy.spatial.distance import cdist\n",
    "import matplotlib.pyplot as plt\n",
    "import numba\n",
    "from numba import jit, njit, prange\n",
    "\n",
    "%matplotlib inline\n",
    "%load_ext cython\n",
    "%load_ext snakeviz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "<br></br>\n",
    "# Speed up your python code\n",
    "## An application with pairwise distances"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "### The problem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let ``X`` and ``Y`` be two sets of points. For all points in `X`, find its closest point in `Y` w.r.t some given distance ``d``.\n",
    "\n",
    "&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;$\\forall x\\in X, c(x) = \\underset{y\\in Y}{\\operatorname{argmin}}{d(x,y)}$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to\n",
    "- compute distances between all points in ``X`` and all points in ``Y``.\n",
    "  Done using scipy ``cdist``.\n",
    "- for each point in ``X`` find which point in ``Y`` minimize this distance.\n",
    "  Done using numpy ``argmin``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cdist_argmin(X, Y, metric):\n",
    "    distances = cdist(X, Y, metric=metric)\n",
    "    argmins = np.argmin(distances, axis=1)\n",
    "    \n",
    "    return argmins"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example we can use it to draw the Voronoi diagram of the set of points ``Y``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot(metric, title):\n",
    "    Xv = np.array([[0.005*(i//200), 0.005*(i%200)] for i in range(40000)])\n",
    "    Yv = np.random.RandomState(42).random_sample((20, 2))\n",
    "\n",
    "    indices = cdist_argmin(Xv,Yv, metric=metric)\n",
    "    plt.scatter(Xv[:,0], Xv[:,1], c=indices, s=10, marker='s')\n",
    "    plt.scatter(Yv[:,0], Yv[:,1], color='red')\n",
    "    plt.axis('equal')\n",
    "    plt.rcParams['figure.figsize'] = [10, 10]\n",
    "    plt.title(title)\n",
    "\n",
    "plot('euclidean', 'Euclidean metric')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot('cityblock', 'Cityblock metric')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "### Efficiency considerations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's generate some data for our benchmarks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = np.random.random_sample((10000, 100))\n",
    "Y = np.random.random_sample((100, 100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the ``cdist_argmin`` function, the critical part is the computation of the pairwise distances as shown in the following profiling. If we need to optimize something, this is the part of the code we want to look at."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%snakeviz cdist_argmin(X, Y, metric='cityblock')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "``cdist`` is very efficient when using a predefined metric but very slow when using a custom metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit cdist(X, Y, metric='cityblock')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cityblock(x, y):\n",
    "    v = np.abs(x - y)\n",
    "    \n",
    "    return v.sum()\n",
    "\n",
    "%timeit -r 3 -n 1 cdist(X, Y, metric=cityblock)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The reason is that for a custom metric ``cdist`` relies on generic python code whereas for a predinefined metric it can rely on a specific C implementation.\n",
    "\n",
    "It's a good candidate to move to compiled code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "### Our custom metric"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Through the tutorial we'll use the following metric:\n",
    "\n",
    "&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;$d(x,y) = \\sum_{i=0}^{k}{(i + 1) |x_i - y_i|}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def custom_metric(x, y):\n",
    "    v = np.abs(x - y)\n",
    "    w = np.arange(1, x.shape[0] + 1)\n",
    "    \n",
    "    return v.dot(w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(custom_metric, 'Custom metric')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Helpers for accuracy testing and benchmarking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(func, X, Y):\n",
    "    true_indices = cdist_argmin(X, Y)\n",
    "    indices = func(X, Y)\n",
    "    if np.all(true_indices == indices):\n",
    "        print('correct result \\o/')\n",
    "    else:\n",
    "        print('incorrect result TT')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "benchs = pd.DataFrame(columns=['version', 'time (s)', 'speedup'])\n",
    "\n",
    "\n",
    "def bench(func, X, Y, title):\n",
    "    t = 0\n",
    "    for i in range(20):\n",
    "        t_ = time.time()\n",
    "        func(X, Y)\n",
    "        t_ = time.time() - t_\n",
    "        t += t_\n",
    "        if t_ > 0.5 and i >=3:\n",
    "            break\n",
    "    t /= (i + 1)\n",
    "\n",
    "    i = benchs.shape[0]\n",
    "    if i == 0:\n",
    "        speedup = 1\n",
    "    else:\n",
    "        speedup = benchs.iloc[0]['time (s)'] / t\n",
    "    benchs.loc[i] = [title, t, speedup]\n",
    "\n",
    "    print(benchs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "### Baseline: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cdist_argmin(X, Y):\n",
    "    distances = cdist(X, Y, metric=custom_metric)\n",
    "    argmins = np.argmin(distances, axis=1)\n",
    "    \n",
    "    return argmins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(cdist_argmin, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(cdist_argmin, X, Y, 'baseline')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "### Cython"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### - Python is valid Cython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -f\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def custom_metric(x, y):\n",
    "    v = np.abs(x - y)\n",
    "    w = np.arange(1, x.shape[0] + 1)\n",
    "    \n",
    "    return np.dot(v, w)\n",
    "\n",
    "def cython_cdist_argmin(X, Y):\n",
    "    distances = np.empty((X.shape[0], Y.shape[0]))\n",
    "    \n",
    "    for i, x in enumerate(X):\n",
    "        for j, y in enumerate(Y):\n",
    "            distances[i, j] = custom_metric(x, y)\n",
    "    \n",
    "    return np.argmin(distances, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(cython_cdist_argmin, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(cython_cdist_argmin, X, Y, \"cython pure python\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "**Cython annotations** (-a option) show which line contains interaction with python objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -a -f\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def custom_metric(x, y):\n",
    "    v = np.abs(x - y)\n",
    "    w = np.arange(1, x.shape[0] + 1)\n",
    "\n",
    "    return np.dot(v, w)\n",
    "\n",
    "\n",
    "def cython_cdist_argmin(X, Y):\n",
    "    distances = np.empty((X.shape[0], Y.shape[0]))\n",
    "    \n",
    "    for i, x in enumerate(X):\n",
    "        for j, y in enumerate(Y):\n",
    "            distances[i, j] = custom_metric(x, y)\n",
    "    \n",
    "    return np.argmin(distances, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "#### - Get rid of python interactions within the nested loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Move from python variables to typed variables  \n",
    "``cdef <type> var = ...``"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Python ``range`` keyword is compiled into efficient C code  \n",
    "e.g. ``for i in range(10)``  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -a -f\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def custom_metric(x, y):\n",
    "    v = np.abs(x - y)\n",
    "    w = np.arange(1, x.shape[0] + 1)\n",
    "\n",
    "    return np.dot(v, w)\n",
    "\n",
    "\n",
    "def cython_cdist_argmin(X, Y):\n",
    "    cdef:\n",
    "        int nx = X.shape[0]\n",
    "        int ny = Y.shape[0]\n",
    "        int i, j\n",
    "\n",
    "    distances = np.empty((X.shape[0], Y.shape[0]))\n",
    "    \n",
    "    for i in range(nx):\n",
    "        for j in range(ny):\n",
    "            distances[i, j] = custom_metric(X[i], Y[j])\n",
    "    \n",
    "    return np.argmin(distances, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "``distances`` is still a python object (numpy array). We can get an efficient access to its data using a **typed memoryview**.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -a -f\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def custom_metric(x, y):\n",
    "    v = np.abs(x - y)\n",
    "    w = np.arange(1, x.shape[0] + 1)\n",
    "\n",
    "    return np.dot(v, w)\n",
    "\n",
    "\n",
    "def cython_cdist_argmin(X, Y):\n",
    "    cdef:\n",
    "        int nx = X.shape[0]\n",
    "        int ny = Y.shape[0]\n",
    "        int i, j\n",
    "\n",
    "        double[:, :] distances = np.empty((nx, ny))\n",
    "    \n",
    "    for i in range(nx):\n",
    "        for j in range(ny):\n",
    "            distances[i, j] = custom_metric(X[i], Y[j])\n",
    "    \n",
    "    return np.argmin(distances, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "``custom_metric`` is a python function with python arguments implemented in pure python.  \n",
    "- functions maint to be used in python should be defined using `def`  \n",
    "- functions maint to be used within cython should be defined using `cdef`, and typed.\n",
    "\n",
    "We also need to use memoryviews for X and Y."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -a -f\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "cdef double custom_metric(double[:] x, double[:] y):\n",
    "    cdef:\n",
    "        int n = x.shape[0]\n",
    "        double tmp\n",
    "        double res = 0.0\n",
    "        int i\n",
    "\n",
    "    for i in range(n):\n",
    "        tmp = x[i] - y[i]\n",
    "        if tmp < 0:\n",
    "            tmp = -tmp\n",
    "        res += (i + 1) * tmp\n",
    "        \n",
    "    return res\n",
    "\n",
    "\n",
    "def cython_cdist_argmin(double[:, :] X, double[:, :] Y):\n",
    "    cdef:\n",
    "        int nx = X.shape[0]\n",
    "        int ny = Y.shape[0]\n",
    "        int i, j\n",
    "\n",
    "        double[:, :] distances = np.empty((nx, ny))\n",
    "    \n",
    "    for i in range(nx):\n",
    "        for j in range(ny):\n",
    "            distances[i, j] = custom_metric(X[i], Y[j])\n",
    "    \n",
    "    return np.argmin(np.asarray(distances), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(cython_cdist_argmin, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(cython_cdist_argmin, X, Y, \"cdef custom_metric\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "Cython checks that indexing memoryviews stays within bounds and allows negative indexing. We can disable both using ``boundscheck`` and ``wraparound`` **cython directives**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -a -f\n",
    "#cython: boundscheck=False, wraparound=False\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "cdef double custom_metric(double[:] x, double[:] y):\n",
    "    cdef:\n",
    "        int n = x.shape[0]\n",
    "        double tmp\n",
    "        double res = 0.0\n",
    "        int i\n",
    "\n",
    "    for i in range(n):\n",
    "        tmp = x[i] - y[i]\n",
    "        if tmp < 0:\n",
    "            tmp = -tmp\n",
    "        res += (i + 1) * tmp\n",
    "        \n",
    "    return res\n",
    "\n",
    "\n",
    "def cython_cdist_argmin(double[:, :] X, double[:, :] Y):\n",
    "    cdef:\n",
    "        int nx = X.shape[0]\n",
    "        int ny = Y.shape[0]\n",
    "        int i, j\n",
    "\n",
    "        double[:, :] distances = np.empty((nx, ny))\n",
    "    \n",
    "    for i in range(nx):\n",
    "        for j in range(ny):\n",
    "            distances[i, j] = custom_metric(X[i], Y[j])\n",
    "    \n",
    "    return np.argmin(np.asarray(distances), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(cython_cdist_argmin, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(cython_cdist_argmin, X, Y, \"cython directives\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "#### - Getting the most out of Cython"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we know X and Y will always have the same **memory layout** we can give that information to cython which will produce specialized code.  \n",
    "Let's assume X and Y are C ordered."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -a -f\n",
    "#cython: boundscheck=False, wraparound=False\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "cdef double custom_metric(double[::1] x, double[::1] y):\n",
    "    cdef:\n",
    "        int n = x.shape[0]\n",
    "        double tmp\n",
    "        double res = 0.0\n",
    "        int i\n",
    "\n",
    "    for i in range(n):\n",
    "        tmp = x[i] - y[i]\n",
    "        if tmp < 0:\n",
    "            tmp = -tmp\n",
    "        res += (i + 1) * tmp\n",
    "        \n",
    "    return res\n",
    "\n",
    "\n",
    "def cython_cdist_argmin(double[:, ::1] X, double[:, ::1] Y):\n",
    "    cdef:\n",
    "        int nx = X.shape[0]\n",
    "        int ny = Y.shape[0]\n",
    "        int i, j\n",
    "\n",
    "        double[:, ::1] distances = np.empty((nx, ny))\n",
    "    \n",
    "    for i in range(nx):\n",
    "        for j in range(ny):\n",
    "            distances[i, j] = custom_metric(X[i], Y[j])\n",
    "    \n",
    "    return np.argmin(np.asarray(distances), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(cython_cdist_argmin, X, Y, \"C ordered\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use functions already implemented in the C standard library: ``fabs``.  \n",
    "C functions need to be imported with the ``cimport`` keyword."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -f\n",
    "#cython: boundscheck=False, wraparound=False\n",
    "\n",
    "import numpy as np\n",
    "from libc.math cimport fabs\n",
    "\n",
    "\n",
    "cdef double custom_metric(double[::1] x, double[::1] y):\n",
    "    cdef:\n",
    "        int n = x.shape[0]\n",
    "        double res = 0.0\n",
    "        int i\n",
    "\n",
    "    for i in range(n):\n",
    "        res += (i + 1) * fabs(x[i] - y[i])\n",
    "        \n",
    "    return res\n",
    "\n",
    "\n",
    "def cython_cdist_argmin(double[:, ::1] X, double[:, ::1] Y):\n",
    "    cdef:\n",
    "        int nx = X.shape[0]\n",
    "        int ny = Y.shape[0]\n",
    "        int i, j\n",
    "\n",
    "        double[:, ::1] distances = np.empty((nx, ny))\n",
    "    \n",
    "    for i in range(nx):\n",
    "        for j in range(ny):\n",
    "            distances[i, j] = custom_metric(X[i], Y[j])\n",
    "    \n",
    "    return np.argmin(np.asarray(distances), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(cython_cdist_argmin, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(cython_cdist_argmin, X, Y, \"fabs\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "GCC can't **vectorize** reduction.\n",
    "We can either:\n",
    "- tell GCC that it's fine to reorder operations\n",
    "- manually unroll the loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can specify compiler flags through the ``-c`` option."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -c=-Ofast -f\n",
    "#cython: boundscheck=False, wraparound=False\n",
    "\n",
    "import numpy as np\n",
    "from libc.math cimport fabs\n",
    "\n",
    "\n",
    "cdef double custom_metric(double[:] x, double[:] y):\n",
    "    cdef:\n",
    "        int n = x.shape[0]\n",
    "        double res = 0.0\n",
    "        int i\n",
    "\n",
    "    for i in range(n):\n",
    "        res += (i + 1) * fabs(x[i] - y[i])\n",
    "        \n",
    "    return res\n",
    "\n",
    "\n",
    "def cython_cdist_argmin(double[:, :] X, double[:, :] Y):\n",
    "    cdef:\n",
    "        int nx = X.shape[0]\n",
    "        int ny = Y.shape[0]\n",
    "        int i, j\n",
    "\n",
    "        double[:, :] distances = np.empty((nx, ny))\n",
    "    \n",
    "    for i in range(nx):\n",
    "        for j in range(ny):\n",
    "            distances[i, j] = custom_metric(X[i], Y[j])\n",
    "    \n",
    "    return np.argmin(np.asarray(distances), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(cython_cdist_argmin, X, Y, \"-Ofast\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Manually unroll the loop will be worth only if we use pointers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -f\n",
    "#cython: boundscheck=False, wraparound=False\n",
    "\n",
    "import numpy as np\n",
    "from libc.math cimport fabs\n",
    "\n",
    "\n",
    "cdef double custom_metric(double* x, double* y, int k):\n",
    "    cdef:\n",
    "        double res = 0.0\n",
    "        int i\n",
    "\n",
    "    for i in range(k):\n",
    "        res += (i + 1) * fabs(x[i] - y[i])\n",
    "\n",
    "    return res\n",
    "\n",
    "\n",
    "def cython_cdist_argmin(double[:, ::1] X, double[:, ::1] Y):\n",
    "    cdef:\n",
    "        int nx = X.shape[0]\n",
    "        int ny = Y.shape[0]\n",
    "        int k = X.shape[1]\n",
    "        int i, j\n",
    "\n",
    "        double[:, ::1] distances = np.empty((nx, ny))\n",
    "    \n",
    "    for i in range(nx):\n",
    "        for j in range(ny):\n",
    "            distances[i, j] = custom_metric(&X[i, 0], &Y[j, 0], k)\n",
    "    \n",
    "    return np.argmin(np.asarray(distances), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(cython_cdist_argmin, X, Y, \"pointers\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -f\n",
    "#cython: boundscheck=False, wraparound=False\n",
    "\n",
    "import numpy as np\n",
    "from libc.math cimport fabs\n",
    "\n",
    "\n",
    "cdef double custom_metric(double* x, double* y, int k):\n",
    "    cdef:\n",
    "        int rem = k % 4\n",
    "        double res = 0.0\n",
    "        int i\n",
    "\n",
    "    for i in range(0, k-rem, 4):\n",
    "        res += ((i + 1) * fabs(x[i] - y[i])\n",
    "               +(i + 2) * fabs(x[i + 1] - y[i + 1])\n",
    "               +(i + 3) * fabs(x[i + 2] - y[i + 2])\n",
    "               +(i + 4) * fabs(x[i + 3] - y[i + 3]))\n",
    "    \n",
    "    for i in range(k-rem, k):\n",
    "        res += (i + 1) * fabs(x[i] - y[i])\n",
    "     \n",
    "    return res\n",
    "\n",
    "\n",
    "def cython_cdist_argmin(double[:, ::1] X, double[:, ::1] Y):\n",
    "    cdef:\n",
    "        int nx = X.shape[0]\n",
    "        int ny = Y.shape[0]\n",
    "        int k = X.shape[1]\n",
    "        int i, j\n",
    "\n",
    "        double[:, ::1] distances = np.empty((nx, ny))\n",
    "    \n",
    "    for i in range(nx):\n",
    "        for j in range(ny):\n",
    "            distances[i, j] = custom_metric(&X[i, 0], &Y[j, 0], k)\n",
    "    \n",
    "    return np.argmin(np.asarray(distances), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(cython_cdist_argmin, X, Y, \"pointers + unrolled loop\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "### More cython"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Profiling cython code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Enabling the ``profile`` cython directive will allow us to profile our code with ``snakeviz``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -f\n",
    "#cython: boundscheck=False, wraparound=False, profile=True\n",
    "\n",
    "import numpy as np\n",
    "from libc.math cimport fabs\n",
    "\n",
    "\n",
    "cdef double custom_metric(double* x, double* y, int k):\n",
    "    cdef:\n",
    "        int rem = k % 4\n",
    "        double res = 0.0\n",
    "        int i\n",
    "\n",
    "    for i in range(0, k, 4):\n",
    "        res += ((i + 1) * fabs(x[i] - y[i])\n",
    "               +(i + 2) * fabs(x[i + 1] - y[i + 1])\n",
    "               +(i + 3) * fabs(x[i + 2] - y[i + 2])\n",
    "               +(i + 4) * fabs(x[i + 3] - y[i + 3]))\n",
    "    \n",
    "    for i in range(i, rem):\n",
    "        res += (i + 1) * fabs(x[i] - y[i])\n",
    "     \n",
    "    return res\n",
    "\n",
    "\n",
    "def cython_cdist_argmin(double[:, ::1] X, double[:, ::1] Y):\n",
    "    cdef:\n",
    "        int nx = X.shape[0]\n",
    "        int ny = Y.shape[0]\n",
    "        int k = X.shape[1]\n",
    "        int i, j\n",
    "\n",
    "        double[:, ::1] distances = np.empty((nx, ny))\n",
    "    \n",
    "    for i in range(nx):\n",
    "        for j in range(ny):\n",
    "            distances[i, j] = custom_metric(&X[i, 0], &Y[j, 0], k)\n",
    "    \n",
    "    return np.argmin(np.asarray(distances), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%snakeviz cython_cdist_argmin(X, Y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "#### Generic code with fused types"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our current code is specific to ``double`` dtype (np.float64)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cython_cdist_argmin(X.astype(np.float32), Y.astype(np.float32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -f\n",
    "#cython: boundscheck=False, wraparound=False\n",
    "\n",
    "import numpy as np\n",
    "from libc.math cimport fabs\n",
    "\n",
    "ctypedef fused floating:\n",
    "    float\n",
    "    double\n",
    "\n",
    "cdef floating custom_metric(floating* x, floating* y, int k):\n",
    "    cdef:\n",
    "        int rem = k % 4\n",
    "        floating res = 0.0\n",
    "        int i\n",
    "\n",
    "    for i in range(0, k, 4):\n",
    "        res += ((i + 1) * fabs(x[i] - y[i])\n",
    "               +(i + 2) * fabs(x[i + 1] - y[i + 1])\n",
    "               +(i + 3) * fabs(x[i + 2] - y[i + 2])\n",
    "               +(i + 4) * fabs(x[i + 3] - y[i + 3]))\n",
    "    \n",
    "    for i in range(i, rem):\n",
    "        res += (i + 1) * fabs(x[i] - y[i])\n",
    "     \n",
    "    return res\n",
    "\n",
    "\n",
    "def cython_cdist_argmin(floating[:, ::1] X, floating[:, ::1] Y):\n",
    "    cdef:\n",
    "        int nx = X.shape[0]\n",
    "        int ny = Y.shape[0]\n",
    "        int k = X.shape[1]\n",
    "        int i, j\n",
    "\n",
    "    dtype = np.float32 if floating is float else np.float64\n",
    "    cdef floating[:, ::1] distances = np.empty((nx, ny), dtype=dtype)\n",
    "    \n",
    "    for i in range(nx):\n",
    "        for j in range(ny):\n",
    "            distances[i, j] = custom_metric(&X[i, 0], &Y[j, 0], k)\n",
    "    \n",
    "    return np.argmin(np.asarray(distances), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cython_cdist_argmin(X.astype(np.float32), Y.astype(np.float32))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "#### Parallelism with OpenMP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simply replace the ``range`` keyword by the ``prange`` keyword."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -c=-fopenmp -f\n",
    "#cython: boundscheck=False, wraparound=False\n",
    "\n",
    "import numpy as np\n",
    "from libc.math cimport fabs\n",
    "from cython.parallel cimport prange\n",
    "\n",
    "ctypedef fused floating:\n",
    "    float\n",
    "    double\n",
    "\n",
    "cdef floating custom_metric(floating* x, floating* y, int k) nogil:\n",
    "    cdef:\n",
    "        int rem = k % 4\n",
    "        floating res = 0.0\n",
    "        int i\n",
    "\n",
    "    for i in range(0, k, 4):\n",
    "        res += ((i + 1) * fabs(x[i] - y[i])\n",
    "               +(i + 2) * fabs(x[i + 1] - y[i + 1])\n",
    "               +(i + 3) * fabs(x[i + 2] - y[i + 2])\n",
    "               +(i + 4) * fabs(x[i + 3] - y[i + 3]))\n",
    "    \n",
    "    for i in range(i, rem):\n",
    "        res += (i + 1) * fabs(x[i] - y[i])\n",
    "     \n",
    "    return res\n",
    "\n",
    "\n",
    "def cython_cdist_argmin(floating[:, ::1] X, floating[:, ::1] Y):\n",
    "    cdef:\n",
    "        int nx = X.shape[0]\n",
    "        int ny = Y.shape[0]\n",
    "        int k = X.shape[1]\n",
    "        int i, j\n",
    "\n",
    "    dtype = np.float32 if floating is float else np.float64\n",
    "    cdef floating[:, ::1] distances = np.empty((nx, ny), dtype=dtype)\n",
    "    \n",
    "    for i in prange(nx, num_threads=2, nogil=True):\n",
    "        for j in range(ny):\n",
    "            distances[i, j] = custom_metric(&X[i, 0], &Y[j, 0], k)\n",
    "    \n",
    "    return np.argmin(np.asarray(distances), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(cython_cdist_argmin, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(cython_cdist_argmin, X, Y, \"prange\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "<br></br>\n",
    "<br></br>\n",
    "<br></br>\n",
    "### Numba"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Like Cython, Numba can work with Python objects but is much more efficient if it can compile a function without involving the Python interpreter. It's called the **nopython** mode.\n",
    "\n",
    "To tell Numba which function needs to be compiled, you just have to decorate it: ``@jit``.   \n",
    "``@jit(nopython=True)`` or ``@njit`` enforces the nopython mode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit\n",
    "def f(x, y):\n",
    "    return x.dot(y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "``f`` is compiled when called for the first time so the first call is longer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time f(X[0], Y[0])\n",
    "%time f(X[0], Y[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "#### - Type inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Numba infers the types of the function and its argument at first compilation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f.inspect_types()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "#### - Numba and Numpy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Numba understands Numpy arrays and supports a lot of Numpy features, meaning using Numpy functions in nopython mode will be compiled into efficient code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit\n",
    "def numba_custom_metric(x, y):\n",
    "    v = np.abs(x - y)\n",
    "    w = np.arange(1, x.shape[0] + 1)\n",
    "    \n",
    "    return v.dot(w)\n",
    "\n",
    "\n",
    "@njit\n",
    "def numba_cdist_argmin(X, Y):\n",
    "    nx = X.shape[0]\n",
    "    ny = Y.shape[0]\n",
    "\n",
    "    distances = np.empty((nx, ny))\n",
    "    \n",
    "    for i in range(nx):\n",
    "        for j in range(ny):\n",
    "            distances[i, j] = numba_custom_metric(X[i], Y[j])\n",
    "    \n",
    "    return np.argmin(distances, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However this code doesn't compile because:\n",
    "- Numba support for ``np.dot`` needs to be called with arrays of same dtype\n",
    "- Numba support for ``np.argmin`` does not handle the ``axis`` parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit\n",
    "def numba_custom_metric(x, y):\n",
    "    v = np.abs(x - y)\n",
    "    w = np.arange(1, x.shape[0] + 1).astype(x.dtype)\n",
    "    \n",
    "    return v.dot(w)\n",
    "\n",
    "\n",
    "@njit\n",
    "def numba_cdist(X, Y):\n",
    "    nx = X.shape[0]\n",
    "    ny = Y.shape[0]\n",
    "\n",
    "    distances = np.empty((nx, ny))\n",
    "    \n",
    "    for i in range(nx):\n",
    "        for j in range(ny):\n",
    "            distances[i, j] = numba_custom_metric(X[i], Y[j])\n",
    "    \n",
    "    return distances\n",
    "\n",
    "            \n",
    "def numba_cdist_argmin(X, Y):\n",
    "    distances = numba_cdist(X, Y)\n",
    "\n",
    "    return np.argmin(distances, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "test(numba_cdist_argmin, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(numba_cdist_argmin, X, Y, \"naive numba\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "#### - Fastmath optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Like with cython, telling the compiler that it can reorder operation as it wants can speed things up significantly. This is done through the **fastmath** option.\n",
    "\n",
    "However we need to rewrite our function in a more C-like way to benefit from it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit(fastmath=True)\n",
    "def numba_custom_metric(x, y):\n",
    "    res = 0\n",
    "\n",
    "    for i in range(x.shape[0]):\n",
    "        v = x[i] - y[i]\n",
    "        if v < 0:\n",
    "            v = - v\n",
    "        res += (i + 1) * v\n",
    "    \n",
    "    return res\n",
    "\n",
    "\n",
    "@njit(fastmath=True)\n",
    "def numba_cdist(X, Y):\n",
    "    nx = X.shape[0]\n",
    "    ny = Y.shape[0]\n",
    "\n",
    "    distances = np.empty((nx, ny))\n",
    "    \n",
    "    for i in range(nx):\n",
    "        for j in range(ny):\n",
    "            distances[i, j] = numba_custom_metric(X[i], Y[j])\n",
    "    \n",
    "    return distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(numba_cdist_argmin, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(numba_cdist_argmin, X, Y, \"numba fastmath\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "### More Numba"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### - Profiling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can profile Numba compiled functions without additional code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%snakeviz numba_cdist_argmin(X, Y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "#### - Parallel loop in Numba"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simply replace ``range`` by ``prange`` and enable the ``parallel`` option.\n",
    "Currently, the only way to specify the number of threads to be use for the ``prange`` is through the ``NUMBA_NUM_THREADS`` environment variable (before imports)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit(parallel=True)\n",
    "def numba_cdist(X, Y):\n",
    "    nx = X.shape[0]\n",
    "    ny = Y.shape[0]\n",
    "\n",
    "    distances = np.empty((nx, ny))\n",
    "    \n",
    "    for i in prange(nx):\n",
    "        for j in range(ny):\n",
    "            distances[i, j] = numba_custom_metric(X[i], Y[j])\n",
    "    \n",
    "    return distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(numba_cdist_argmin, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(numba_cdist_argmin, X, Y, \"numba prange\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Numba supports the OpenMP and TBB **threading layers** (and a simple built-in threading layer in case none of those is found).\n",
    "\n",
    "``pip install tbb`` will install TBB."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "numba.threading_layer()"
   ]
  },
  {
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}