Kullback-Leibler_divergences_in_native_Python__Cython_and_Numba.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "toc": "true"
   },
   "source": [
    "# Table of Contents\n",
    " <p><div class=\"lev1 toc-item\"><a href=\"#Introduction\" data-toc-modified-id=\"Introduction-1\"><span class=\"toc-item-num\">1&nbsp;&nbsp;</span>Introduction</a></div><div class=\"lev1 toc-item\"><a href=\"#KL-divergences-and-KL-UCB-indexes,-in-naive-Python\" data-toc-modified-id=\"KL-divergences-and-KL-UCB-indexes,-in-naive-Python-2\"><span class=\"toc-item-num\">2&nbsp;&nbsp;</span>KL divergences and KL-UCB indexes, in naive Python</a></div><div class=\"lev2 toc-item\"><a href=\"#KL-divergences\" data-toc-modified-id=\"KL-divergences-21\"><span class=\"toc-item-num\">2.1&nbsp;&nbsp;</span>KL divergences</a></div><div class=\"lev3 toc-item\"><a href=\"#Bernoulli-distributions\" data-toc-modified-id=\"Bernoulli-distributions-211\"><span class=\"toc-item-num\">2.1.1&nbsp;&nbsp;</span>Bernoulli distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Binomial-distributions\" data-toc-modified-id=\"Binomial-distributions-212\"><span class=\"toc-item-num\">2.1.2&nbsp;&nbsp;</span>Binomial distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Poisson-distributions\" data-toc-modified-id=\"Poisson-distributions-213\"><span class=\"toc-item-num\">2.1.3&nbsp;&nbsp;</span>Poisson distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Exponential-distributions\" data-toc-modified-id=\"Exponential-distributions-214\"><span class=\"toc-item-num\">2.1.4&nbsp;&nbsp;</span>Exponential distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Gamma-distributions\" data-toc-modified-id=\"Gamma-distributions-215\"><span class=\"toc-item-num\">2.1.5&nbsp;&nbsp;</span>Gamma distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Negative-binomial-distributions\" data-toc-modified-id=\"Negative-binomial-distributions-216\"><span class=\"toc-item-num\">2.1.6&nbsp;&nbsp;</span>Negative binomial distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Gaussian-distributions\" data-toc-modified-id=\"Gaussian-distributions-217\"><span class=\"toc-item-num\">2.1.7&nbsp;&nbsp;</span>Gaussian distributions</a></div><div class=\"lev2 toc-item\"><a href=\"#Generic-KL-UCB-indexes,-with-a-bisection-search\" data-toc-modified-id=\"Generic-KL-UCB-indexes,-with-a-bisection-search-22\"><span class=\"toc-item-num\">2.2&nbsp;&nbsp;</span>Generic KL-UCB indexes, with a bisection search</a></div><div class=\"lev2 toc-item\"><a href=\"#Distribution-specific-KL-UCB-indexes\" data-toc-modified-id=\"Distribution-specific-KL-UCB-indexes-23\"><span class=\"toc-item-num\">2.3&nbsp;&nbsp;</span>Distribution-specific KL-UCB indexes</a></div><div class=\"lev3 toc-item\"><a href=\"#Gaussian\" data-toc-modified-id=\"Gaussian-231\"><span class=\"toc-item-num\">2.3.1&nbsp;&nbsp;</span>Gaussian</a></div><div class=\"lev3 toc-item\"><a href=\"#Bernoulli\" data-toc-modified-id=\"Bernoulli-232\"><span class=\"toc-item-num\">2.3.2&nbsp;&nbsp;</span>Bernoulli</a></div><div class=\"lev3 toc-item\"><a href=\"#Poisson\" data-toc-modified-id=\"Poisson-233\"><span class=\"toc-item-num\">2.3.3&nbsp;&nbsp;</span>Poisson</a></div><div class=\"lev3 toc-item\"><a href=\"#Exponential\" data-toc-modified-id=\"Exponential-234\"><span class=\"toc-item-num\">2.3.4&nbsp;&nbsp;</span>Exponential</a></div><div class=\"lev3 toc-item\"><a href=\"#Others\" data-toc-modified-id=\"Others-235\"><span class=\"toc-item-num\">2.3.5&nbsp;&nbsp;</span>Others</a></div><div class=\"lev1 toc-item\"><a href=\"#With-Numba\" data-toc-modified-id=\"With-Numba-3\"><span class=\"toc-item-num\">3&nbsp;&nbsp;</span>With Numba</a></div><div class=\"lev2 toc-item\"><a href=\"#KL-divergences\" data-toc-modified-id=\"KL-divergences-31\"><span class=\"toc-item-num\">3.1&nbsp;&nbsp;</span>KL divergences</a></div><div class=\"lev3 toc-item\"><a href=\"#Bernoulli-distributions\" data-toc-modified-id=\"Bernoulli-distributions-311\"><span class=\"toc-item-num\">3.1.1&nbsp;&nbsp;</span>Bernoulli distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Binomial-distributions\" data-toc-modified-id=\"Binomial-distributions-312\"><span class=\"toc-item-num\">3.1.2&nbsp;&nbsp;</span>Binomial distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Poisson-distributions\" data-toc-modified-id=\"Poisson-distributions-313\"><span class=\"toc-item-num\">3.1.3&nbsp;&nbsp;</span>Poisson distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Exponential-distributions\" data-toc-modified-id=\"Exponential-distributions-314\"><span class=\"toc-item-num\">3.1.4&nbsp;&nbsp;</span>Exponential distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Gamma-distributions\" data-toc-modified-id=\"Gamma-distributions-315\"><span class=\"toc-item-num\">3.1.5&nbsp;&nbsp;</span>Gamma distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Negative-binomial-distributions\" data-toc-modified-id=\"Negative-binomial-distributions-316\"><span class=\"toc-item-num\">3.1.6&nbsp;&nbsp;</span>Negative binomial distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Gaussian-distributions\" data-toc-modified-id=\"Gaussian-distributions-317\"><span class=\"toc-item-num\">3.1.7&nbsp;&nbsp;</span>Gaussian distributions</a></div><div class=\"lev2 toc-item\"><a href=\"#Generic-KL-UCB-indexes,-with-a-bisection-search\" data-toc-modified-id=\"Generic-KL-UCB-indexes,-with-a-bisection-search-32\"><span class=\"toc-item-num\">3.2&nbsp;&nbsp;</span>Generic KL-UCB indexes, with a bisection search</a></div><div class=\"lev2 toc-item\"><a href=\"#Distribution-specific-KL-UCB-indexes\" data-toc-modified-id=\"Distribution-specific-KL-UCB-indexes-33\"><span class=\"toc-item-num\">3.3&nbsp;&nbsp;</span>Distribution-specific KL-UCB indexes</a></div><div class=\"lev3 toc-item\"><a href=\"#Gaussian\" data-toc-modified-id=\"Gaussian-331\"><span class=\"toc-item-num\">3.3.1&nbsp;&nbsp;</span>Gaussian</a></div><div class=\"lev3 toc-item\"><a href=\"#Bernoulli\" data-toc-modified-id=\"Bernoulli-332\"><span class=\"toc-item-num\">3.3.2&nbsp;&nbsp;</span>Bernoulli</a></div><div class=\"lev3 toc-item\"><a href=\"#Poisson\" data-toc-modified-id=\"Poisson-333\"><span class=\"toc-item-num\">3.3.3&nbsp;&nbsp;</span>Poisson</a></div><div class=\"lev3 toc-item\"><a href=\"#Exponential\" data-toc-modified-id=\"Exponential-334\"><span class=\"toc-item-num\">3.3.4&nbsp;&nbsp;</span>Exponential</a></div><div class=\"lev1 toc-item\"><a href=\"#With-Cython\" data-toc-modified-id=\"With-Cython-4\"><span class=\"toc-item-num\">4&nbsp;&nbsp;</span>With Cython</a></div><div class=\"lev2 toc-item\"><a href=\"#KL-divergences\" data-toc-modified-id=\"KL-divergences-41\"><span class=\"toc-item-num\">4.1&nbsp;&nbsp;</span>KL divergences</a></div><div class=\"lev3 toc-item\"><a href=\"#Bernoulli-distributions\" data-toc-modified-id=\"Bernoulli-distributions-411\"><span class=\"toc-item-num\">4.1.1&nbsp;&nbsp;</span>Bernoulli distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Binomial-distributions\" data-toc-modified-id=\"Binomial-distributions-412\"><span class=\"toc-item-num\">4.1.2&nbsp;&nbsp;</span>Binomial distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Poisson-distributions\" data-toc-modified-id=\"Poisson-distributions-413\"><span class=\"toc-item-num\">4.1.3&nbsp;&nbsp;</span>Poisson distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Exponential-distributions\" data-toc-modified-id=\"Exponential-distributions-414\"><span class=\"toc-item-num\">4.1.4&nbsp;&nbsp;</span>Exponential distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Gamma-distributions\" data-toc-modified-id=\"Gamma-distributions-415\"><span class=\"toc-item-num\">4.1.5&nbsp;&nbsp;</span>Gamma distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Negative-binomial-distributions\" data-toc-modified-id=\"Negative-binomial-distributions-416\"><span class=\"toc-item-num\">4.1.6&nbsp;&nbsp;</span>Negative binomial distributions</a></div><div class=\"lev3 toc-item\"><a href=\"#Gaussian-distributions\" data-toc-modified-id=\"Gaussian-distributions-417\"><span class=\"toc-item-num\">4.1.7&nbsp;&nbsp;</span>Gaussian distributions</a></div><div class=\"lev2 toc-item\"><a href=\"#Generic-KL-UCB-indexes,-with-a-bisection-search\" data-toc-modified-id=\"Generic-KL-UCB-indexes,-with-a-bisection-search-42\"><span class=\"toc-item-num\">4.2&nbsp;&nbsp;</span>Generic KL-UCB indexes, with a bisection search</a></div><div class=\"lev1 toc-item\"><a href=\"#With-the-C-API-for-Python\" data-toc-modified-id=\"With-the-C-API-for-Python-5\"><span class=\"toc-item-num\">5&nbsp;&nbsp;</span>With the C API for Python</a></div><div class=\"lev1 toc-item\"><a href=\"#Tests-and-benchmarks\" data-toc-modified-id=\"Tests-and-benchmarks-6\"><span class=\"toc-item-num\">6&nbsp;&nbsp;</span>Tests and benchmarks</a></div><div class=\"lev2 toc-item\"><a href=\"#KL-divergences\" data-toc-modified-id=\"KL-divergences-61\"><span class=\"toc-item-num\">6.1&nbsp;&nbsp;</span>KL divergences</a></div><div class=\"lev3 toc-item\"><a href=\"#Bernoulli\" data-toc-modified-id=\"Bernoulli-611\"><span class=\"toc-item-num\">6.1.1&nbsp;&nbsp;</span>Bernoulli</a></div><div class=\"lev3 toc-item\"><a href=\"#Binomial\" data-toc-modified-id=\"Binomial-612\"><span class=\"toc-item-num\">6.1.2&nbsp;&nbsp;</span>Binomial</a></div><div class=\"lev3 toc-item\"><a href=\"#Poisson\" data-toc-modified-id=\"Poisson-613\"><span class=\"toc-item-num\">6.1.3&nbsp;&nbsp;</span>Poisson</a></div><div class=\"lev3 toc-item\"><a href=\"#Exponential\" data-toc-modified-id=\"Exponential-614\"><span class=\"toc-item-num\">6.1.4&nbsp;&nbsp;</span>Exponential</a></div><div class=\"lev3 toc-item\"><a href=\"#Gamma\" data-toc-modified-id=\"Gamma-615\"><span class=\"toc-item-num\">6.1.5&nbsp;&nbsp;</span>Gamma</a></div><div class=\"lev3 toc-item\"><a href=\"#Negative-binomial\" data-toc-modified-id=\"Negative-binomial-616\"><span class=\"toc-item-num\">6.1.6&nbsp;&nbsp;</span>Negative binomial</a></div><div class=\"lev3 toc-item\"><a href=\"#Gaussian\" data-toc-modified-id=\"Gaussian-617\"><span class=\"toc-item-num\">6.1.7&nbsp;&nbsp;</span>Gaussian</a></div><div class=\"lev2 toc-item\"><a href=\"#KL-UCB-indexes\" data-toc-modified-id=\"KL-UCB-indexes-62\"><span class=\"toc-item-num\">6.2&nbsp;&nbsp;</span>KL-UCB indexes</a></div><div class=\"lev3 toc-item\"><a href=\"#Gaussian\" data-toc-modified-id=\"Gaussian-621\"><span class=\"toc-item-num\">6.2.1&nbsp;&nbsp;</span>Gaussian</a></div><div class=\"lev3 toc-item\"><a href=\"#Bernoulli\" data-toc-modified-id=\"Bernoulli-622\"><span class=\"toc-item-num\">6.2.2&nbsp;&nbsp;</span>Bernoulli</a></div><div class=\"lev3 toc-item\"><a href=\"#Poisson\" data-toc-modified-id=\"Poisson-623\"><span class=\"toc-item-num\">6.2.3&nbsp;&nbsp;</span>Poisson</a></div><div class=\"lev3 toc-item\"><a href=\"#Exponential\" data-toc-modified-id=\"Exponential-624\"><span class=\"toc-item-num\">6.2.4&nbsp;&nbsp;</span>Exponential</a></div><div class=\"lev2 toc-item\"><a href=\"#Clean-up\" data-toc-modified-id=\"Clean-up-63\"><span class=\"toc-item-num\">6.3&nbsp;&nbsp;</span>Clean up</a></div><div class=\"lev1 toc-item\"><a href=\"#Conclusion\" data-toc-modified-id=\"Conclusion-7\"><span class=\"toc-item-num\">7&nbsp;&nbsp;</span>Conclusion</a></div><div class=\"lev2 toc-item\"><a href=\"#Take-away-messages\" data-toc-modified-id=\"Take-away-messages-71\"><span class=\"toc-item-num\">7.1&nbsp;&nbsp;</span>Take away messages</a></div><div class=\"lev2 toc-item\"><a href=\"#Using-Cython-for-real-?\" data-toc-modified-id=\"Using-Cython-for-real-?-72\"><span class=\"toc-item-num\">7.2&nbsp;&nbsp;</span>Using Cython <em>for real</em> ?</a></div><div class=\"lev2 toc-item\"><a href=\"#Using-C-for-real-?\" data-toc-modified-id=\"Using-C-for-real-?-73\"><span class=\"toc-item-num\">7.3&nbsp;&nbsp;</span>Using C <em>for real</em> ?</a></div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----\n",
    "# Introduction\n",
    "\n",
    "In this small notebook, I implement various [Kullback-Leibler divergence functions](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence), in [Python](https://www.python.org/), using different approaches: naive Python, and using Numba and Cython.\n",
    "\n",
    "I also implement KL-UCB indexes, in the three approaches, and finally I present some basic benchmarks to compare the time and memory efficiency of the different approaches, for each function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Requirements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Lilian Besson (Naereen) \n",
      "\n",
      "CPython 3.6.3\n",
      "IPython 6.3.1\n",
      "\n",
      "numpy 1.14.2\n",
      "numba 0.37.0\n",
      "\n",
      "compiler   : GCC 7.2.0\n",
      "system     : Linux\n",
      "release    : 4.13.0-38-generic\n",
      "machine    : x86_64\n",
      "processor  : x86_64\n",
      "CPU cores  : 4\n",
      "interpreter: 64bit\n",
      "Git hash   : f56a7701503fe66afca4dcc7ed96827d462f773d\n"
     ]
    }
   ],
   "source": [
    "%load_ext watermark\n",
    "%watermark -v -m -a \"Lilian Besson (Naereen)\" -p numpy,numba -g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----\n",
    "# KL divergences and KL-UCB indexes, in naive Python\n",
    "\n",
    "I will copy and paste parts of [this file](https://github.com/SMPyBandits/SMPyBandits/blob/master/SMPyBandits/Policies/kullback.py) from my [SMPyBandits](https://github.com/SMPyBandits/SMPyBandits/) library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I will include docstrings and examples only for the naive implementation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## KL divergences"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bernoulli distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def klBern(x, y):\n",
    "    r\"\"\" Kullback-Leibler divergence for Bernoulli distributions. https://en.wikipedia.org/wiki/Bernoulli_distribution#Kullback.E2.80.93Leibler_divergence\n",
    "\n",
    "    .. math:: \\mathrm{KL}(\\mathcal{B}(x), \\mathcal{B}(y)) = x \\log(\\frac{x}{y}) + (1-x) \\log(\\frac{1-x}{1-y}).\"\"\"\n",
    "    x = min(max(x, eps), 1 - eps)\n",
    "    y = min(max(y, eps), 1 - eps)\n",
    "    return x * np.log(x / y) + (1 - x) * np.log((1 - x) / (1 - y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "1.7577796618689758"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "1.7577796618689758"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.020135513550688863"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "4.503217453131898"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "34.53957599234081"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "klBern(0.5, 0.5)\n",
    "klBern(0.1, 0.9)\n",
    "klBern(0.9, 0.1)\n",
    "klBern(0.4, 0.5)\n",
    "klBern(0.01, 0.99)\n",
    "klBern(0, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Binomial distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def klBin(x, y, n):\n",
    "    r\"\"\" Kullback-Leibler divergence for Binomial distributions. https://math.stackexchange.com/questions/320399/kullback-leibner-divergence-of-binomial-distributions\n",
    "\n",
    "    - It is simply the n times :func:`klBern` on x and y.\n",
    "\n",
    "    .. math:: \\mathrm{KL}(\\mathrm{Bin}(x, n), \\mathrm{Bin}(y, n)) = n \\times \\left(x \\log(\\frac{x}{y}) + (1-x) \\log(\\frac{1-x}{1-y}) \\right).\n",
    "\n",
    "    .. warning:: The two distributions must have the same parameter n, and x, y are p, q in (0, 1).\n",
    "    \"\"\"\n",
    "    x = min(max(x, eps), 1 - eps)\n",
    "    y = min(max(y, eps), 1 - eps)\n",
    "    return n * (x * np.log(x / y) + (1 - x) * np.log((1 - x) / (1 - y)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "17.57779661868976"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "17.57779661868976"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.20135513550688863"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "45.03217453131897"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "345.3957599234081"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "klBin(0.5, 0.5, 10)\n",
    "klBin(0.1, 0.9, 10)\n",
    "klBin(0.9, 0.1, 10)\n",
    "klBin(0.4, 0.5, 10)\n",
    "klBin(0.01, 0.99, 10)\n",
    "klBin(0, 1, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Poisson distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def klPoisson(x, y):\n",
    "    r\"\"\" Kullback-Leibler divergence for Poison distributions. https://en.wikipedia.org/wiki/Poisson_distribution#Kullback.E2.80.93Leibler_divergence\n",
    "\n",
    "    .. math:: \\mathrm{KL}(\\mathrm{Poisson}(x), \\mathrm{Poisson}(y)) = y - x + x \\times \\log(\\frac{x}{y}).\n",
    "    \"\"\"\n",
    "    x = max(x, eps)\n",
    "    y = max(y, eps)\n",
    "    return y - x + x * np.log(x / y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.3862943611198906"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.3068528194400547"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.9205584583201643"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.2739075652893146"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "33.538776394910684"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "klPoisson(3, 3)\n",
    "klPoisson(2, 1)\n",
    "klPoisson(1, 2)\n",
    "klPoisson(3, 6)\n",
    "klPoisson(6, 8)\n",
    "klPoisson(1, 0)\n",
    "klPoisson(0, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exponential distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def klExp(x, y):\n",
    "    r\"\"\" Kullback-Leibler divergence for exponential distributions. https://en.wikipedia.org/wiki/Exponential_distribution#Kullback.E2.80.93Leibler_divergence\n",
    "\n",
    "    .. math::\n",
    "\n",
    "        \\mathrm{KL}(\\mathrm{Exp}(x), \\mathrm{Exp}(y)) = \\begin{cases}\n",
    "        \\frac{x}{y} - 1 - \\log(\\frac{x}{y}) & \\text{if} x > 0, y > 0\\\\\n",
    "        +\\infty & \\text{otherwise}\n",
    "        \\end{cases}\n",
    "    \"\"\"\n",
    "    if x <= 0 or y <= 0:\n",
    "        return float('+inf')\n",
    "    else:\n",
    "        x = max(x, eps)\n",
    "        y = max(y, eps)\n",
    "        return x / y - 1 - np.log(x / y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.1931471805599453"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.1931471805599453"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.3068528194400547"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.3068528194400547"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.0376820724517809"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "inf"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "inf"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "inf"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "klExp(3, 3)\n",
    "klExp(3, 6)\n",
    "klExp(1, 2)\n",
    "klExp(2, 1)\n",
    "klExp(4, 2)\n",
    "klExp(6, 8)\n",
    "klExp(-3, 2)\n",
    "klExp(3, -2)\n",
    "klExp(-3, -2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gamma distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def klGamma(x, y, a=1):\n",
    "    r\"\"\" Kullback-Leibler divergence for gamma distributions. https://en.wikipedia.org/wiki/Gamma_distribution#Kullback.E2.80.93Leibler_divergence\n",
    "\n",
    "    - It is simply the a times :func:`klExp` on x and y.\n",
    "\n",
    "    .. math::\n",
    "\n",
    "        \\mathrm{KL}(\\Gamma(x, a), \\Gamma(y, a)) = \\begin{cases}\n",
    "        a \\times \\left( \\frac{x}{y} - 1 - \\log(\\frac{x}{y}) \\right) & \\text{if} x > 0, y > 0\\\\\n",
    "        +\\infty & \\text{otherwise}\n",
    "        \\end{cases}\n",
    "\n",
    "    .. warning:: The two distributions must have the same parameter a.\n",
    "    \"\"\"\n",
    "    if x <= 0 or y <= 0:\n",
    "        return float('+inf')\n",
    "    else:\n",
    "        x = max(x, eps)\n",
    "        y = max(y, eps)\n",
    "        return a * (x / y - 1 - np.log(x / y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.1931471805599453"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.1931471805599453"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.3068528194400547"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.3068528194400547"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.0376820724517809"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "inf"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "inf"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "inf"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "klGamma(3, 3)\n",
    "klGamma(3, 6)\n",
    "klGamma(1, 2)\n",
    "klGamma(2, 1)\n",
    "klGamma(4, 2)\n",
    "klGamma(6, 8)\n",
    "klGamma(-3, 2)\n",
    "klGamma(3, -2)\n",
    "klGamma(-3, -2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Negative binomial distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def klNegBin(x, y, r=1):\n",
    "    r\"\"\" Kullback-Leibler divergence for negative binomial distributions. https://en.wikipedia.org/wiki/Negative_binomial_distribution\n",
    "\n",
    "    .. math:: \\mathrm{KL}(\\mathrm{NegBin}(x, r), \\mathrm{NegBin}(y, r)) = r \\times \\log((r + x) / (r + y)) - x \\times \\log(y \\times (r + x) / (x \\times (r + y))).\n",
    "\n",
    "    .. warning:: The two distributions must have the same parameter r.\n",
    "    \"\"\"\n",
    "    x = max(x, eps)\n",
    "    y = max(y, eps)\n",
    "    return r * np.log((r + x) / (r + y)) - x * np.log(y * (r + x) / (x * (r + y)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "-0.7116117934648849"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "2.0321564902394043"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "-0.13065314341785483"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "-0.7173536633057466"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "34.53957599234081"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "-0.8329919030334189"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "-0.9148905602182661"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "2.332552851091954"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "-0.15457261175809217"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "-0.8362571425112515"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "klNegBin(0.5, 0.5)\n",
    "klNegBin(0.1, 0.9)\n",
    "klNegBin(0.9, 0.1)\n",
    "klNegBin(0.4, 0.5)\n",
    "klNegBin(0.01, 0.99)\n",
    "klBern(0, 1)\n",
    "klNegBin(0.5, 0.5, r=2)\n",
    "klNegBin(0.1, 0.9, r=2)\n",
    "klNegBin(0.1, 0.9, r=4)\n",
    "klNegBin(0.9, 0.1, r=2)\n",
    "klNegBin(0.4, 0.5, r=2)\n",
    "klNegBin(0.01, 0.99, r=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gaussian distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def klGauss(x, y, sig2x=0.25, sig2y=None):\n",
    "    r\"\"\" Kullback-Leibler divergence for Gaussian distributions of means ``x`` and ``y`` and variances ``sig2x`` and ``sig2y``, :math:`\\nu_1 = \\mathcal{N}(x, \\sigma_x^2)` and :math:`\\nu_2 = \\mathcal{N}(y, \\sigma_x^2)`:\n",
    "\n",
    "    .. math:: \\mathrm{KL}(\\nu_1, \\nu_2) = \\frac{(x - y)^2}{2 \\sigma_y^2} + \\frac{1}{2}\\left( \\frac{\\sigma_x^2}{\\sigma_y^2} - 1 \\log\\left(\\frac{\\sigma_x^2}{\\sigma_y^2}\\right) \\right).\n",
    "\n",
    "    See https://en.wikipedia.org/wiki/Normal_distribution#Other_properties\n",
    "\n",
    "    - By default, sig2y is assumed to be sig2x (same variance).\n",
    "    \"\"\"\n",
    "    if sig2y is None or - eps < (sig2y - sig2x) < eps:\n",
    "        return (x - y) ** 2 / (2. * sig2x)\n",
    "    else:\n",
    "        return (x - y) ** 2 / (2. * sig2y) + 0.5 * ((sig2x/sig2y)**2 - 1 - np.log(sig2x/sig2y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "18.0"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "2.0"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "2.0"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "8.0"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "8.0"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "50.0"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "50.0"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "2.0"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "2.0"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.45"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.05"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.05"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.2"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.2"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "-0.028426409720027357"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.2243971805599453"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "1.1534264097200273"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.9715735902799727"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.7243971805599453"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "3.1534264097200273"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.9715735902799727"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.7243971805599453"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "3.1534264097200273"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "klGauss(3, 3)\n",
    "klGauss(3, 6)\n",
    "klGauss(1, 2)\n",
    "klGauss(2, 1)\n",
    "klGauss(4, 2)\n",
    "klGauss(6, 8)\n",
    "klGauss(-3, 2)\n",
    "klGauss(3, -2)\n",
    "klGauss(-3, -2)\n",
    "klGauss(3, 2)\n",
    "klGauss(3, 3, sig2x=10)\n",
    "klGauss(3, 6, sig2x=10)\n",
    "klGauss(1, 2, sig2x=10)\n",
    "klGauss(2, 1, sig2x=10)\n",
    "klGauss(4, 2, sig2x=10)\n",
    "klGauss(6, 8, sig2x=10)\n",
    "klGauss(0, 0, sig2x=0.25, sig2y=0.5)\n",
    "klGauss(0, 0, sig2x=0.25, sig2y=1.0)\n",
    "klGauss(0, 0, sig2x=0.5, sig2y=0.25)\n",
    "klGauss(0, 1, sig2x=0.25, sig2y=0.5)\n",
    "klGauss(0, 1, sig2x=0.25, sig2y=1.0)\n",
    "klGauss(0, 1, sig2x=0.5, sig2y=0.25)\n",
    "klGauss(1, 0, sig2x=0.25, sig2y=0.5)\n",
    "klGauss(1, 0, sig2x=0.25, sig2y=1.0)\n",
    "klGauss(1, 0, sig2x=0.5, sig2y=0.25)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generic KL-UCB indexes, with a bisection search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def klucb(x, d, kl, upperbound, lowerbound=float('-inf'), precision=1e-6, max_iterations=50):\n",
    "    \"\"\" The generic KL-UCB index computation.\n",
    "\n",
    "    - x: value of the cum reward,\n",
    "    - d: upper bound on the divergence,\n",
    "    - kl: the KL divergence to be used (:func:`klBern`, :func:`klGauss`, etc),\n",
    "    - upperbound, lowerbound=float('-inf'): the known bound of the values x,\n",
    "    - precision=1e-6: the threshold from where to stop the research,\n",
    "    - max_iterations: max number of iterations of the loop (safer to bound it to reduce time complexity).\n",
    "\n",
    "    .. note:: It uses a **bisection search**, and one call to ``kl`` for each step of the bisection search.\n",
    "    \"\"\"\n",
    "    value = max(x, lowerbound)\n",
    "    u = upperbound\n",
    "    _count_iteration = 0\n",
    "    while _count_iteration < max_iterations and u - value > precision:\n",
    "        _count_iteration += 1\n",
    "        m = (value + u) / 2.\n",
    "        if kl(x, m) > d:\n",
    "            u = m\n",
    "        else:\n",
    "            value = m\n",
    "    return (value + u) / 2."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, for `klucbBern`, the two steps are to first compute an upperbound (as precise as possible) and the compute the kl-UCB index:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.994140625"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.9944824218750001"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.994140625"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.9944896697998048"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x, d = 0.9, 0.2\n",
    "upperbound = 1\n",
    "klucb(x, d, klBern, upperbound, lowerbound=0, precision=1e-3, max_iterations=10)\n",
    "klucb(x, d, klBern, upperbound, lowerbound=0, precision=1e-6, max_iterations=10)\n",
    "klucb(x, d, klBern, upperbound, lowerbound=0, precision=1e-3, max_iterations=50)\n",
    "klucb(x, d, klBern, upperbound, lowerbound=0, precision=1e-6, max_iterations=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distribution-specific KL-UCB indexes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gaussian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def klucbGauss(x, d, sig2x=0.25, precision=0.):\n",
    "    \"\"\" KL-UCB index computation for Gaussian distributions.\n",
    "\n",
    "    - Note that it does not require any search.\n",
    "\n",
    "    .. warning:: it works only if the good variance constant is given.\n",
    "    \"\"\"\n",
    "    return x + np.sqrt(2 * sig2x * d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.416227766016838"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.816227766016838"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "1.216227766016838"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.547213595499958"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.7708203932499369"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.9472135954999579"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "1.170820393249937"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "1.347213595499958"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "1.570820393249937"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "klucbGauss(0.1, 0.2)\n",
    "klucbGauss(0.5, 0.2)\n",
    "klucbGauss(0.9, 0.2)\n",
    "klucbGauss(0.1, 0.4)\n",
    "klucbGauss(0.1, 0.9)\n",
    "klucbGauss(0.5, 0.4)\n",
    "klucbGauss(0.5, 0.9)\n",
    "klucbGauss(0.9, 0.4)\n",
    "klucbGauss(0.9, 0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bernoulli"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def klucbBern(x, d, precision=1e-6):\n",
    "    \"\"\" KL-UCB index computation for Bernoulli distributions, using :func:`klucb`.\"\"\"\n",
    "    upperbound = min(1., klucbGauss(x, d, sig2x=0.25))  # variance 1/4 for [0,1] bounded distributions\n",
    "    # upperbound = min(1., klucbPoisson(x, d))  # also safe, and better ?\n",
    "    return klucb(x, d, klBern, upperbound, precision)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.37839145109809247"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.7870889692292777"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.9944896697998048"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.5194755673450786"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.7347148310932183"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.871035844022684"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.9568095207214355"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.9992855072021485"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.9999950408935546"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "klucbBern(0.1, 0.2)\n",
    "klucbBern(0.5, 0.2)\n",
    "klucbBern(0.9, 0.2)\n",
    "klucbBern(0.1, 0.4)\n",
    "klucbBern(0.1, 0.9)\n",
    "klucbBern(0.5, 0.4)\n",
    "klucbBern(0.5, 0.9)\n",
    "klucbBern(0.9, 0.4)\n",
    "klucbBern(0.9, 0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Poisson"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def klucbPoisson(x, d, precision=1e-6):\n",
    "    \"\"\" KL-UCB index computation for Poisson distributions, using :func:`klucb`.\"\"\"\n",
    "    upperbound = x + d + np.sqrt(d * d + 2 * x * d)  # looks safe, to check: left (Gaussian) tail of Poisson dev\n",
    "    return klucb(x, d, klPoisson, upperbound, precision)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.45052392780119604"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "1.0893765430263218"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "1.6401128559741487"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.6936844019642616"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "1.2527967047658155"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "1.4229339603816749"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "2.122985165630671"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "2.033691887156203"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "2.8315738094979777"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "klucbPoisson(0.1, 0.2)\n",
    "klucbPoisson(0.5, 0.2)\n",
    "klucbPoisson(0.9, 0.2)\n",
    "klucbPoisson(0.1, 0.4)\n",
    "klucbPoisson(0.1, 0.9)\n",
    "klucbPoisson(0.5, 0.4)\n",
    "klucbPoisson(0.5, 0.9)\n",
    "klucbPoisson(0.9, 0.4)\n",
    "klucbPoisson(0.9, 0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exponential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "def klucbExp(x, d, precision=1e-6):\n",
    "    \"\"\" KL-UCB index computation for exponential distributions, using :func:`klucb`.\"\"\"\n",
    "    if d < 0.77:  # XXX where does this value come from?\n",
    "        upperbound = x / (1 + 2. / 3 * d - np.sqrt(4. / 9 * d * d + 2 * d))\n",
    "        # safe, klexp(x,y) >= e^2/(2*(1-2e/3)) if x=y(1-e)\n",
    "    else:\n",
    "        upperbound = x * np.exp(d + 1)\n",
    "    if d > 1.61:  # XXX where does this value come from?\n",
    "        lowerbound = x * np.exp(d)\n",
    "    else:\n",
    "        lowerbound = x / (1 + d - np.sqrt(d * d + 2 * d))\n",
    "    return klucb(x, d, klGamma, upperbound, lowerbound, precision)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.20274118449172676"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "1.013706285168157"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "1.8246716397412546"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.2857928251730546"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.5590884945251575"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "1.428962647183463"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "2.7954420946912126"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "2.572132498767508"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "5.031795430303065"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "klucbExp(0.1, 0.2)\n",
    "klucbExp(0.5, 0.2)\n",
    "klucbExp(0.9, 0.2)\n",
    "klucbExp(0.1, 0.4)\n",
    "klucbExp(0.1, 0.9)\n",
    "klucbExp(0.5, 0.4)\n",
    "klucbExp(0.5, 0.9)\n",
    "klucbExp(0.9, 0.4)\n",
    "klucbExp(0.9, 0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Others\n",
    "We could do the same for more distributions, but that's enough."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----\n",
    "# With Numba\n",
    "\n",
    "It will be *exactly* the same code as above, except that the [`numba.jit`](http://numba.pydata.org/numba-doc/latest/user/jit.html) decorator will be used for each functions, to let [numba](http://numba.pydata.org/) *try* to speed up the code!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import jit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As much as possible, one should call `@jit(nopython=True)` to be sure that numba does not fall back silently to naive Python code. With `nopython=True`, any call to the generated function will fail if the compilation could not succeed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## KL divergences"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bernoulli distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def klBern_numba(x, y):\n",
    "    x = min(max(x, eps), 1 - eps)\n",
    "    y = min(max(y, eps), 1 - eps)\n",
    "    return x * np.log(x / y) + (1 - x) * np.log((1 - x) / (1 - y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Binomial distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def klBin_numba(x, y, n):\n",
    "    x = min(max(x, eps), 1 - eps)\n",
    "    y = min(max(y, eps), 1 - eps)\n",
    "    return n * (x * np.log(x / y) + (1 - x) * np.log((1 - x) / (1 - y)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Poisson distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def klPoisson_numba(x, y):\n",
    "    x = max(x, eps)\n",
    "    y = max(y, eps)\n",
    "    return y - x + x * np.log(x / y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exponential distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def klExp_numba(x, y):\n",
    "    if x <= 0 or y <= 0:\n",
    "        return inf\n",
    "    else:\n",
    "        x = max(x, eps)\n",
    "        y = max(y, eps)\n",
    "        return x / y - 1 - np.log(x / y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gamma distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def klGamma_numba(x, y, a=1):\n",
    "    if x <= 0 or y <= 0:\n",
    "        return inf\n",
    "    else:\n",
    "        x = max(x, eps)\n",
    "        y = max(y, eps)\n",
    "        return a * (x / y - 1 - np.log(x / y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Negative binomial distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def klNegBin_numba(x, y, r=1):\n",
    "    x = max(x, eps)\n",
    "    y = max(y, eps)\n",
    "    return r * np.log((r + x) / (r + y)) - x * np.log(y * (r + x) / (x * (r + y)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gaussian distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def klGauss_numba(x, y, sig2x=0.25, sig2y=0.25):\n",
    "    if - eps < (sig2y - sig2x) and (sig2y - sig2x) < eps:\n",
    "        return (x - y) ** 2 / (2. * sig2x)\n",
    "    else:\n",
    "        return (x - y) ** 2 / (2. * sig2y) + 0.5 * ((sig2x/sig2y)**2 - 1 - np.log(sig2x/sig2y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generic KL-UCB indexes, with a bisection search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit\n",
    "def klucb_numba(x, d, kl, upperbound,\n",
    "                lowerbound=float('-inf'), precision=1e-6, max_iterations=50):\n",
    "    value = max(x, lowerbound)\n",
    "    u = upperbound\n",
    "    _count_iteration = 0\n",
    "    while _count_iteration < max_iterations and u - value > precision:\n",
    "        _count_iteration += 1\n",
    "        m = (value + u) / 2.\n",
    "        if kl(x, m) > d:\n",
    "            u = m\n",
    "        else:\n",
    "            value = m\n",
    "    return (value + u) / 2."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, for `klucbBern`, the two steps are to first compute an upperbound (as precise as possible) and the compute the kl-UCB index:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.994140625"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.9944824218750001"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.994140625"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.9944896697998048"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x, d = 0.9, 0.2\n",
    "upperbound = 1\n",
    "klucb_numba(x, d, klBern_numba, upperbound, lowerbound=0, precision=1e-3, max_iterations=10)\n",
    "klucb_numba(x, d, klBern_numba, upperbound, lowerbound=0, precision=1e-6, max_iterations=10)\n",
    "klucb_numba(x, d, klBern_numba, upperbound, lowerbound=0, precision=1e-3, max_iterations=50)\n",
    "klucb_numba(x, d, klBern_numba, upperbound, lowerbound=0, precision=1e-6, max_iterations=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distribution-specific KL-UCB indexes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gaussian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def klucbGauss_numba(x, d, sig2x=0.25, precision=0.):\n",
    "    return x + np.sqrt(2 * sig2x * d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bernoulli\n",
    "\n",
    "Here, the `nopython=True` fails as numba has a hard time typing linked function calls."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit\n",
    "def klucbBern_numba(x, d, precision=1e-6):\n",
    "    upperbound = min(1., klucbGauss_numba(x, d, sig2x=0.25))  # variance 1/4 for [0,1] bounded distributions\n",
    "    # upperbound = min(1., klucbPoisson(x, d))  # also safe, and better ?\n",
    "    return klucb_numba(x, d, klBern_numba, upperbound, precision)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Poisson"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit\n",
    "def klucbPoisson_numba(x, d, precision=1e-6):\n",
    "    upperbound = x + d + np.sqrt(d * d + 2 * x * d)  # looks safe, to check: left (Gaussian) tail of Poisson dev\n",
    "    return klucb_numba(x, d, klPoisson_numba, upperbound, precision)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exponential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def klucbExp_numba(x, d, precision=1e-6):\n",
    "    if d < 0.77:  # XXX where does this value come from?\n",
    "        upperbound = x / (1 + 2. / 3 * d - np.sqrt(4. / 9 * d * d + 2 * d))\n",
    "        # safe, klexp(x,y) >= e^2/(2*(1-2e/3)) if x=y(1-e)\n",
    "    else:\n",
    "        upperbound = x * np.exp(d + 1)\n",
    "    if d > 1.61:  # XXX where does this value come from?\n",
    "        lowerbound = x * np.exp(d)\n",
    "    else:\n",
    "        lowerbound = x / (1 + d - np.sqrt(d * d + 2 * d))\n",
    "    return klucb_numba(x, d, klGamma_numba, upperbound, lowerbound, precision)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----\n",
    "# With Cython\n",
    "\n",
    "It will be *almost* exactly the same code, by using the [`cython`]() magic to have cells written in [Cython](http://cython.org/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext cython"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A cell can now be written in Cython.\n",
    "For instance, we can define a simple example function in Python, and then write a Cython version, simply by declaring variables and tagging their types, like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "def some_loop(n: int) -> int:\n",
    "    s = 0\n",
    "    for i in range(0, n, 2):\n",
    "        s += i\n",
    "    return s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython\n",
    "def some_loop_cython(int n) -> int:\n",
    "    cdef int s = 0\n",
    "    cdef int i = 0\n",
    "    for i in range(0, n, 2):\n",
    "        s += i\n",
    "    return s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.72 µs ± 106 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n",
      "13.1 µs ± 1.25 µs per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n",
      "2.14 µs ± 209 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit np.random.randint(1000)\n",
    "%timeit some_loop(np.random.randint(1000))\n",
    "%timeit some_loop_cython(np.random.randint(1000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we observe a large speed-up. But how large? $6$ times or $50$ times?\n",
    "\n",
    "It's really important to include the time taken by the Pseudo-Random Number Generator:\n",
    "\n",
    "- Wrong computation of the speed-up gives about $6$ times faster:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6.606334841628959"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "14.6 / 2.21"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- But if we remove the time taken by the PRNG (which takes the same time for both the naive Python and the Cython function), we get a larger speed-up, closer to reality, about $50$ times and not just $6$ times faster!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "48.65384615384615"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(14.6 - 1.95) / (2.21 - 1.95)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## KL divergences"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bernoulli distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython\n",
    "from libc.math cimport log\n",
    "eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]\n",
    "\n",
    "def klBern_cython(float x, float y) -> float:\n",
    "    x = min(max(x, eps), 1 - eps)\n",
    "    y = min(max(y, eps), 1 - eps)\n",
    "    return x * log(x / y) + (1 - x) * log((1 - x) / (1 - y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Binomial distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython\n",
    "from libc.math cimport log\n",
    "eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]\n",
    "\n",
    "def klBin_cython(float x, float y, int n) -> float:\n",
    "    x = min(max(x, eps), 1 - eps)\n",
    "    y = min(max(y, eps), 1 - eps)\n",
    "    return n * (x * log(x / y) + (1 - x) * log((1 - x) / (1 - y)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Poisson distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython\n",
    "from libc.math cimport log\n",
    "eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]\n",
    "\n",
    "def klPoisson_cython(float x, float y) -> float:\n",
    "    x = max(x, eps)\n",
    "    y = max(y, eps)\n",
    "    return y - x + x * log(x / y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exponential distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython\n",
    "from libc.math cimport log\n",
    "eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]\n",
    "\n",
    "def klExp_cython(float x, float y) -> float:\n",
    "    if x <= 0 or y <= 0:\n",
    "        return float('+inf')\n",
    "    else:\n",
    "        x = max(x, eps)\n",
    "        y = max(y, eps)\n",
    "        return x / y - 1 - log(x / y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gamma distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython\n",
    "from libc.math cimport log\n",
    "eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]\n",
    "\n",
    "def klGamma_cython(float x, float y, float a=1) -> float:\n",
    "    if x <= 0 or y <= 0:\n",
    "        return float('+inf')\n",
    "    else:\n",
    "        x = max(x, eps)\n",
    "        y = max(y, eps)\n",
    "        return a * (x / y - 1 - log(x / y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Negative binomial distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython\n",
    "from libc.math cimport log\n",
    "eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]\n",
    "\n",
    "def klNegBin_cython(float x, float y, float r=1) -> float:\n",
    "    x = max(x, eps)\n",
    "    y = max(y, eps)\n",
    "    return r * log((r + x) / (r + y)) - x * log(y * (r + x) / (x * (r + y)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gaussian distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython\n",
    "from libc.math cimport log\n",
    "eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]\n",
    "\n",
    "def klGauss_cython(float x, float y, float sig2x=0.25, float sig2y=0.25) -> float:\n",
    "    if - eps < (sig2y - sig2x) < eps:\n",
    "        return (x - y) ** 2 / (2. * sig2x)\n",
    "    else:\n",
    "        return (x - y) ** 2 / (2. * sig2y) + 0.5 * ((sig2x/sig2y)**2 - 1 - log(sig2x/sig2y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generic KL-UCB indexes, with a bisection search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For these, they need previously defined functions, which have to be rewritten from inside the `cython` cell to be accessible from Cython.\n",
    "To minimize repetitions, I use only one cell to define all functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython\n",
    "from libc.math cimport sqrt, log, exp\n",
    "eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]\n",
    "\n",
    "\n",
    "def klucbGauss_cython(float x, float d, float sig2x=0.25, float precision=0.) -> float:\n",
    "    return x + sqrt(2 * sig2x * d)\n",
    "\n",
    "cdef float klucbGauss_cython_x(float x, float d, float sig2x=0.25, float precision=0.):\n",
    "    return x + sqrt(2 * sig2x * d)\n",
    "\n",
    "\n",
    "def klucb_cython(float x, float d, kl, float upperbound,\n",
    "                 float lowerbound=float('-inf'),\n",
    "                 float precision=1e-6, int max_iterations=50) -> float:\n",
    "    cdef float value = max(x, lowerbound)\n",
    "    cdef float u = upperbound\n",
    "    cdef int _count_iteration = 0\n",
    "    cdef float m = 0\n",
    "    while _count_iteration < max_iterations and u - value > precision:\n",
    "        _count_iteration += 1\n",
    "        m = (value + u) / 2.\n",
    "        if kl(x, m) > d:\n",
    "            u = m\n",
    "        else:\n",
    "            value = m\n",
    "    return (value + u) / 2.\n",
    "\n",
    "\n",
    "cdef float klBern_cython_x(float x, float y):\n",
    "    x = min(max(x, eps), 1 - eps)\n",
    "    y = min(max(y, eps), 1 - eps)\n",
    "    return x * log(x / y) + (1 - x) * log((1 - x) / (1 - y))\n",
    "\n",
    "def klucbBern_cython(float x, float d, float precision=1e-6) -> float:\n",
    "    cdef float upperbound = min(1., klucbGauss_cython_x(x, d, sig2x=0.25))  # variance 1/4 for [0,1] bounded distributions\n",
    "    # upperbound = min(1., klucbPoisson(x, d))  # also safe, and better ?\n",
    "    return klucb_cython(x, d, klBern_cython_x, upperbound, precision)\n",
    "\n",
    "\n",
    "cdef float klPoisson_cython_x(float x, float y):\n",
    "    x = max(x, eps)\n",
    "    y = max(y, eps)\n",
    "    return y - x + x * log(x / y)\n",
    "\n",
    "def klucbPoisson_cython(float x, float d, float precision=1e-6) -> float:\n",
    "    cdef float upperbound = x + d + sqrt(d * d + 2 * x * d)  # looks safe, to check: left (Gaussian) tail of Poisson dev\n",
    "    return klucb_cython(x, d, klPoisson_cython_x, upperbound, precision)\n",
    "\n",
    "\n",
    "cdef float klGamma_cython_x(float x, float y):\n",
    "    if x <= 0 or y <= 0:\n",
    "        return float('+inf')\n",
    "    else:\n",
    "        x = max(x, eps)\n",
    "        y = max(y, eps)\n",
    "        return x / y - 1 - log(x / y)\n",
    "\n",
    "def klucbExp_cython(float x, float d, float precision=1e-6) -> float:\n",
    "    cdef float upperbound = 1\n",
    "    cdef float lowerbound = 0\n",
    "    if d < 0.77:  # XXX where does this value come from?\n",
    "        upperbound = x / (1 + 2. / 3 * d - sqrt(4. / 9 * d * d + 2 * d))\n",
    "        # safe, klexp(x,y) >= e^2/(2*(1-2e/3)) if x=y(1-e)\n",
    "    else:\n",
    "        upperbound = x * exp(d + 1)\n",
    "    if d > 1.61:  # XXX where does this value come from?\n",
    "        lowerbound = x * exp(d)\n",
    "    else:\n",
    "        lowerbound = x / (1 + d - sqrt(d * d + 2 * d))\n",
    "    return klucb_cython(x, d, klGamma_cython_x, upperbound, lowerbound, precision)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, for `klucbBern_cython`, the two steps are to first compute an upperbound (as precise as possible) and the compute the kl-UCB index:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.994140625"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.9944823980331421"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.994140625"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.9944896697998047"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x, d = 0.9, 0.2\n",
    "upperbound = 1\n",
    "klucb_cython(x, d, klBern_cython, upperbound, lowerbound=0, precision=1e-3, max_iterations=10)\n",
    "klucb_cython(x, d, klBern_cython, upperbound, lowerbound=0, precision=1e-6, max_iterations=10)\n",
    "klucb_cython(x, d, klBern_cython, upperbound, lowerbound=0, precision=1e-3, max_iterations=50)\n",
    "klucb_cython(x, d, klBern_cython, upperbound, lowerbound=0, precision=1e-6, max_iterations=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----\n",
    "# With the C API for Python\n",
    "\n",
    "It is more tedious, and won't be included here, but Python can easily [be extended](https://docs.python.org/3/c-api/) using C.\n",
    "It is the best way to obtain close-to-optimal performance for some parts of your code, and I will let you read [the introduction to the official documentation](https://docs.python.org/3/c-api/intro.html) if you are curious.\n",
    "\n",
    "For my [SMPyBandits](https://github.com/SMPyBandits/SMPyBandits/), I reused [some code from the py/maBandits](http://mloss.org/software/view/415/) project, and the authors implemented some of the previously defined KL-divergences and KL-UCB indexes in [pure Python](https://github.com/SMPyBandits/SMPyBandits/tree/master/SMPyBandits/Policies/kullback.py) as well as in [C optimized](https://github.com/SMPyBandits/SMPyBandits/tree/master/SMPyBandits/Policies/C/).\n",
    "I copied the compiled library in the current directory, and it can be imported:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-rw-r--r-- 1 lilian lilian 54K oct.  28 10:52 kullback.cpython-36m-x86_64-linux-gnu.so\n",
      "-rw-r--r-- 1 lilian lilian 19K janv.  9 17:31 kullback.py\n",
      "'kullback.py' -> 'kullback.py.old'\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "ls -larth *kullback*\n",
    "[ -f kullback.py ] && mv -vf kullback.py kullback.py.old"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-rw-r--r-- 1 lilian lilian 54K oct.  28 10:52 \u0000 kullback.cpython-36m-x86_64-linux-gnu.so\r\n"
     ]
    }
   ],
   "source": [
    "!ls -larth kullback*.so"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "import kullback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Help on built-in function klBern in module kullback:\n",
      "\n",
      "klBern(...)\n",
      "    klBern(x, y): Calculate the binary Kullback-Leibler divergence.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "help(kullback.klBern)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['klBern',\n",
       " 'klBin',\n",
       " 'klExp',\n",
       " 'klGamma',\n",
       " 'klGauss',\n",
       " 'klPoisson',\n",
       " 'klucbBern',\n",
       " 'klucbExp',\n",
       " 'klucbGamma',\n",
       " 'klucbGauss',\n",
       " 'klucbPoisson',\n",
       " 'maxEV']"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[ s for s in dir(kullback) if not s.startswith('_') ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "klBern_c = kullback.klBern\n",
    "klBin_c = kullback.klBin\n",
    "klExp_c = kullback.klExp\n",
    "klGamma_c = kullback.klGamma\n",
    "klGauss_c = kullback.klGauss\n",
    "klPoisson_c = kullback.klPoisson\n",
    "klucbBern_c = kullback.klucbBern\n",
    "klucbExp_c = kullback.klucbExp\n",
    "klucbGamma_c = kullback.klucbGamma\n",
    "klucbGauss_c = kullback.klucbGauss\n",
    "klucbPoisson_c = kullback.klucbPoisson"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to reproduce this notebook, download the [`kullback_py3.c`](https://github.com/SMPyBandits/SMPyBandits/blob/master/SMPyBandits/Policies/C/kullback_py3.c) and follow the [build instructions](https://github.com/SMPyBandits/SMPyBandits/blob/master/SMPyBandits/Policies/C/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----\n",
    "# Tests and benchmarks\n",
    "\n",
    "For each of the functions defined in three approaches above, I will do some numerical tests to compare their speed − and memory −  efficiency. Simple.\n",
    "\n",
    "The benchmark will be to test the computation time on random entries.\n",
    "It includes a constant time: creating random values! So I also compare the time to simply generate the values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = np.random.random\n",
    "rn = lambda: np.random.randint(1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "576 ns ± 8.53 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n",
      "2.35 µs ± 19.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit (r(), r())\n",
    "%timeit (r(), r(), rn())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- The time to generate random numbers like this is small, but not zero!\n",
    "- Generating a uniform integer, in particular, takes some time (more than 1 µs is not something that can be ignored!).\n",
    "\n",
    "$\\implies$ we will remove this $700$ ns or $2.5$ µs overhead when computing speed-up ratio between naive Python and numb or Cython versions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But we also need to test that the three versions of each function gives the same results (up-to approximation errors less than \n",
    "$10^{-6}$ (at least))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_fs(fs, inputs, tolerance=1e-5, nb_tests=100):\n",
    "    for _ in range(nb_tests):\n",
    "        args = inputs()\n",
    "        ref_f = fs[0]  # Python version\n",
    "        output = ref_f(*args)\n",
    "        for other_f in fs[1:]:\n",
    "            other_output = other_f(*args)\n",
    "            if abs(output) > 1:\n",
    "                rel_diff = (output - other_output) / output\n",
    "            else:\n",
    "                rel_diff = (output - other_output)\n",
    "            assert abs(rel_diff) <= tolerance, \"Error: function {} gave {} and function {} gave {} on inputs {}, and the two outputs are too different.\".format(ref_f, output, other_f, other_output, args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:red;\">WARNING</span> in the following, I use a very manual approach: I copied the time of each '%timeit' example, to compare speed-up ratios.\n",
    "So when I rerun the cells, the times might vary (a little bit), and I cannot keep an up-to-date versions of the computations of each ratio, so bear with me the (tiny) incoherences."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## KL divergences"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bernoulli"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_fs([klBern, klBern_numba, klBern_cython, klBern_c], lambda: (r(), r()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6.28 µs ± 919 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klBern(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 µs ± 145 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klBern_numba(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "882 ns ± 40.2 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klBern_cython(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "811 ns ± 40.6 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klBern_c(r(), r())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a speed-up ratio of about $12$ times faster for Numba and Cython, and $25$ times faster for the C version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "13.452830188679245"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "18.640522875816995"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "24.272340425531915"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(6280 - 576) / (1000 - 576)  # for Python vs numba\n",
    "(6280 - 576) / (882 - 576)   # for Python vs Cython\n",
    "(6280 - 576) / (811 - 576)   # for Python vs C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Binomial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "ename": "AssertionError",
     "evalue": "Error: function <function klBin at 0x7f8ab53acae8> gave 0.0011060494622968124 and function <built-in function klBin_cython> gave 0.0011183457989155888 on inputs (0.14985671245165777, 0.14884809478960037, 276), and the two outputs are too different.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-76-cad9d3d2157b>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mtest_fs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mklBin\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mklBin_numba\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mklBin_cython\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mklBin_c\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mlambda\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-63-12acbf032a73>\u001b[0m in \u001b[0;36mtest_fs\u001b[0;34m(fs, inputs, tolerance, nb_tests)\u001b[0m\n\u001b[1;32m     10\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m                 \u001b[0mrel_diff\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mother_output\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 12\u001b[0;31m             \u001b[0;32massert\u001b[0m \u001b[0mabs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrel_diff\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m<=\u001b[0m \u001b[0mtolerance\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"Error: function {} gave {} and function {} gave {} on inputs {}, and the two outputs are too different.\"\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mref_f\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother_f\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother_output\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mAssertionError\u001b[0m: Error: function <function klBin at 0x7f8ab53acae8> gave 0.0011060494622968124 and function <built-in function klBin_cython> gave 0.0011183457989155888 on inputs (0.14985671245165777, 0.14884809478960037, 276), and the two outputs are too different."
     ]
    }
   ],
   "source": [
    "test_fs([klBin, klBin_numba, klBin_cython, klBin_c], lambda: (r(), r(), rn()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Too much numerical difference? Let's try again with a larger tolerance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_fs([klBin, klBin_numba, klBin_cython, klBin_c], lambda: (r(), r(), rn()), tolerance=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7.05 µs ± 620 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klBin(r(), r(), rn())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.07 µs ± 146 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klBin_numba(r(), r(), rn())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.31 µs ± 258 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klBin_cython(r(), r(), rn())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.98 µs ± 148 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klBin_c(r(), r(), rn())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a speed-up ratio of about $5$ times faster for both Numba and Cython. Not so great, but still something!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6.465277777777778"
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "4.745158002038736"
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "7.388888888888889"
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(7005 - 2350) / (3070 - 2350)  # for Python vs numba\n",
    "(7005 - 2350) / (3331 - 2350)  # for Python vs Cython\n",
    "(7005 - 2350) / (2980 - 2350)  # for Python vs C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Poisson"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_fs([klPoisson, klPoisson_numba, klPoisson_cython, klPoisson_c], lambda: (r(), r()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.35 µs ± 172 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klPoisson(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "935 ns ± 30.7 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klPoisson_numba(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "859 ns ± 37.7 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klPoisson_cython(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "811 ns ± 36.1 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klPoisson_c(r(), r())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a speed-up ratio of about $7.5$ times faster for Numba, and about $7$ times for Cython and C."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4.9415041782729805"
      ]
     },
     "execution_count": 95,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "6.268551236749117"
      ]
     },
     "execution_count": 95,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "7.548936170212766"
      ]
     },
     "execution_count": 95,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(2350 - 576) / (935 - 576)  # for Python vs numba\n",
    "(2350 - 576) / (859 - 576)  # for Python vs Cython\n",
    "(2350 - 576) / (811 - 576)  # for Python vs C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exponential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_fs([klExp, klExp_numba, klExp_cython, klExp_c], lambda: (r(), r()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.21 µs ± 47.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klExp(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.07 µs ± 211 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klExp_numba(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "842 ns ± 32.2 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klExp_cython(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "869 ns ± 56 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klExp_c(r(), r())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a speed-up ratio of about $3$ times faster for Numba and $6$ times faster for Cython.\n",
    "Cython starts to win the race!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3.3076923076923075"
      ]
     },
     "execution_count": 101,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "6.142857142857143"
      ]
     },
     "execution_count": 101,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "5.57679180887372"
      ]
     },
     "execution_count": 101,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(2210 - 576) / (1070 - 576)  # for Python vs numba\n",
    "(2210 - 576) / (842 - 576)   # for Python vs Cython\n",
    "(2210 - 576) / (869 - 576)   # for Python vs C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gamma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [],
   "source": [
    "klGamma_c = lambda x, y: kullback.klGamma(x, y, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_fs([klGamma, klGamma_numba, klGamma_cython, klGamma_c], lambda: (r(), r()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.7 µs ± 163 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klGamma(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.07 µs ± 21.1 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klGamma_numba(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "889 ns ± 87.6 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klGamma_cython(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "997 ns ± 57.2 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klGamma_c(r(), r())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a speed-up ratio of about $6$ times faster for Numba, and $6$ or $5$ times faster for Cython and C.\n",
    "Note that the C version probably looses a little bit due to this `lambda , y: ...` trick."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4.299595141700405"
      ]
     },
     "execution_count": 111,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "6.785942492012779"
      ]
     },
     "execution_count": 111,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "5.045130641330166"
      ]
     },
     "execution_count": 111,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(2700 - 576) / (1070 - 576)  # for Python vs numba\n",
    "(2700 - 576) / (889 - 576)   # for Python vs Cython\n",
    "(2700 - 576) / (997 - 576)   # for Python vs C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Negative binomial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_fs([klNegBin, klNegBin_numba, klNegBin_cython], lambda: (r(), r()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.89 µs ± 249 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klNegBin(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.16 µs ± 76.8 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klNegBin_numba(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "901 ns ± 19.2 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klNegBin_cython(r(), r())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a speed-up ratio of about $5$ times faster for Numba and $10$ times faster for Cython."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5.674657534246576"
      ]
     },
     "execution_count": 116,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "10.196923076923078"
      ]
     },
     "execution_count": 116,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(3890 - 576) / (1160 - 576)  # for Python vs numba\n",
    "(3890 - 576) / (901 - 576)   # for Python vs Cython"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gaussian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {},
   "outputs": [],
   "source": [
    "klGauss_c = lambda x, y: kullback.klGauss(x, y, 0.25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_fs([klGauss, klGauss_numba, klGauss_cython, klGauss_c], lambda: (r(), r()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "852 ns ± 40.5 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klGauss(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.07 µs ± 54 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klGauss_numba(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "745 ns ± 40.6 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klGauss_cython(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "911 ns ± 56.6 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klGauss_c(r(), r())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a speed-up ratio of about $45$ times faster for Cython, but Numba completely here!\n",
    "Why? No idea!\n",
    "C is also slower than the Python version here, due to this `lambda x,y: ...` trick."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.5587044534412956"
      ]
     },
     "execution_count": 126,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "1.6331360946745561"
      ]
     },
     "execution_count": 126,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.8238805970149253"
      ]
     },
     "execution_count": 126,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(852 - 576) / (1070 - 576)  # for Python vs numba\n",
    "(852 - 576) / (745 - 576)   # for Python vs Cython\n",
    "(852 - 576) / (911 - 576)   # for Python vs C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## KL-UCB indexes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gaussian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {},
   "outputs": [],
   "source": [
    "klucbGauss_c = lambda x, y: kullback.klucbGauss(x, y, 0.25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_fs([klucbGauss, klucbGauss_numba, klucbGauss_cython, klucbGauss_c], lambda: (r(), r()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.96 µs ± 184 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klucbGauss(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "31.3 µs ± 950 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klucbGauss_numba(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "676 ns ± 25.7 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klucbGauss_cython(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "939 ns ± 23.6 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klucbGauss_c(r(), r())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a speed-up ratio of about $14$ times faster for Cython and $4$ times for C, and one more failure case for Numba.\n",
    "The C version looses again due to the `lambda x,y` trick."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.04504621794037235"
      ]
     },
     "execution_count": 135,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "13.84"
      ]
     },
     "execution_count": 135,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "3.81267217630854"
      ]
     },
     "execution_count": 135,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(1960 - 576) / (31300 - 576)  # for Python vs numba\n",
    "(1960 - 576) / (676 - 576)    # for Python vs Cython\n",
    "(1960 - 576) / (939 - 576)    # for Python vs C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bernoulli"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "metadata": {},
   "outputs": [],
   "source": [
    "klucbBern_c = lambda x, y: kullback.klucbBern(x, y, 1e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_fs([klucbBern, klucbBern_numba, klucbBern_cython, klucbBern_c], lambda: (r(), r()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "91.9 µs ± 4.7 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klucbBern(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "170 µs ± 8.43 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klucbBern_numba(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6.93 µs ± 815 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klucbBern_cython(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.14 µs ± 154 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klucbBern_c(r(), r())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a speed-up ratio of about $15$ times faster for Cython, and one more failure case for Numba.\n",
    "The speed-up for the C version is CRAZY here, and it shows that optimizing loops is the most important!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.5390263480970818"
      ]
     },
     "execution_count": 146,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "14.372678627636136"
      ]
     },
     "execution_count": 146,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "348.5648854961832"
      ]
     },
     "execution_count": 146,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(91900 - 576) / (170000 - 576)  # for Python vs numba\n",
    "(91900 - 576) / (6930 - 576)    # for Python vs Cython\n",
    "(91900 - 576) / abs(314 - 576)  # for Python vs C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Poisson"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "metadata": {},
   "outputs": [],
   "source": [
    "klucbPoisson_c = lambda x, y: kullback.klucbPoisson(x, y, 1e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_fs([klucbPoisson, klucbPoisson_numba, klucbPoisson_cython, klucbPoisson_c], lambda: (r(), r()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "72.6 µs ± 1.69 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klucbPoisson(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "167 µs ± 16.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klucbPoisson_numba(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.33 µs ± 382 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klucbPoisson_cython(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.18 µs ± 37.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klucbPoisson_c(r(), r())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a speed-up ratio of about $15$ times faster for Cython, and one more failure case for Numba.\n",
    "It's again a strong victory for the C version, with a speed up of about $45$ !"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.43277411911743496"
      ]
     },
     "execution_count": 154,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "15.150189314261674"
      ]
     },
     "execution_count": 154,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "44.90274314214464"
      ]
     },
     "execution_count": 154,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(72600 - 576) / (167000 - 576)  # for Python vs numba\n",
    "(72600 - 576) / (5330 - 576)    # for Python vs Cython\n",
    "(72600 - 576) / (2180 - 576)    # for Python vs Cython"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exponential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "metadata": {},
   "outputs": [],
   "source": [
    "klucbExp_c = lambda x, y: kullback.klucbExp(x, y, 1e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_fs([klucbExp, klucbExp_numba, klucbExp_cython, klucbExp_c], lambda: (r(), r()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "78.7 µs ± 898 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klucbExp(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "156 µs ± 10.1 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klucbExp_numba(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.41 µs ± 401 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klucbExp_cython(r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.04 µs ± 122 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit klucbExp_c(r(), r())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a speed-up ratio of about $17$ times faster for Cython, and one more failure case for Numba.\n",
    "Strong victory for the C version, a speed-up of $50$ is impressive!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.5026508132592135"
      ]
     },
     "execution_count": 161,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "20.37663015127804"
      ]
     },
     "execution_count": 161,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "53.36338797814208"
      ]
     },
     "execution_count": 161,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(78700 - 576) / (156000 - 576)  # for Python vs numba\n",
    "(78700 - 576) / (4410 - 576)    # for Python vs Cython\n",
    "(78700 - 576) / (2040 - 576)    # for Python vs Cython"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Clean up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 162,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'kullback.py.old' -> 'kullback.py'\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "[ -f kullback.py.old ] && mv -vf kullback.py.old kullback.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----\n",
    "# Conclusion\n",
    "\n",
    "- As expected, the Numba, Cython and C versions are *way* faster than the naive Python versions, on very simple functions,\n",
    "- The simpler the function, the closer the speed-up is between Numba and Cython or C,\n",
    "- Cython and C always give the best improvement,\n",
    "- On less simple functions, Numba can fail to produce `nopython` code, and on some examples the `nopython` code can be *slower* than naive Python (like, crazily slower). No idea why, and the point was precisely not to try too much optimizing this use of Numba.\n",
    "- Cython gives speed-up factors typically between $100$ and $12$ times faster than naive Python.\n",
    "- C alsways gives the best speed-up, with a speed-up between $150$ and $50$ times faster than naive Python, and usually $10$ times faster than Cython!\n",
    "\n",
    "## Take away messages\n",
    "The take away messages are the following:\n",
    "\n",
    "1. if your code makes a heavy use of a few small and not-too-complicated functions, it is probably worth using `numba.jit` to speed them up,\n",
    "2. but be careful, and do some basic benchmark on each \"possibly optimized\" function, to check that using Numba actually speeds it up instead of slowing it down!\n",
    "3. if Numba is not enough to speed up your code, try to write a Cython version of the bottleneck functions.\n",
    "4. if your bottleneck is *still* too slow, try to write a C extension, but it requires a good knowledge of both the internals of the C language, as well as a (basic) knowledge of the [C-API of CPython](https://docs.python.org/3/c-api/).\n",
    "\n",
    "## Using Cython *for real* ?\n",
    "My advice for using Cython are the following:\n",
    "\n",
    "1. First try in a notebook, using this [`%%cython` magic is very easy!](https://cython.readthedocs.io/en/latest/src/quickstart/build.html#using-the-jupyter-notebook)\n",
    "2. Then if you are happy about your implementation, save it to a `.pyx` file, and use [`pyximport`](https://cython.readthedocs.io/en/latest/src/userguide/source_files_and_compilation.html#pyximport) from your Python code to automatically compile and import it. It works perfectly fine, believe me!\n",
    "\n",
    "## Using C *for real* ?\n",
    "It's not too hard, but it's certainly not as easy as Cython.\n",
    "My approach from now will to not even consider writing C code, as Cython offers a very good speedup when carefully used. And it's much easier to just write `import pyximport` before importing your Cython-accelerated module, in comparison to writing a `setupy.py` and requiring a compilation step for a C-optimized module!\n",
    "\n",
    "> That's it for today, folks! See [this page](https://github.com/Naereen/notebooks) for other notebooks I wrote recently."
   ]
  }
 ],
 "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.3"
  },
  "toc": {
   "colors": {
    "hover_highlight": "#DAA520",
    "running_highlight": "#FF0000",
    "selected_highlight": "#FFD700"
   },
   "moveMenuLeft": true,
   "nav_menu": {
    "height": "11px",
    "width": "251px"
   },
   "navigate_menu": true,
   "number_sections": true,
   "sideBar": false,
   "threshold": 4,
   "toc_cell": true,
   "toc_position": {
    "height": "196px",
    "left": "985.125px",
    "right": "20px",
    "top": "96px",
    "width": "229px"
   },
   "toc_section_display": "block",
   "toc_window_display": true
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}