import numpy as np
from numba import autojit, jit, double

def sum(arr):
    M, N = arr.shape
    sum = 0.0
    for i in range(M):
        for j in range(N):
            sum += arr[i,j]
    return sum
fastsum = jit('f8(f8[:,:])')(sum)
flexsum = autojit(sum)

arr2d = np.arange(600,dtype=float).reshape(20,30)
print sum(arr2d)
print fastsum(arr2d)
print flexsum(arr2d)
print flexsum(arr2d.astype(int))

%timeit sum(arr2d)

%timeit fastsum(arr2d)

393 / 2.88 # speedup

%timeit arr2d.sum() 

7.9 / 2.88 # even provides a speedup over general-purpose NumPy surm

@jit('void(f8[:,:],f8[:,:],f8[:,:])')
def filter(image, filt, output):
    M, N = image.shape
    m, n = filt.shape
    for i in range(m//2, M-m//2):
        for j in range(n//2, N-n//2):
            result = 0.0
            for k in range(m):
                for l in range(n):
                    result += image[i+k-m//2,j+l-n//2]*filt[k, l]
            output[i,j] = result

from scipy.misc import lena
import time
image = lena().astype('double')
filt = np.ones((15,15),dtype='double')
filt /= filt.sum()
output = image.copy()
filter(image, filt, output)
gray()
imshow(output)
start = time.time()
filter(image[:100,:100], filt, output[:100,:100])
fast = time.time() - start
start = time.time()
filter.py_func(image[:100,:100], filt, output[:100,:100])
slow = time.time() - start
print "Python: %f s; Numba: %f ms; Speed up is %f" % (slow, fast*1000, slow / fast)

@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**2 + c
        if abs(z)**2 >= 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

image = np.zeros((500, 750), dtype=np.uint8)
imshow(create_fractal(-2.0, 1.0, -1.0, 1.0, image, 20))
jet()

%timeit create_fractal(-2.0, 1.0, -1.0, 1.0, image, 20)

%timeit create_fractal.py_func(-2.0, 1.0, -1.0, 1.0, image, 20)

18.9/ 1.28 # speedup of compiling outer-loop (inner-loop mandel call is still optimized)

class MyClass(object):
    def mymethod(self, arg):
        return arg * 2
    
@autojit(locals=dict(mydouble=double)) # specify types for local variables
def call_method(obj):
    print obj.mymethod("hello")   # object result
    mydouble = obj.mymethod(10.2) # native double
    print mydouble * 2            # native multiplication
    
call_method(MyClass())

@autojit
def complex_support(real, imag):
    c = complex(real, imag)
    return (c ** 2).conjugate()

c = 2.0 + 4.0j
complex_support(c.real, c.imag), (c**2).conjugate()

from numba import struct, jit2
record_type = struct([('x', double), ('y', double)])
record_dtype = record_type.get_dtype()
a = np.array([(1.0, 2.0), (3.0, 4.0)], dtype=record_dtype)

@jit2(argtypes=[record_type[:]])
def pyth(data):
    result = np.empty_like(data, dtype=np.float64) # return types of numpy functions are inferred
    for i in range(data.shape[0]):
        result[i] = np.sqrt(data[i].x ** 2 + data[i].y ** 2)

    return result

print pyth(a)

print pyth.signature # inspect function signature, note inferred return type

[line for line in str(pyth.lfunc).splitlines() if 'sqrt' in line] # note native math calls