%load_ext watermark
%watermark -v -m -a "Lilian Besson (Naereen)" -p numpy,numba -g

import numpy as np

"""
Explanation: Table of Contents
<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>

----
# Introduction

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.

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.

Requirements:
End of explanation
"""


eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]

"""
Explanation: KL divergences and KL-UCB indexes, in naive Python
I will copy and paste parts of this file from my SMPyBandits library.
End of explanation
"""


def klBern(x, y):
    r""" Kullback-Leibler divergence for Bernoulli distributions. https://en.wikipedia.org/wiki/Bernoulli_distribution#Kullback.E2.80.93Leibler_divergence

    .. math:: \mathrm{KL}(\mathcal{B}(x), \mathcal{B}(y)) = x \log(\frac{x}{y}) + (1-x) \log(\frac{1-x}{1-y})."""
    x = min(max(x, eps), 1 - eps)
    y = min(max(y, eps), 1 - eps)
    return x * np.log(x / y) + (1 - x) * np.log((1 - x) / (1 - y))

klBern(0.5, 0.5)
klBern(0.1, 0.9)
klBern(0.9, 0.1)
klBern(0.4, 0.5)
klBern(0.01, 0.99)
klBern(0, 1)

"""
Explanation: I will include docstrings and examples only for the naive implementation.
KL divergences
Bernoulli distributions
End of explanation
"""


def klBin(x, y, n):
    r""" Kullback-Leibler divergence for Binomial distributions. https://math.stackexchange.com/questions/320399/kullback-leibner-divergence-of-binomial-distributions

    - It is simply the n times :func:`klBern` on x and y.

    .. 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).

    .. warning:: The two distributions must have the same parameter n, and x, y are p, q in (0, 1).
    """
    x = min(max(x, eps), 1 - eps)
    y = min(max(y, eps), 1 - eps)
    return n * (x * np.log(x / y) + (1 - x) * np.log((1 - x) / (1 - y)))

klBin(0.5, 0.5, 10)
klBin(0.1, 0.9, 10)
klBin(0.9, 0.1, 10)
klBin(0.4, 0.5, 10)
klBin(0.01, 0.99, 10)
klBin(0, 1, 10)

"""
Explanation: Binomial distributions
End of explanation
"""


def klPoisson(x, y):
    r""" Kullback-Leibler divergence for Poison distributions. https://en.wikipedia.org/wiki/Poisson_distribution#Kullback.E2.80.93Leibler_divergence

    .. math:: \mathrm{KL}(\mathrm{Poisson}(x), \mathrm{Poisson}(y)) = y - x + x \times \log(\frac{x}{y}).
    """
    x = max(x, eps)
    y = max(y, eps)
    return y - x + x * np.log(x / y)

klPoisson(3, 3)
klPoisson(2, 1)
klPoisson(1, 2)
klPoisson(3, 6)
klPoisson(6, 8)
klPoisson(1, 0)
klPoisson(0, 0)

"""
Explanation: Poisson distributions
End of explanation
"""


def klExp(x, y):
    r""" Kullback-Leibler divergence for exponential distributions. https://en.wikipedia.org/wiki/Exponential_distribution#Kullback.E2.80.93Leibler_divergence

    .. math::

        \mathrm{KL}(\mathrm{Exp}(x), \mathrm{Exp}(y)) = \begin{cases}
        \frac{x}{y} - 1 - \log(\frac{x}{y}) & \text{if} x > 0, y > 0\\
        +\infty & \text{otherwise}
        \end{cases}
    """
    if x <= 0 or y <= 0:
        return float('+inf')
    else:
        x = max(x, eps)
        y = max(y, eps)
        return x / y - 1 - np.log(x / y)

klExp(3, 3)
klExp(3, 6)
klExp(1, 2)
klExp(2, 1)
klExp(4, 2)
klExp(6, 8)
klExp(-3, 2)
klExp(3, -2)
klExp(-3, -2)

"""
Explanation: Exponential distributions
End of explanation
"""


def klGamma(x, y, a=1):
    r""" Kullback-Leibler divergence for gamma distributions. https://en.wikipedia.org/wiki/Gamma_distribution#Kullback.E2.80.93Leibler_divergence

    - It is simply the a times :func:`klExp` on x and y.

    .. math::

        \mathrm{KL}(\Gamma(x, a), \Gamma(y, a)) = \begin{cases}
        a \times \left( \frac{x}{y} - 1 - \log(\frac{x}{y}) \right) & \text{if} x > 0, y > 0\\
        +\infty & \text{otherwise}
        \end{cases}

    .. warning:: The two distributions must have the same parameter a.
    """
    if x <= 0 or y <= 0:
        return float('+inf')
    else:
        x = max(x, eps)
        y = max(y, eps)
        return a * (x / y - 1 - np.log(x / y))

klGamma(3, 3)
klGamma(3, 6)
klGamma(1, 2)
klGamma(2, 1)
klGamma(4, 2)
klGamma(6, 8)
klGamma(-3, 2)
klGamma(3, -2)
klGamma(-3, -2)

"""
Explanation: Gamma distributions
End of explanation
"""


def klNegBin(x, y, r=1):
    r""" Kullback-Leibler divergence for negative binomial distributions. https://en.wikipedia.org/wiki/Negative_binomial_distribution

    .. 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))).

    .. warning:: The two distributions must have the same parameter r.
    """
    x = max(x, eps)
    y = max(y, eps)
    return r * np.log((r + x) / (r + y)) - x * np.log(y * (r + x) / (x * (r + y)))

klNegBin(0.5, 0.5)
klNegBin(0.1, 0.9)
klNegBin(0.9, 0.1)
klNegBin(0.4, 0.5)
klNegBin(0.01, 0.99)
klBern(0, 1)
klNegBin(0.5, 0.5, r=2)
klNegBin(0.1, 0.9, r=2)
klNegBin(0.1, 0.9, r=4)
klNegBin(0.9, 0.1, r=2)
klNegBin(0.4, 0.5, r=2)
klNegBin(0.01, 0.99, r=2)

"""
Explanation: Negative binomial distributions
End of explanation
"""


def klGauss(x, y, sig2x=0.25, sig2y=None):
    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)`:

    .. 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).

    See https://en.wikipedia.org/wiki/Normal_distribution#Other_properties

    - By default, sig2y is assumed to be sig2x (same variance).
    """
    if sig2y is None or - eps < (sig2y - sig2x) < eps:
        return (x - y) ** 2 / (2. * sig2x)
    else:
        return (x - y) ** 2 / (2. * sig2y) + 0.5 * ((sig2x/sig2y)**2 - 1 - np.log(sig2x/sig2y))

klGauss(3, 3)
klGauss(3, 6)
klGauss(1, 2)
klGauss(2, 1)
klGauss(4, 2)
klGauss(6, 8)
klGauss(-3, 2)
klGauss(3, -2)
klGauss(-3, -2)
klGauss(3, 2)
klGauss(3, 3, sig2x=10)
klGauss(3, 6, sig2x=10)
klGauss(1, 2, sig2x=10)
klGauss(2, 1, sig2x=10)
klGauss(4, 2, sig2x=10)
klGauss(6, 8, sig2x=10)
klGauss(0, 0, sig2x=0.25, sig2y=0.5)
klGauss(0, 0, sig2x=0.25, sig2y=1.0)
klGauss(0, 0, sig2x=0.5, sig2y=0.25)
klGauss(0, 1, sig2x=0.25, sig2y=0.5)
klGauss(0, 1, sig2x=0.25, sig2y=1.0)
klGauss(0, 1, sig2x=0.5, sig2y=0.25)
klGauss(1, 0, sig2x=0.25, sig2y=0.5)
klGauss(1, 0, sig2x=0.25, sig2y=1.0)
klGauss(1, 0, sig2x=0.5, sig2y=0.25)

"""
Explanation: Gaussian distributions
End of explanation
"""


def klucb(x, d, kl, upperbound, lowerbound=float('-inf'), precision=1e-6, max_iterations=50):
    """ The generic KL-UCB index computation.

    - x: value of the cum reward,
    - d: upper bound on the divergence,
    - kl: the KL divergence to be used (:func:`klBern`, :func:`klGauss`, etc),
    - upperbound, lowerbound=float('-inf'): the known bound of the values x,
    - precision=1e-6: the threshold from where to stop the research,
    - max_iterations: max number of iterations of the loop (safer to bound it to reduce time complexity).

    .. note:: It uses a **bisection search**, and one call to ``kl`` for each step of the bisection search.
    """
    value = max(x, lowerbound)
    u = upperbound
    _count_iteration = 0
    while _count_iteration < max_iterations and u - value > precision:
        _count_iteration += 1
        m = (value + u) / 2.
        if kl(x, m) > d:
            u = m
        else:
            value = m
    return (value + u) / 2.

"""
Explanation: Generic KL-UCB indexes, with a bisection search
End of explanation
"""


x, d = 0.9, 0.2
upperbound = 1
klucb(x, d, klBern, upperbound, lowerbound=0, precision=1e-3, max_iterations=10)
klucb(x, d, klBern, upperbound, lowerbound=0, precision=1e-6, max_iterations=10)
klucb(x, d, klBern, upperbound, lowerbound=0, precision=1e-3, max_iterations=50)
klucb(x, d, klBern, upperbound, lowerbound=0, precision=1e-6, max_iterations=100)

"""
Explanation: For example, for klucbBern, the two steps are to first compute an upperbound (as precise as possible) and the compute the kl-UCB index:
End of explanation
"""


def klucbGauss(x, d, sig2x=0.25, precision=0.):
    """ KL-UCB index computation for Gaussian distributions.

    - Note that it does not require any search.

    .. warning:: it works only if the good variance constant is given.
    """
    return x + np.sqrt(2 * sig2x * d)

klucbGauss(0.1, 0.2)
klucbGauss(0.5, 0.2)
klucbGauss(0.9, 0.2)
klucbGauss(0.1, 0.4)
klucbGauss(0.1, 0.9)
klucbGauss(0.5, 0.4)
klucbGauss(0.5, 0.9)
klucbGauss(0.9, 0.4)
klucbGauss(0.9, 0.9)

"""
Explanation: Distribution-specific KL-UCB indexes
Gaussian
End of explanation
"""


def klucbBern(x, d, precision=1e-6):
    """ KL-UCB index computation for Bernoulli distributions, using :func:`klucb`."""
    upperbound = min(1., klucbGauss(x, d, sig2x=0.25))  # variance 1/4 for [0,1] bounded distributions
    # upperbound = min(1., klucbPoisson(x, d))  # also safe, and better ?
    return klucb(x, d, klBern, upperbound, precision)

klucbBern(0.1, 0.2)
klucbBern(0.5, 0.2)
klucbBern(0.9, 0.2)
klucbBern(0.1, 0.4)
klucbBern(0.1, 0.9)
klucbBern(0.5, 0.4)
klucbBern(0.5, 0.9)
klucbBern(0.9, 0.4)
klucbBern(0.9, 0.9)

"""
Explanation: Bernoulli
End of explanation
"""


def klucbPoisson(x, d, precision=1e-6):
    """ KL-UCB index computation for Poisson distributions, using :func:`klucb`."""
    upperbound = x + d + np.sqrt(d * d + 2 * x * d)  # looks safe, to check: left (Gaussian) tail of Poisson dev
    return klucb(x, d, klPoisson, upperbound, precision)

klucbPoisson(0.1, 0.2)
klucbPoisson(0.5, 0.2)
klucbPoisson(0.9, 0.2)
klucbPoisson(0.1, 0.4)
klucbPoisson(0.1, 0.9)
klucbPoisson(0.5, 0.4)
klucbPoisson(0.5, 0.9)
klucbPoisson(0.9, 0.4)
klucbPoisson(0.9, 0.9)

"""
Explanation: Poisson
End of explanation
"""


def klucbExp(x, d, precision=1e-6):
    """ KL-UCB index computation for exponential distributions, using :func:`klucb`."""
    if d < 0.77:  # XXX where does this value come from?
        upperbound = x / (1 + 2. / 3 * d - np.sqrt(4. / 9 * d * d + 2 * d))
        # safe, klexp(x,y) >= e^2/(2*(1-2e/3)) if x=y(1-e)
    else:
        upperbound = x * np.exp(d + 1)
    if d > 1.61:  # XXX where does this value come from?
        lowerbound = x * np.exp(d)
    else:
        lowerbound = x / (1 + d - np.sqrt(d * d + 2 * d))
    return klucb(x, d, klGamma, upperbound, lowerbound, precision)

klucbExp(0.1, 0.2)
klucbExp(0.5, 0.2)
klucbExp(0.9, 0.2)
klucbExp(0.1, 0.4)
klucbExp(0.1, 0.9)
klucbExp(0.5, 0.4)
klucbExp(0.5, 0.9)
klucbExp(0.9, 0.4)
klucbExp(0.9, 0.9)

"""
Explanation: Exponential
End of explanation
"""


from numba import jit

"""
Explanation: Others
We could do the same for more distributions, but that's enough.

With Numba
It will be exactly the same code as above, except that the numba.jit decorator will be used for each functions, to let numba try to speed up the code!
End of explanation
"""


@jit(nopython=True)
def klBern_numba(x, y):
    x = min(max(x, eps), 1 - eps)
    y = min(max(y, eps), 1 - eps)
    return x * np.log(x / y) + (1 - x) * np.log((1 - x) / (1 - y))

"""
Explanation: 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.
KL divergences
Bernoulli distributions
End of explanation
"""


@jit(nopython=True)
def klBin_numba(x, y, n):
    x = min(max(x, eps), 1 - eps)
    y = min(max(y, eps), 1 - eps)
    return n * (x * np.log(x / y) + (1 - x) * np.log((1 - x) / (1 - y)))

"""
Explanation: Binomial distributions
End of explanation
"""


@jit(nopython=True)
def klPoisson_numba(x, y):
    x = max(x, eps)
    y = max(y, eps)
    return y - x + x * np.log(x / y)

"""
Explanation: Poisson distributions
End of explanation
"""


@jit(nopython=True)
def klExp_numba(x, y):
    if x <= 0 or y <= 0:
        return inf
    else:
        x = max(x, eps)
        y = max(y, eps)
        return x / y - 1 - np.log(x / y)

"""
Explanation: Exponential distributions
End of explanation
"""


@jit(nopython=True)
def klGamma_numba(x, y, a=1):
    if x <= 0 or y <= 0:
        return inf
    else:
        x = max(x, eps)
        y = max(y, eps)
        return a * (x / y - 1 - np.log(x / y))

"""
Explanation: Gamma distributions
End of explanation
"""


@jit(nopython=True)
def klNegBin_numba(x, y, r=1):
    x = max(x, eps)
    y = max(y, eps)
    return r * np.log((r + x) / (r + y)) - x * np.log(y * (r + x) / (x * (r + y)))

"""
Explanation: Negative binomial distributions
End of explanation
"""


@jit(nopython=True)
def klGauss_numba(x, y, sig2x=0.25, sig2y=0.25):
    if - eps < (sig2y - sig2x) and (sig2y - sig2x) < eps:
        return (x - y) ** 2 / (2. * sig2x)
    else:
        return (x - y) ** 2 / (2. * sig2y) + 0.5 * ((sig2x/sig2y)**2 - 1 - np.log(sig2x/sig2y))

"""
Explanation: Gaussian distributions
End of explanation
"""


@jit
def klucb_numba(x, d, kl, upperbound,
                lowerbound=float('-inf'), precision=1e-6, max_iterations=50):
    value = max(x, lowerbound)
    u = upperbound
    _count_iteration = 0
    while _count_iteration < max_iterations and u - value > precision:
        _count_iteration += 1
        m = (value + u) / 2.
        if kl(x, m) > d:
            u = m
        else:
            value = m
    return (value + u) / 2.

"""
Explanation: Generic KL-UCB indexes, with a bisection search
End of explanation
"""


x, d = 0.9, 0.2
upperbound = 1
klucb_numba(x, d, klBern_numba, upperbound, lowerbound=0, precision=1e-3, max_iterations=10)
klucb_numba(x, d, klBern_numba, upperbound, lowerbound=0, precision=1e-6, max_iterations=10)
klucb_numba(x, d, klBern_numba, upperbound, lowerbound=0, precision=1e-3, max_iterations=50)
klucb_numba(x, d, klBern_numba, upperbound, lowerbound=0, precision=1e-6, max_iterations=100)

"""
Explanation: For example, for klucbBern, the two steps are to first compute an upperbound (as precise as possible) and the compute the kl-UCB index:
End of explanation
"""


@jit(nopython=True)
def klucbGauss_numba(x, d, sig2x=0.25, precision=0.):
    return x + np.sqrt(2 * sig2x * d)

"""
Explanation: Distribution-specific KL-UCB indexes
Gaussian
End of explanation
"""


@jit
def klucbBern_numba(x, d, precision=1e-6):
    upperbound = min(1., klucbGauss_numba(x, d, sig2x=0.25))  # variance 1/4 for [0,1] bounded distributions
    # upperbound = min(1., klucbPoisson(x, d))  # also safe, and better ?
    return klucb_numba(x, d, klBern_numba, upperbound, precision)

"""
Explanation: Bernoulli
Here, the nopython=True fails as numba has a hard time typing linked function calls.
End of explanation
"""


@jit
def klucbPoisson_numba(x, d, precision=1e-6):
    upperbound = x + d + np.sqrt(d * d + 2 * x * d)  # looks safe, to check: left (Gaussian) tail of Poisson dev
    return klucb_numba(x, d, klPoisson_numba, upperbound, precision)

"""
Explanation: Poisson
End of explanation
"""


@jit
def klucbExp_numba(x, d, precision=1e-6):
    if d < 0.77:  # XXX where does this value come from?
        upperbound = x / (1 + 2. / 3 * d - np.sqrt(4. / 9 * d * d + 2 * d))
        # safe, klexp(x,y) >= e^2/(2*(1-2e/3)) if x=y(1-e)
    else:
        upperbound = x * np.exp(d + 1)
    if d > 1.61:  # XXX where does this value come from?
        lowerbound = x * np.exp(d)
    else:
        lowerbound = x / (1 + d - np.sqrt(d * d + 2 * d))
    return klucb_numba(x, d, klGamma_numba, upperbound, lowerbound, precision)

"""
Explanation: Exponential
End of explanation
"""


%load_ext cython

"""
Explanation: With Cython
It will be almost exactly the same code, by using the cython magic to have cells written in Cython.
End of explanation
"""


def some_loop(n: int) -> int:
    s = 0
    for i in range(0, n, 2):
        s += i
    return s

%%cython
def some_loop_cython(int n) -> int:
    cdef int s = 0
    cdef int i = 0
    for i in range(0, n, 2):
        s += i
    return s

%timeit np.random.randint(1000)
%timeit some_loop(np.random.randint(1000))
%timeit some_loop_cython(np.random.randint(1000))

"""
Explanation: A cell can now be written in Cython.
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:
End of explanation
"""


14.6 / 2.21

"""
Explanation: Here we observe a large speed-up. But how large? $6$ times or $50$ times?
It's really important to include the time taken by the Pseudo-Random Number Generator:

Wrong computation of the speed-up gives about $6$ times faster:
End of explanation
"""


(14.6 - 1.95) / (2.21 - 1.95)

"""
Explanation: 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!
End of explanation
"""


%%cython
from libc.math cimport log
eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]

def klBern_cython(float x, float y) -> float:
    x = min(max(x, eps), 1 - eps)
    y = min(max(y, eps), 1 - eps)
    return x * log(x / y) + (1 - x) * log((1 - x) / (1 - y))

"""
Explanation: KL divergences
Bernoulli distributions
End of explanation
"""


%%cython
from libc.math cimport log
eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]

def klBin_cython(float x, float y, int n) -> float:
    x = min(max(x, eps), 1 - eps)
    y = min(max(y, eps), 1 - eps)
    return n * (x * log(x / y) + (1 - x) * log((1 - x) / (1 - y)))

"""
Explanation: Binomial distributions
End of explanation
"""


%%cython
from libc.math cimport log
eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]

def klPoisson_cython(float x, float y) -> float:
    x = max(x, eps)
    y = max(y, eps)
    return y - x + x * log(x / y)

"""
Explanation: Poisson distributions
End of explanation
"""


%%cython
from libc.math cimport log
eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]

def klExp_cython(float x, float y) -> float:
    if x <= 0 or y <= 0:
        return float('+inf')
    else:
        x = max(x, eps)
        y = max(y, eps)
        return x / y - 1 - log(x / y)

"""
Explanation: Exponential distributions
End of explanation
"""


%%cython
from libc.math cimport log
eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]

def klGamma_cython(float x, float y, float a=1) -> float:
    if x <= 0 or y <= 0:
        return float('+inf')
    else:
        x = max(x, eps)
        y = max(y, eps)
        return a * (x / y - 1 - log(x / y))

"""
Explanation: Gamma distributions
End of explanation
"""


%%cython
from libc.math cimport log
eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]

def klNegBin_cython(float x, float y, float r=1) -> float:
    x = max(x, eps)
    y = max(y, eps)
    return r * log((r + x) / (r + y)) - x * log(y * (r + x) / (x * (r + y)))

"""
Explanation: Negative binomial distributions
End of explanation
"""


%%cython
from libc.math cimport log
eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]

def klGauss_cython(float x, float y, float sig2x=0.25, float sig2y=0.25) -> float:
    if - eps < (sig2y - sig2x) < eps:
        return (x - y) ** 2 / (2. * sig2x)
    else:
        return (x - y) ** 2 / (2. * sig2y) + 0.5 * ((sig2x/sig2y)**2 - 1 - log(sig2x/sig2y))

"""
Explanation: Gaussian distributions
End of explanation
"""


%%cython
from libc.math cimport sqrt, log, exp
eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]


def klucbGauss_cython(float x, float d, float sig2x=0.25, float precision=0.) -> float:
    return x + sqrt(2 * sig2x * d)

cdef float klucbGauss_cython_x(float x, float d, float sig2x=0.25, float precision=0.):
    return x + sqrt(2 * sig2x * d)


def klucb_cython(float x, float d, kl, float upperbound,
                 float lowerbound=float('-inf'),
                 float precision=1e-6, int max_iterations=50) -> float:
    cdef float value = max(x, lowerbound)
    cdef float u = upperbound
    cdef int _count_iteration = 0
    cdef float m = 0
    while _count_iteration < max_iterations and u - value > precision:
        _count_iteration += 1
        m = (value + u) / 2.
        if kl(x, m) > d:
            u = m
        else:
            value = m
    return (value + u) / 2.


cdef float klBern_cython_x(float x, float y):
    x = min(max(x, eps), 1 - eps)
    y = min(max(y, eps), 1 - eps)
    return x * log(x / y) + (1 - x) * log((1 - x) / (1 - y))

def klucbBern_cython(float x, float d, float precision=1e-6) -> float:
    cdef float upperbound = min(1., klucbGauss_cython_x(x, d, sig2x=0.25))  # variance 1/4 for [0,1] bounded distributions
    # upperbound = min(1., klucbPoisson(x, d))  # also safe, and better ?
    return klucb_cython(x, d, klBern_cython_x, upperbound, precision)


cdef float klPoisson_cython_x(float x, float y):
    x = max(x, eps)
    y = max(y, eps)
    return y - x + x * log(x / y)

def klucbPoisson_cython(float x, float d, float precision=1e-6) -> float:
    cdef float upperbound = x + d + sqrt(d * d + 2 * x * d)  # looks safe, to check: left (Gaussian) tail of Poisson dev
    return klucb_cython(x, d, klPoisson_cython_x, upperbound, precision)


cdef float klGamma_cython_x(float x, float y):
    if x <= 0 or y <= 0:
        return float('+inf')
    else:
        x = max(x, eps)
        y = max(y, eps)
        return x / y - 1 - log(x / y)

def klucbExp_cython(float x, float d, float precision=1e-6) -> float:
    cdef float upperbound = 1
    cdef float lowerbound = 0
    if d < 0.77:  # XXX where does this value come from?
        upperbound = x / (1 + 2. / 3 * d - sqrt(4. / 9 * d * d + 2 * d))
        # safe, klexp(x,y) >= e^2/(2*(1-2e/3)) if x=y(1-e)
    else:
        upperbound = x * exp(d + 1)
    if d > 1.61:  # XXX where does this value come from?
        lowerbound = x * exp(d)
    else:
        lowerbound = x / (1 + d - sqrt(d * d + 2 * d))
    return klucb_cython(x, d, klGamma_cython_x, upperbound, lowerbound, precision)

"""
Explanation: Generic KL-UCB indexes, with a bisection search
For these, they need previously defined functions, which have to be rewritten from inside the cython cell to be accessible from Cython.
To minimize repetitions, I use only one cell to define all functions.
End of explanation
"""


x, d = 0.9, 0.2
upperbound = 1
klucb_cython(x, d, klBern_cython, upperbound, lowerbound=0, precision=1e-3, max_iterations=10)
klucb_cython(x, d, klBern_cython, upperbound, lowerbound=0, precision=1e-6, max_iterations=10)
klucb_cython(x, d, klBern_cython, upperbound, lowerbound=0, precision=1e-3, max_iterations=50)
klucb_cython(x, d, klBern_cython, upperbound, lowerbound=0, precision=1e-6, max_iterations=100)

"""
Explanation: 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:
End of explanation
"""


%%bash
ls -larth *kullback*
[ -f kullback.py ] && mv -vf kullback.py kullback.py.old

!ls -larth kullback*.so

import kullback

help(kullback.klBern)

[ s for s in dir(kullback) if not s.startswith('_') ]

klBern_c = kullback.klBern
klBin_c = kullback.klBin
klExp_c = kullback.klExp
klGamma_c = kullback.klGamma
klGauss_c = kullback.klGauss
klPoisson_c = kullback.klPoisson
klucbBern_c = kullback.klucbBern
klucbExp_c = kullback.klucbExp
klucbGamma_c = kullback.klucbGamma
klucbGauss_c = kullback.klucbGauss
klucbPoisson_c = kullback.klucbPoisson

"""
Explanation: With the C API for Python
It is more tedious, and won't be included here, but Python can easily be extended using C.
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 if you are curious.
For my SMPyBandits, I reused some code from the py/maBandits project, and the authors implemented some of the previously defined KL-divergences and KL-UCB indexes in pure Python as well as in C optimized.
I copied the compiled library in the current directory, and it can be imported:
End of explanation
"""


r = np.random.random
rn = lambda: np.random.randint(1000)

%timeit (r(), r())
%timeit (r(), r(), rn())

"""
Explanation: If you want to reproduce this notebook, download the kullback_py3.c and follow the build instructions.

Tests and benchmarks
For each of the functions defined in three approaches above, I will do some numerical tests to compare their speed − and memory −  efficiency. Simple.
The benchmark will be to test the computation time on random entries.
It includes a constant time: creating random values! So I also compare the time to simply generate the values.
End of explanation
"""


def test_fs(fs, inputs, tolerance=1e-5, nb_tests=100):
    for _ in range(nb_tests):
        args = inputs()
        ref_f = fs[0]  # Python version
        output = ref_f(*args)
        for other_f in fs[1:]:
            other_output = other_f(*args)
            if abs(output) > 1:
                rel_diff = (output - other_output) / output
            else:
                rel_diff = (output - other_output)
            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)

"""
Explanation: The time to generate random numbers like this is small, but not zero!
Generating a uniform integer, in particular, takes some time (more than 1 µs is not something that can be ignored!).

$\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.
But we also need to test that the three versions of each function gives the same results (up-to approximation errors less than 
$10^{-6}$ (at least)).
End of explanation
"""


test_fs([klBern, klBern_numba, klBern_cython, klBern_c], lambda: (r(), r()))

%timeit klBern(r(), r())

%timeit klBern_numba(r(), r())

%timeit klBern_cython(r(), r())

%timeit klBern_c(r(), r())

"""
Explanation: <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.
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.
KL divergences
Bernoulli
End of explanation
"""


(6280 - 576) / (1000 - 576)  # for Python vs numba
(6280 - 576) / (882 - 576)   # for Python vs Cython
(6280 - 576) / (811 - 576)   # for Python vs C

"""
Explanation: This is a speed-up ratio of about $12$ times faster for Numba and Cython, and $25$ times faster for the C version.
End of explanation
"""


test_fs([klBin, klBin_numba, klBin_cython, klBin_c], lambda: (r(), r(), rn()))

"""
Explanation: Binomial
End of explanation
"""


test_fs([klBin, klBin_numba, klBin_cython, klBin_c], lambda: (r(), r(), rn()), tolerance=1e-3)

%timeit klBin(r(), r(), rn())

%timeit klBin_numba(r(), r(), rn())

%timeit klBin_cython(r(), r(), rn())

%timeit klBin_c(r(), r(), rn())

"""
Explanation: Too much numerical difference? Let's try again with a larger tolerance:
End of explanation
"""


(7005 - 2350) / (3070 - 2350)  # for Python vs numba
(7005 - 2350) / (3331 - 2350)  # for Python vs Cython
(7005 - 2350) / (2980 - 2350)  # for Python vs C

"""
Explanation: This is a speed-up ratio of about $5$ times faster for both Numba and Cython. Not so great, but still something!
End of explanation
"""


test_fs([klPoisson, klPoisson_numba, klPoisson_cython, klPoisson_c], lambda: (r(), r()))

%timeit klPoisson(r(), r())

%timeit klPoisson_numba(r(), r())

%timeit klPoisson_cython(r(), r())

%timeit klPoisson_c(r(), r())

"""
Explanation: Poisson
End of explanation
"""


(2350 - 576) / (935 - 576)  # for Python vs numba
(2350 - 576) / (859 - 576)  # for Python vs Cython
(2350 - 576) / (811 - 576)  # for Python vs C

"""
Explanation: This is a speed-up ratio of about $7.5$ times faster for Numba, and about $7$ times for Cython and C.
End of explanation
"""


test_fs([klExp, klExp_numba, klExp_cython, klExp_c], lambda: (r(), r()))

%timeit klExp(r(), r())

%timeit klExp_numba(r(), r())

%timeit klExp_cython(r(), r())

%timeit klExp_c(r(), r())

"""
Explanation: Exponential
End of explanation
"""


(2210 - 576) / (1070 - 576)  # for Python vs numba
(2210 - 576) / (842 - 576)   # for Python vs Cython
(2210 - 576) / (869 - 576)   # for Python vs C

"""
Explanation: This is a speed-up ratio of about $3$ times faster for Numba and $6$ times faster for Cython.
Cython starts to win the race!
End of explanation
"""


klGamma_c = lambda x, y: kullback.klGamma(x, y, 1)

test_fs([klGamma, klGamma_numba, klGamma_cython, klGamma_c], lambda: (r(), r()))

%timeit klGamma(r(), r())

%timeit klGamma_numba(r(), r())

%timeit klGamma_cython(r(), r())

%timeit klGamma_c(r(), r())

"""
Explanation: Gamma
End of explanation
"""


(2700 - 576) / (1070 - 576)  # for Python vs numba
(2700 - 576) / (889 - 576)   # for Python vs Cython
(2700 - 576) / (997 - 576)   # for Python vs C

"""
Explanation: This is a speed-up ratio of about $6$ times faster for Numba, and $6$ or $5$ times faster for Cython and C.
Note that the C version probably looses a little bit due to this lambda , y: ... trick.
End of explanation
"""


test_fs([klNegBin, klNegBin_numba, klNegBin_cython], lambda: (r(), r()))

%timeit klNegBin(r(), r())

%timeit klNegBin_numba(r(), r())

%timeit klNegBin_cython(r(), r())

"""
Explanation: Negative binomial
End of explanation
"""


(3890 - 576) / (1160 - 576)  # for Python vs numba
(3890 - 576) / (901 - 576)   # for Python vs Cython

"""
Explanation: This is a speed-up ratio of about $5$ times faster for Numba and $10$ times faster for Cython.
End of explanation
"""


klGauss_c = lambda x, y: kullback.klGauss(x, y, 0.25)

test_fs([klGauss, klGauss_numba, klGauss_cython, klGauss_c], lambda: (r(), r()))

%timeit klGauss(r(), r())

%timeit klGauss_numba(r(), r())

%timeit klGauss_cython(r(), r())

%timeit klGauss_c(r(), r())

"""
Explanation: Gaussian
End of explanation
"""


(852 - 576) / (1070 - 576)  # for Python vs numba
(852 - 576) / (745 - 576)   # for Python vs Cython
(852 - 576) / (911 - 576)   # for Python vs C

"""
Explanation: This is a speed-up ratio of about $45$ times faster for Cython, but Numba completely here!
Why? No idea!
C is also slower than the Python version here, due to this lambda x,y: ... trick.
End of explanation
"""


klucbGauss_c = lambda x, y: kullback.klucbGauss(x, y, 0.25)

test_fs([klucbGauss, klucbGauss_numba, klucbGauss_cython, klucbGauss_c], lambda: (r(), r()))

%timeit klucbGauss(r(), r())

%timeit klucbGauss_numba(r(), r())

%timeit klucbGauss_cython(r(), r())

%timeit klucbGauss_c(r(), r())

"""
Explanation: KL-UCB indexes
Gaussian
End of explanation
"""


(1960 - 576) / (31300 - 576)  # for Python vs numba
(1960 - 576) / (676 - 576)    # for Python vs Cython
(1960 - 576) / (939 - 576)    # for Python vs C

"""
Explanation: 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.
The C version looses again due to the lambda x,y trick.
End of explanation
"""


klucbBern_c = lambda x, y: kullback.klucbBern(x, y, 1e-6)

test_fs([klucbBern, klucbBern_numba, klucbBern_cython, klucbBern_c], lambda: (r(), r()))

%timeit klucbBern(r(), r())

%timeit klucbBern_numba(r(), r())

%timeit klucbBern_cython(r(), r())

%timeit klucbBern_c(r(), r())

"""
Explanation: Bernoulli
End of explanation
"""


(91900 - 576) / (170000 - 576)  # for Python vs numba
(91900 - 576) / (6930 - 576)    # for Python vs Cython
(91900 - 576) / abs(314 - 576)  # for Python vs C

"""
Explanation: This is a speed-up ratio of about $15$ times faster for Cython, and one more failure case for Numba.
The speed-up for the C version is CRAZY here, and it shows that optimizing loops is the most important!
End of explanation
"""


klucbPoisson_c = lambda x, y: kullback.klucbPoisson(x, y, 1e-6)

test_fs([klucbPoisson, klucbPoisson_numba, klucbPoisson_cython, klucbPoisson_c], lambda: (r(), r()))

%timeit klucbPoisson(r(), r())

%timeit klucbPoisson_numba(r(), r())

%timeit klucbPoisson_cython(r(), r())

%timeit klucbPoisson_c(r(), r())

"""
Explanation: Poisson
End of explanation
"""


(72600 - 576) / (167000 - 576)  # for Python vs numba
(72600 - 576) / (5330 - 576)    # for Python vs Cython
(72600 - 576) / (2180 - 576)    # for Python vs Cython

"""
Explanation: This is a speed-up ratio of about $15$ times faster for Cython, and one more failure case for Numba.
It's again a strong victory for the C version, with a speed up of about $45$ !
End of explanation
"""


klucbExp_c = lambda x, y: kullback.klucbExp(x, y, 1e-6)

test_fs([klucbExp, klucbExp_numba, klucbExp_cython, klucbExp_c], lambda: (r(), r()))

%timeit klucbExp(r(), r())

%timeit klucbExp_numba(r(), r())

%timeit klucbExp_cython(r(), r())

%timeit klucbExp_c(r(), r())

"""
Explanation: Exponential
End of explanation
"""


(78700 - 576) / (156000 - 576)  # for Python vs numba
(78700 - 576) / (4410 - 576)    # for Python vs Cython
(78700 - 576) / (2040 - 576)    # for Python vs Cython

"""
Explanation: This is a speed-up ratio of about $17$ times faster for Cython, and one more failure case for Numba.
Strong victory for the C version, a speed-up of $50$ is impressive!
End of explanation
"""


%%bash
[ -f kullback.py.old ] && mv -vf kullback.py.old kullback.py

"""
Explanation: Clean up
End of explanation
"""


