%matplotlib inline
from itertools import product
import numpy
from matplotlib import pyplot

# grid parameters
res = 500, 500  # resolution of the grid (n x n)
re_min, re_max = -1., 1.
im_min, im_max = -1., 1.
re_step_size = float(re_max - re_min)/res[0]
im_step_size = float(im_max - im_min)/res[0]

# Construct the complex coordinate plane.
arr = numpy.empty(res, dtype=complex)
for i in range(res[0]):
    for j in range(res[1]):
        arr[i, j] = complex(i*im_step_size + im_min,
                            j*re_step_size + re_min)
        
pyplot.matshow(numpy.absolute(arr))

@numpy.vectorize
def julia(z):
    # julia set parameters
    c = complex(0., .75)
    n_max = 100
    z_max = 10
    n = 0
    while abs(z) < z_max and n < n_max:
        z = z**2 + c
        n += 1
    return n


arr = julia(arr)
pyplot.matshow(numpy.absolute(arr))

from distarray import Context

# context is our connection to the engines, and emulates
# the numpy module object. With no arguments it uses all
# of  the available ipython parallel engines.
context = Context()

# dist describes the distribution of our array.
dist={0: 'c', 1: 'c'}
# 'c': cyclic, 'b': block

distarray = context.empty(res, dtype=complex, dist=dist)

@context.local
def draw_coord(arr, re_min, re_max, im_min, im_max, res):
    """Draw the complex coordinate plane."""
    re_step = float(re_max - re_min) / res[0]
    im_step = float(im_max - im_min) / res[1]
    
    for i in arr.maps[0].global_index:
        for j in arr.maps[1].global_index:
            arr[i, j] = complex(re_min + re_step*i,
                                im_min + im_step*j)
    return arr

distarray = draw_coord(distarray, re_min, re_max, im_min, im_max, res)
pyplot.matshow(numpy.absolute(distarray.tondarray()))

@context.vectorize
def julia(z):
    # julia set parameters
    c = complex(0., .75)
    n_max = 100
    z_max = 10
    n = 0
    while abs(z) < z_max and n < n_max:
        z = z**2 + c
        n += 1
    return n
distarray = julia(distarray)

pyplot.matshow(numpy.absolute(distarray.tondarray()))