06-numpy-buffers-fused-types.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cython and NumPy, buffers, fused types\n",
    "\n",
    "We'll be re-implementing the `scipy.stats.poisson.entropy()` function / method using Cython and its NumPy support, and we'll see what sort of speedup we can get for our efforts.\n",
    "\n",
    "Obligatory XKCD:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"https://imgs.xkcd.com/comics/password_strength.png\" width=600 height=600 />"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext Cython\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pure Python implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# https://en.wiktionary.org/wiki/Shannon_entropy\n",
    "\n",
    "def shannon_entropy_py(p_x):\n",
    "    return - np.sum(p_x * np.log(p_x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cythonized version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -a\n",
    "\n",
    "import numpy as np\n",
    "cimport numpy as cnp\n",
    "\n",
    "def shannon_entropy_cy(cnp.ndarray p_x):\n",
    "    return - np.sum(p_x * np.log(p_x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What's `cimport`?\n",
    "\n",
    "Cython introduces a new keyword, `cimport` that allows interfacing with other Cython code _at the C-level_ and _at compile time_.\n",
    "\n",
    "This is in distinction to Python's regular `import` statement, which interfaces with other _Python_ modules at _runtime_.\n",
    "\n",
    "The statement `cimport numpy as cnp` allows Cython to interface with NumPy arrays at the C level and at compile time for improved performance.\n",
    "\n",
    "The `cimport numpy` statement causes the Cython compiler to look for a `numpy.pxd` Cython declaration file at compile time.  This is where we can find the C-level declarations of the NumPy C-API.  Here we're using the `ndarray` object from `numpy.pxd` to declare the argument of `shannon_entropy_cy()`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scipy.stats comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from scipy.stats import poisson\n",
    "poi = poisson(10.0)\n",
    "n = 100\n",
    "pmf = poi.pmf(np.arange(n))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(poi.entropy())\n",
    "print(shannon_entropy_py(pmf))\n",
    "print(shannon_entropy_cy(pmf))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "poi.entropy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "shannon_entropy_py(pmf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "shannon_entropy_cy(pmf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Explicit `for` loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -a\n",
    "\n",
    "cimport numpy as cnp\n",
    "from libc.math cimport log as clog\n",
    "\n",
    "def shannon_entropy_v1(cnp.ndarray p_x):\n",
    "    cdef double res = 0.0\n",
    "    cdef int n = p_x.shape[0]\n",
    "    cdef int i\n",
    "    for i in range(n):\n",
    "        res += p_x[i] * clog(p_x[i])\n",
    "    return -res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "shannon_entropy_v1(pmf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What's `from libc.math cimport log`?\n",
    "\n",
    "Cython allows us to `cimport` C (and C++) functions from the C and C++ standard (template) libraries.\n",
    "\n",
    "To access functions in the C stdlib `math.h` header file, we simply do\n",
    "\n",
    "```\n",
    "from libc.math cimport exp, log, sqrt\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NumPy buffer special declaration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -a\n",
    "\n",
    "cimport numpy as cnp\n",
    "from libc.math cimport log as clog\n",
    "\n",
    "def shannon_entropy_v2(cnp.ndarray[double] p_x):\n",
    "    cdef double res = 0.0\n",
    "    cdef int n = p_x.shape[0]\n",
    "    cdef int i\n",
    "    for i in range(n):\n",
    "        res += p_x[i] * clog(p_x[i])\n",
    "    return -res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "shannon_entropy_v2(pmf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The `cnp.ndarray[double]` syntax\n",
    "\n",
    "The `cnp.ndarray[double]` declares a NumPy array _buffer_ object.  Cython knows how to interact with this array-like object efficiently.  The `double` in square brackets is the (scalar) dtype of the array elements."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Turn off boundschecking and wraparound checking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -a\n",
    "\n",
    "cimport cython\n",
    "cimport numpy as cnp\n",
    "from libc.math cimport log\n",
    "\n",
    "@cython.boundscheck(False)\n",
    "@cython.wraparound(False)\n",
    "def shannon_entropy_v3(cnp.ndarray[double] p_x):\n",
    "    cdef double res = 0.0\n",
    "    cdef int n = p_x.shape[0]\n",
    "    cdef int i\n",
    "    for i in range(n):\n",
    "        res += p_x[i] * log(p_x[i])\n",
    "    return -res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "shannon_entropy_v3(pmf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The `cython` cimported magic module\n",
    "\n",
    "Cython allows us to control its compile-time semantics and behavior via the magic `cython` module that we `cimport`.\n",
    "\n",
    "We can then use Cython directives like:\n",
    "\n",
    "```python\n",
    "@cython.boundscheck(False)\n",
    "@cython.wraparound(False)\n",
    "def func(...):\n",
    "  ...\n",
    "```\n",
    "\n",
    "To control how code is generated.\n",
    "\n",
    "In this case, we're telling cython to not generate code that does boundschecking or wraparound checking (negative indexing)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Typed memoryview syntax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -a\n",
    "\n",
    "cimport cython\n",
    "from libc.math cimport log\n",
    "\n",
    "@cython.boundscheck(False)\n",
    "@cython.wraparound(False)\n",
    "def shannon_entropy_mv(double[::1] p_x):\n",
    "    cdef double res = 0.0\n",
    "    cdef int n = p_x.shape[0]\n",
    "    cdef int i\n",
    "    for i in range(n):\n",
    "        res += p_x[i] * log(p_x[i])\n",
    "    return -res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "shannon_entropy_mv(pmf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Typed memoryview syntax\n",
    "\n",
    "The declaration\n",
    "\n",
    "```python\n",
    "def shannon_entropy_mv(double[::1] p_x):\n",
    "   ...\n",
    "```\n",
    "\n",
    "Declares `p_x` to be a one dimensional contiguous typed memoryview -- a Cython only construct that's compatible with NumPy arrays and PEP 3118 buffer objects.  We include it here because you may see this synax used in other situations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fused types example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython -a\n",
    "\n",
    "cimport cython\n",
    "from libc.math cimport log\n",
    "\n",
    "@cython.boundscheck(False)\n",
    "@cython.wraparound(False)\n",
    "def shannon_entropy_mv(cython.floating[::1] p_x):\n",
    "    cdef double res = 0.0\n",
    "    cdef int n = p_x.shape[0]\n",
    "    cdef int i\n",
    "    for i in range(n):\n",
    "        if p_x[i] > 0.0: # Have to guard against underflow...\n",
    "            res += p_x[i] * log(p_x[i])\n",
    "    return -res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(shannon_entropy_mv(pmf.astype('f8')))\n",
    "print(shannon_entropy_mv(pmf.astype('f4')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}