# Same as mendelbrot.py except it is in the form of a function
# so that we can call it multiple times.
#
import numpy as np
import matplotlib.pyplot as plt


ITERATIONS = 10
DENSITY = 1000 # warning: execution speed decreases with square of DENSITY

def mendelbrot():
    x_min, x_max = -2, 1
    y_min, y_max = -1.5, 1.5
    
    x, y = np.meshgrid(np.linspace(x_min, x_max, DENSITY),
                       np.linspace(y_min, y_max, DENSITY))
    
    c = x + 1j*y # complex grid
    z = c.copy()
    fractal = np.zeros(z.shape, dtype=np.uint8) + 255
    
    for n in range(ITERATIONS):
        print "Iteration %d" % n
        
        # --- Uncomment to see different sets ---
        
        # Tricorn
        # z = z.conj()
        
        # Burning ship
        # z = abs(z.real) + 1j*abs(z.imag)
        
        # ---
        
        # Leave the lines below in place
        z *= z
        z += c
        
        mask = (fractal == 255) & (abs(z) > 10)
        fractal[mask] = 254 * n / float(ITERATIONS)
        
    plt.imshow(np.log(fractal), cmap=plt.cm.hot,
               extent=(x_min, x_max, y_min, y_max))
    plt.title('Mandelbrot Set')
    plt.xlabel('Re(z)')
    plt.ylabel('Im(z)')
    plt.show()
    return 0

mendelbrot()