# necessary imports
%matplotlib inline
from pylab import linspace, zeros, imshow, complex64, subplot, figure

def f(z, n=100):
    f = 0
    for i in range(n):
        f = f**2 + z
    return f

num = 200
grid_x = linspace(-1.7, 0.6, num=num)
grid_y = linspace(-1.4, 1.4, num=num)

data = zeros((grid_x.size, grid_y.size), dtype=complex64)

%%time
for ind_x, x in enumerate(grid_x):
    for ind_y, y in enumerate(grid_y):
        data[ind_x, ind_y] = f(x + 1j * y)

extent = (min(grid_x), max(grid_x), min(grid_y), max(grid_y))
imshow(abs(data.T)**2 < 1000, cmap='gray',  extent=extent)

from IPython.html.widgets import interactive

def plot_vmax(vmax=0.2, cmap='hot'):
    imshow(abs(data.T), cmap=cmap,  vmax=vmax, extent=extent)

interactive(plot_vmax, 
            vmax=(0.1, 2, 0.05), 
            cmap=('hot', 'jet'))

figure(figsize=(10, 10))
for sub, vmax in enumerate([0.1, 0.45, 1.05, 2.]):
    subplot(2, 2, sub + 1)
    plot_vmax(vmax, 'jet')

figure(figsize=(10, 10))
for sub, vmax in enumerate([0.1, 0.45, 1.05, 2.]):
    subplot(2, 2, sub + 1)
    plot_vmax(vmax, 'hot')