#!/usr/bin/env python
# coding: utf-8

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

# # 5.1. Accelerating pure Python code with Numba and Just-In-Time compilation

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

# In[ ]:


import numpy as np


# We initialize the simulation and generate the grid
# in the complex plane.

# In[ ]:


size = 200
iterations = 100


# ## Pure Python version

# The following function generates the fractal.

# In[ ]:


def mandelbrot_python(m, size, iterations):
    for i in range(size):
        for j in range(size):
            c = -2 + 3./size*j + 1j*(1.5-3./size*i)
            z = 0
            for n in range(iterations):
                if np.abs(z) <= 10:
                    z = z*z + c
                    m[i, j] = n
                else:
                    break


# In[ ]:


m = np.zeros((size, size))
mandelbrot_python(m, size, iterations)


# In[ ]:


import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')
plt.imshow(np.log(m), cmap=plt.cm.hot,);
plt.xticks([]); plt.yticks([]);


# In[ ]:


get_ipython().run_cell_magic('timeit', 'm = np.zeros((size, size))', 'mandelbrot_python(m, size, iterations)\n')


# ## Numba version

# We first import Numba.

# In[ ]:


import numba
from numba import jit, complex128


# Now, we just add the `@jit` decorator to the exact same function.

# In[ ]:


@jit(locals=dict(c=complex128, z=complex128))
def mandelbrot_numba(m, size, iterations):
    for i in range(size):
        for j in range(size):
            c = -2 + 3./size*j + 1j*(1.5-3./size*i)
            z = 0
            for n in range(iterations):
                if abs(z) <= 10:
                    z = z*z + c
                    m[i, j] = n
                else:
                    break


# In[ ]:


m = np.zeros((size, size))
mandelbrot_numba(m, size, iterations)


# In[ ]:


get_ipython().run_cell_magic('timeit', 'm = np.zeros((size, size))', 'mandelbrot_numba(m, size, iterations)\n')


# The Numba version is 250 times faster than the pure Python version here!

# > You'll find all the explanations, figures, references, and much more in the book (to be released later this summer).
# 
# > [IPython Cookbook](http://ipython-books.github.io/), by [Cyrille Rossant](http://cyrille.rossant.net), Packt Publishing, 2014 (500 pages).