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

# # Mandelbrot Set
# 
# ## Author [Jean-François Puget](https://www.ibm.com/developerworks/community/blogs/jfp/?lang=en)
# 
# The code used in this notebook is discussed at length in [My Christmas Gift](https://www.ibm.com/developerworks/community/blogs/jfp/entry/my_christmas_gift )

# First, let's load useful extensions and import useful modules

# In[1]:


import numpy as np
from numba import jit
from matplotlib import pyplot as plt
from matplotlib import colors
get_ipython().run_line_magic('matplotlib', 'inline')


# ## The Escape Method
# 
# We first compute Mandelbrodt set with the escape method. 

# In[2]:


@jit
def mandelbrot(c,maxiter):
    z = c
    for n in range(maxiter):
        if abs(z) > 2:
            return n
        z = z*z + c
    return 0

@jit
def mandelbrot_set(xmin,xmax,ymin,ymax,width,height,maxiter):
    r1 = np.linspace(xmin, xmax, width)
    r2 = np.linspace(ymin, ymax, height)
    n3 = np.empty((width,height))
    for i in range(width):
        for j in range(height):
            n3[i,j] = mandelbrot(r1[i] + 1j*r2[j],maxiter)
    return (r1,r2,n3)


# Using Numba makes it rather fast.  Note that time is roughly proportional to the max number of iterations.

# In[3]:


get_ipython().run_line_magic('timeit', 'mandelbrot_set(-2.0,0.5,-1.25,1.25,1000,1000,80)')
get_ipython().run_line_magic('timeit', 'mandelbrot_set(-2.0,0.5,-1.25,1.25,1000,1000,800)')


# ## Display
# 
# Let's plot the result.  We take as input the size of the image in inches, and we compute the number of pixels in each axis using the standard 72 dots per inch used for web images.

# In[4]:


image_counter = 30

def save_image(fig):
    global image_counter
    filename = "mandelbrodt_%d.png" % image_counter
    image_counter += 1
    fig.savefig(filename)


# In[5]:


def mandelbrot_image(xmin,xmax,ymin,ymax,width=10,height=10,maxiter=256):
    dpi = 72
    img_width = dpi * width
    img_height = dpi * height
    x,y,z = mandelbrot_set(xmin,xmax,ymin,ymax,img_width,img_height,maxiter)
    
    fig, ax = plt.subplots(figsize=(width, height),dpi=72)
    ticks = np.arange(0,img_width,3*dpi)
    x_ticks = xmin + (xmax-xmin)*ticks/img_width
    plt.xticks(ticks, x_ticks)
    y_ticks = ymin + (ymax-ymin)*ticks/img_width
    plt.yticks(ticks, y_ticks)
    
    ax.imshow(z.T,origin='lower') 
    
    save_image(fig)


# Let's see what this gives

# In[6]:


mandelbrot_image(-2.0,0.5,-1.25,1.25)


# Let's try different color maps.

# In[7]:


def mandelbrot_image(xmin,xmax,ymin,ymax,width=10,height=10,maxiter=80,cmap='jet'):
    dpi = 72
    img_width = dpi * width
    img_height = dpi * height
    x,y,z = mandelbrot_set(xmin,xmax,ymin,ymax,img_width,img_height,maxiter)
    
    fig, ax = plt.subplots(figsize=(width, height),dpi=72)
    ticks = np.arange(0,img_width,3*dpi)
    x_ticks = xmin + (xmax-xmin)*ticks/img_width
    plt.xticks(ticks, x_ticks)
    y_ticks = ymin + (ymax-ymin)*ticks/img_width
    plt.yticks(ticks, y_ticks)
    ax.set_title(cmap)
    
    ax.imshow(z.T,cmap=cmap,origin='lower') 
    
    save_image(fig)


# In[8]:


mandelbrot_image(-2.0,0.5,-1.25,1.25,cmap='hot')


# That's a bit dark, because most of the pixels have a low value.

# In[9]:


r1,r2,z = mandelbrot_set(-2.0,0.5,-1.25,1.25,100,100,100)
plt.hist(z)
plt.show()


# We can use a power norm to spread the values more evenly.

# In[10]:


z3 = z ** 0.3
plt.hist(z3)
plt.show()


# In[11]:


def mandelbrot_image(xmin,xmax,ymin,ymax,width=10,height=10,maxiter=80,cmap='jet'):
    dpi = 72
    img_width = dpi * width
    img_height = dpi * height
    x,y,z = mandelbrot_set(xmin,xmax,ymin,ymax,img_width,img_height,maxiter)
    z = z ** 0.3
    
    fig, ax = plt.subplots(figsize=(width, height),dpi=72)
    ticks = np.arange(0,img_width,3*dpi)
    x_ticks = xmin + (xmax-xmin)*ticks/img_width
    plt.xticks(ticks, x_ticks)
    y_ticks = ymin + (ymax-ymin)*ticks/img_width
    plt.yticks(ticks, y_ticks)
    ax.set_title(cmap)
    
    ax.imshow(z.T,cmap=cmap,origin='lower') 


# In[12]:


mandelbrot_image(-2.0,0.5,-1.25,1.25,cmap='hot')


# An equivalent way of doing it is to use a normalize object

# In[13]:


def mandelbrot_image(xmin,xmax,ymin,ymax,width=10,height=10,maxiter=80,cmap='jet',gamma=0.3):
    dpi = 72
    img_width = dpi * width
    img_height = dpi * height
    x,y,z = mandelbrot_set(xmin,xmax,ymin,ymax,img_width,img_height,maxiter)
    
    fig, ax = plt.subplots(figsize=(width, height),dpi=72)
    ticks = np.arange(0,img_width,3*dpi)
    x_ticks = xmin + (xmax-xmin)*ticks/img_width
    plt.xticks(ticks, x_ticks)
    y_ticks = ymin + (ymax-ymin)*ticks/img_width
    plt.yticks(ticks, y_ticks)
    ax.set_title(cmap)
    
    norm = colors.PowerNorm(gamma)
    ax.imshow(z.T,cmap=cmap,origin='lower',norm=norm)  
    
    save_image(fig)


# In[14]:


mandelbrot_image(-2.0,0.5,-1.25,1.25,cmap='hot')


# ## Antialiasing
# 
# This is nice but it shows one limit of this technique: the colors aren't evolving smoothly.  
# The followings function computes a floating point value for each point instead of an integer.  Math for it can be found in my [blog]()

# In[15]:


@jit
def mandelbrot(z,maxiter,horizon,log_horizon):
    c = z
    for n in range(maxiter):
        az = abs(z)
        if az > horizon:
            return n - np.log(np.log(az))/np.log(2) + log_horizon
        z = z*z + c
    return 0

@jit
def mandelbrot_set(xmin,xmax,ymin,ymax,width,height,maxiter):
    horizon = 2.0 ** 40
    log_horizon = np.log(np.log(horizon))/np.log(2)
    r1 = np.linspace(xmin, xmax, width)
    r2 = np.linspace(ymin, ymax, height)
    n3 = np.empty((width,height))
    for i in range(width):
        for j in range(height):
            n3[i,j] = mandelbrot(r1[i] + 1j*r2[j],maxiter,horizon, log_horizon)
    return (r1,r2,n3)


# In[16]:


mandelbrot_image(-2.0,0.5,-1.25,1.25,cmap='hot')


# Much better!  The image is antialiased
# 
# ## Zooming
# 
# Let's zoom on the valley between the two largest black areas

# In[17]:


mandelbrot_image(-0.8,-0.7,0,0.1,cmap='hot')


# We see that the precision isn't great.  We need to augment the number of iterations.

# In[18]:


mandelbrot_image(-0.8,-0.7,0,0.1,cmap='hot',maxiter=1024)


# Win.

# Let's zoom further on the 'sea horses on the right' .  We need to further increase the number of iteraitons to get a crisper result.

# In[19]:


mandelbrot_image(-0.755,-0.745,0.06,0.07,cmap='hot',maxiter=2048)


# Let's zoom on one sea horse

# In[20]:


mandelbrot_image(-0.75,-0.747,0.063,0.066,cmap='hot',maxiter=2048)


# Let's zoom on its tail

# In[21]:


mandelbrot_image(-0.749,-0.748,0.065,0.066,cmap='hot',maxiter=2048)


# Let's zoom on the black dot at the bottom left

# In[22]:


mandelbrot_image(-0.74877,-0.74872,0.06505,0.06510,cmap='hot',maxiter=2048)


# In[23]:


mandelbrot_image(-0.74877,-0.74872,0.065053,0.065103,cmap='gnuplot2',maxiter=2048,gamma=0.3,width=14,height=14)


# Guess what, the black dot is similar to the whole Mandelbrodt set itself!  It is this property of self similarity that defines fractal shapes.

# Just for eye pleasure, here is the Mandelbrodt set with another color map I like a lot. It is only available if you upgrade matplotlib to its version 1.5.  You can do it with 
# 
# `pip install --upgrade matplotlib`

# In[24]:


mandelbrot_image(-2.0,0.5,-1.25,1.25,cmap='gnuplot2',maxiter=80,width=14,height=14)


# You can now play with this code using other color maps.  They can be seen at
# 
# http://matplotlib.org/examples/color/colormaps_reference.html