interactive Mandelbrot numba example

mandel.py

"""From the numba examples
(https://github.com/numba/numba/blob/master/examples/mandel.py),
but tweaked to recalculate on the fly as the viewport changes"""
from numba import autojit
import numpy as np
from pylab import imshow, jet, show, ion
import matplotlib.pyplot as plt
from matplotlib import rcParams

rcParams['image.origin'] = 'lower'

@autojit
def mandel(x, y, max_iters):
    """
    Given the real and imaginary parts of a complex number,
    determine if it is a candidate for membership in the Mandelbrot
    set given a fixed number of iterations.
    """
    i = 0
    c = complex(x,y)
    z = 0.0j
    for i in range(max_iters):
        z = z*z + c
        if (z.real*z.real + z.imag*z.imag) >= 4:
            return i

    return 255

@autojit
def create_fractal(min_x, max_x, min_y, max_y, image, iters):
    height = image.shape[0]
    width = image.shape[1]

    pixel_size_x = (max_x - min_x) / width
    pixel_size_y = (max_y - min_y) / height
    for x in range(width):
        real = min_x + x * pixel_size_x
        for y in range(height):
            imag = min_y + y * pixel_size_y
            color = mandel(real, imag, iters)
            image[y, x] = color

    return image

def update_axes(ax):
    xlim = ax.get_xlim()
    ylim = ax.get_ylim()
    trans = ax.transAxes
    ext = trans.transform([[1, 1]]) - trans.transform([[0, 0]])
    nx = ext[0, 0]
    ny = ext[0, 1]
    image = np.zeros((ny, nx), dtype=np.uint8)
    create_fractal(xlim[0], xlim[1], ylim[0], ylim[1], image, 255)
    ax.clear()
    ax.imshow(image, extent=[xlim[0], xlim[1], ylim[0], ylim[1]], zorder=0)
    ax.figure.canvas.draw()

image = np.zeros((500, 750), dtype=np.uint8)
artist = plt.imshow(create_fractal(-2.0, 1.0, -1.0, 1.0, image, 255),
                    extent=[-2.0, 1.0, -1.0, 1.0])

ax = plt.gca()
ax.figure.canvas.mpl_connect('button_release_event', lambda x: update_axes(ax))
ax.set_title('Zoom in')
plt.show()

Nice! What does autojit do?

I think you should replace xranges with np.arange and the for loops with Numpy functions as they are much faster.
The loops in create_fractal can be wrapped in a utility function and then use np.apply_over_axes.

@peterhil take a look at the Numba project, where autojit comes from. It compiles python functions into LLVM, so the loops run at ~C speed. The code runs about 1000x faster after applying the autojit decorator. This code is lifted from the numba repo, with the update_axes function added for fun times.

Somehow I get mirrored zooming in this, unless I swap the ylim parameters in line 51, create_fractal(xlim[0], xlim[1], ylim[1], ylim[0], image, 255)

@petrushy You probably have 'image.origin = upper' somewhere in your matplotlibrc file (or maybe that's the default, and it's overridden to lower in mine). Anyways, I think the new version should resolve that problem.