mandelbrot.py

from numpy import empty


def mandelbrot_escape(x, y, n):
    """ Mandelbrot set escape time algorithm in real and complex components
    """
    z_x = x
    z_y = y
    for i in range(n):
        z_x, z_y = z_x**2 - z_y**2 + x, 2*z_x*z_y + y
        if z_x**2 + z_y**2 >= 4.0:
            break
    else:
        i = -1
    return i


def generate_mandelbrot(xs, ys, n):
    """ Generate a mandelbrot set """
    d = empty(shape=(len(ys), len(xs)))
    for j in range(len(ys)):
        for i in range(len(xs)):
            d[j, i] = mandelbrot_escape(xs[i], ys[j], n)
    return d


if __name__ == '__main__':
    import time
    from numpy import r_
    from matplotlib.pyplot import imshow, show, cm

    x = r_[-2:1:50j]
    y = r_[-1.5:1.5:50j]

    t = time.time()
    d = generate_mandelbrot(x, y, 100)
    print 'Execution time:', time.time() - t

    imshow(d, extent=[-2, 1, -1.5, 1.5], cmap=cm.gist_stern)
    show()