#!/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.4. Accelerating Python code with Cython

# We use Cython to accelerate the generation of the Mandelbrot fractal.

# In[ ]:


import numpy as np


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

# In[ ]:


size = 200
iterations = 100


# ## Pure Python

# 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[ ]:


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


# ## Cython versions

# We first import Cython.

# In[ ]:


#%load_ext cythonmagic
get_ipython().run_line_magic('load_ext', 'Cython')


# ### Take 1

# First, we just add the %%cython magic.

# In[ ]:


get_ipython().run_cell_magic('cython', '-a', 'import numpy as np\n\ndef mandelbrot_cython(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\n')


# In[ ]:


get_ipython().run_cell_magic('timeit', '-n1 -r1 m = np.zeros((size, size), dtype=np.int32)', 'mandelbrot_cython(m, size, iterations)\n')


# Virtually no speedup.

# ### Take 2

# Now, we add type information, using memory views for NumPy arrays.

# In[ ]:


get_ipython().run_cell_magic('cython', '-a', 'import numpy as np\n\ndef mandelbrot_cython(int[:,::1] m, \n                      int size, \n                      int iterations):\n    cdef int i, j, n\n    cdef complex z, c\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 z.real**2 + z.imag**2 <= 100:\n                    z = z*z + c\n                    m[i, j] = n\n                else:\n                    break\n')


# In[ ]:


get_ipython().run_cell_magic('timeit', '-n1 -r1 m = np.zeros((size, size), dtype=np.int32)', 'mandelbrot_cython(m, size, iterations)\n')


# Interesting speedup!

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