{
 "metadata": {
  "name": "",
  "signature": "sha256:b27bacf4cc6c3b83c0b9ffd97f19b3a230a741dc1815ac1e76c5f97417ff59dd"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "metadata": [],
     "source": [
      "> This is one of the 100 recipes of the [IPython Cookbook](http://ipython-books.github.io/), the definitive guide to high-performance scientific computing and data science in Python.\n"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# 5.1. Accelerating pure Python code with Numba and Just-In-Time compilation"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "In this example, we first write a pure Python version of a function that generates a Mandelbrot fractal. Then, we use Numba to compile it dynamically to native code."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import numpy as np"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "We initialize the simulation and generate the grid\n",
      "in the complex plane."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "size = 200\n",
      "iterations = 100"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "## Pure Python version"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "The following function generates the fractal."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def mandelbrot_python(m, size, iterations):\n",
      "    for i in range(size):\n",
      "        for j in range(size):\n",
      "            c = -2 + 3./size*j + 1j*(1.5-3./size*i)\n",
      "            z = 0\n",
      "            for n in range(iterations):\n",
      "                if np.abs(z) <= 10:\n",
      "                    z = z*z + c\n",
      "                    m[i, j] = n\n",
      "                else:\n",
      "                    break"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "m = np.zeros((size, size))\n",
      "mandelbrot_python(m, size, iterations)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import matplotlib.pyplot as plt\n",
      "%matplotlib inline\n",
      "plt.imshow(np.log(m), cmap=plt.cm.hot,);\n",
      "plt.xticks([]); plt.yticks([]);"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%%timeit m = np.zeros((size, size))\n",
      "mandelbrot_python(m, size, iterations)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "## Numba version"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "We first import Numba."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import numba\n",
      "from numba import jit, complex128"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Now, we just add the `@jit` decorator to the exact same function."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "@jit(locals=dict(c=complex128, z=complex128))\n",
      "def mandelbrot_numba(m, size, iterations):\n",
      "    for i in range(size):\n",
      "        for j in range(size):\n",
      "            c = -2 + 3./size*j + 1j*(1.5-3./size*i)\n",
      "            z = 0\n",
      "            for n in range(iterations):\n",
      "                if abs(z) <= 10:\n",
      "                    z = z*z + c\n",
      "                    m[i, j] = n\n",
      "                else:\n",
      "                    break"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "m = np.zeros((size, size))\n",
      "mandelbrot_numba(m, size, iterations)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%%timeit m = np.zeros((size, size))\n",
      "mandelbrot_numba(m, size, iterations)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "The Numba version is 250 times faster than the pure Python version here!"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "> You'll find all the explanations, figures, references, and much more in the book (to be released later this summer).\n",
      "\n",
      "> [IPython Cookbook](http://ipython-books.github.io/), by [Cyrille Rossant](http://cyrille.rossant.net), Packt Publishing, 2014 (500 pages)."
     ]
    }
   ],
   "metadata": {}
  }
 ]
}