cython_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "from scipy.spatial.distance import cdist\n",
    "import pandas as pd\n",
    "from sklearn.metrics import pairwise_distances_argmin\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "%matplotlib inline\n",
    "%load_ext cython\n",
    "%load_ext snakeviz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cython tutorial\n",
    "## An application with pairwise distances"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Profiling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Computing pairwise distances between points and centers is the critical part of K-means algorithm. Let's check that with the snakeviz profiling tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.cluster import KMeans\n",
    "\n",
    "kmeans = KMeans(n_clusters=100, n_init=1, algorithm='full', init='random')\n",
    "X = np.random.random_sample((10000, 100))\n",
    "\n",
    "%snakeviz kmeans.fit(X)"
   ]
  },
  {
   "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`. Let's add a bit of parietal by doing it w.r.t the L1-Norm.\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}}{||x-y||_1}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot():\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 = pairwise_distances_argmin(Xv,Yv, metric='manhattan')\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",
    "\n",
    "plot()"
   ]
  },
  {
   "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 = pairwise_distances_argmin(X, Y, metric='manhattan')\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",
    "def bench(func, X, Y, title):\n",
    "    if title not in benchs:\n",
    "        t = 0\n",
    "        for _ in range(10):\n",
    "            t_ = time.time()\n",
    "            func(X, Y)\n",
    "            t_ = time.time() - t_\n",
    "            t += t_\n",
    "        t /= 10\n",
    "        \n",
    "        i = benchs.shape[0]\n",
    "        if i == 0:\n",
    "            speedup = 1\n",
    "        else:\n",
    "            speedup = benchs[benchs['version']=='naive python']['time(s)'].values[0] / t\n",
    "        benchs.loc[i] = [title, t, speedup]\n",
    "        \n",
    "        print(benchs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "Let's generate some data to benchmark our experiments."
   ]
  },
  {
   "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": [
    "<br></br>\n",
    "### Baseline: naive python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def naive_python(X, Y):\n",
    "    distances = np.empty((X.shape[0], Y.shape[0]))\n",
    "    for i, x in enumerate(X):\n",
    "        for j, y in enumerate(Y):\n",
    "            distances[i, j] = np.abs(x - y).sum()\n",
    "    \n",
    "    return np.argmin(distances, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(naive_python, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(naive_python, X, Y, 'naive python')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%memit naive_python(X, Y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "### Goal"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Scipy `cdist` is a function implemented in C which computes the pairwise distances between two sets of points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def scipy_pairwise(X, Y):\n",
    "    distances = cdist(X, Y, metric='cityblock')\n",
    "    return np.argmin(distances, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(scipy_pairwise, X, Y, 'cdist')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%memit scipy_pairwise(X, Y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There's a helper in scikit-learn which does the distances and argmin computations (using scipy's cdist under the hood)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sklearn_pairwise(X, Y):\n",
    "    return pairwise_distances_argmin(X, Y, metric='manhattan')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(sklearn_pairwise, X, Y, 'sklearn')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "### numpy only ?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using numpy only requires making use of broadcasting, which may increase memory consumption."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def naive_numpy(X, Y):\n",
    "    distances = np.abs(X[:,:,np.newaxis] - Y.T[np.newaxis,:,:]).sum(axis=1)\n",
    "    \n",
    "    return np.argmin(distances, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(naive_numpy, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(naive_numpy, X, Y, 'naive numpy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%memit naive_numpy(X, Y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "### Cython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%cython -f\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def cython_pairwise(X, Y):\n",
    "    distances = np.empty((X.shape[0], Y.shape[0]))\n",
    "    for i, x in enumerate(X):\n",
    "        for j, y in enumerate(Y):\n",
    "            distances[i, j] = np.abs(x - y).sum()\n",
    "    \n",
    "    return np.argmin(distances, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(cython_pairwise, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(cython_pairwise, X, Y, 'cython nothing')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "## More cython"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fused types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -f\n",
    "import numpy as np\n",
    "from libc.math cimport sqrt\n",
    "\n",
    "\n",
    "cpdef double euclidean_distance(double[::1] x, double[::1] y):\n",
    "    cdef:\n",
    "        int i\n",
    "        double res = 0\n",
    "\n",
    "    for i in range(x.shape[0]):\n",
    "        res += (x[i] - y[i]) * (x[i] - y[i])\n",
    "\n",
    "    return sqrt(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.ones(100)\n",
    "y = np.zeros(100)\n",
    "euclidean_distance(x, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "### Paralellism with OpenMP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First a simple parallel loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(cython_pairwise, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(cython_pairwise, X, Y, 'cython prange')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br></br>\n",
    "Now a parallel loop with local buffer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%cython -c=-fopenmp -f\n",
    "#cython: wraparound=False, boundscheck=False\n",
    "import numpy as np\n",
    "from cython.parallel cimport parallel, prange\n",
    "from libc.math cimport fabs\n",
    "from libc.stdlib cimport malloc, free\n",
    "\n",
    "\n",
    "cdef double manhattan_distance(double *x, double *y, int n) nogil:\n",
    "    cdef:\n",
    "        int i\n",
    "        double res = 0\n",
    "        int k = n // 4\n",
    "        int r = n % 4\n",
    "    \n",
    "    for i in range(k):\n",
    "        res += (fabs(x[0] - y[0])\n",
    "               +fabs(x[1] - y[1])\n",
    "               +fabs(x[2] - y[2])\n",
    "               +fabs(x[3] - y[3]))\n",
    "        x += 4; y += 4\n",
    "    \n",
    "    for i in range(r):\n",
    "        res += fabs(x[i] - y[i])\n",
    "    \n",
    "    return res\n",
    "\n",
    "\n",
    "cdef void manhattan_distances_chunked(double *X, int nx,\n",
    "                                      double *Y, int ny,\n",
    "                                      int n_features,\n",
    "                                      double *distances) nogil:\n",
    "    cdef:\n",
    "        int i, j\n",
    "        \n",
    "    for i in range(nx):\n",
    "        for j in range(ny):\n",
    "            distances[i * ny + j] = manhattan_distance(X + i * n_features, Y + j * n_features, n_features)\n",
    "                \n",
    "\n",
    "cdef void argmin_chunked(double *distances, int m, int n, int *indices) nogil:\n",
    "    cdef:\n",
    "        int i, j\n",
    "        int best_j\n",
    "        double min_dist\n",
    "        \n",
    "    for i in range(m):\n",
    "        min_dist = distances[i * n] \n",
    "        best_j = 0\n",
    "        \n",
    "        for j in range(1, n):\n",
    "            if distances[i * n + j] < min_dist:\n",
    "                min_dist = distances[i * n + j]\n",
    "                best_j = j\n",
    "        \n",
    "        indices[i] = best_j\n",
    "            \n",
    "\n",
    "\n",
    "def cython_parallel_pairwise(double[:, ::1] X, double[:, ::1] Y):\n",
    "    cdef:\n",
    "        int i\n",
    "        int n_samples_X = X.shape[0]\n",
    "        int n_samples_Y = Y.shape[0]\n",
    "        int n_features = X.shape[1]\n",
    "        \n",
    "        int[::1] indices = np.empty(n_samples_X, dtype=np.int32)\n",
    "        \n",
    "        double *local_buffer\n",
    "        int nx = 100\n",
    "        int n_chunks = n_samples_X // nx\n",
    "\n",
    "    with nogil, parallel():\n",
    "        local_buffer = <double *> malloc(nx * n_samples_Y * sizeof(double))\n",
    "        \n",
    "        for i in prange(n_chunks):\n",
    "            manhattan_distances_chunked(&X[i * nx, 0], nx,\n",
    "                                        &Y[0, 0], n_samples_Y,\n",
    "                                        n_features,\n",
    "                                        local_buffer)\n",
    "        \n",
    "            argmin_chunked(local_buffer, nx, n_samples_Y, &indices[i * nx])\n",
    "        \n",
    "        free(local_buffer)\n",
    "            \n",
    "    return np.asarray(indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(cython_parallel_pairwise, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench(cython_parallel_pairwise, X, Y, 'cython prange buffer')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}