gistfile1.py

from numpy import zeros
from scipy import weave
import numpy as np
import time
import timeit
import scipy.weave as weave
import pyximport
pyximport.install(setup_args={'include_dirs':[np.get_include()]})

dx = 0.1
dy = 0.1
dx2 = dx*dx
dy2 = dy*dy


def py_update(u):
    nx, ny = u.shape
    for i in xrange(1,nx-1):
        for j in xrange(1, ny-1):
            u[i,j] = ((u[i+1, j] + u[i-1, j]) * dy2 + (u[i, j+1] + u[i, j-1]) * dx2) / (2*(dx2+dy2))

def num_update(u):
    u[1:-1,1:-1] = ((u[2:,1:-1]+u[:-2,1:-1])*dy2 +(u[1:-1,2:] + u[1:-1,:-2])*dx2) / (2*(dx2+dy2))

def weave_update(u):
    code = """
    int i, j;
    for (i=1; i<Nu[0]-1; i++) {
       for (j=1; j<Nu[1]-1; j++) {
           U2(i,j) = ((U2(i+1, j) + U2(i-1, j))*dy2 + \
                       (U2(i, j+1) + U2(i, j-1))*dx2) / (2*(dx2+dy2));
       }
    }
    """
    weave.inline(code, ['u', 'dx2', 'dy2'])

def calc(N, Niter=100, func=num_update, args=()):
    u = zeros([N, N])
    u[0] = 1
    #print "u:\n",u
    for i in range(Niter):
        func(u,*args)
    return u

if __name__=='__main__':

    t = timeit.Timer('pureLAPLACE.calc(100,func=pureLAPLACE.py_update)','import pureLAPLACE')
    print "Pure python time:",min(t.repeat(1,10))

    t = timeit.Timer('pureLAPLACE.calc(100)','import pureLAPLACE')
    print "Numpy time:",min(t.repeat(5,10))

    t = timeit.Timer('pureLAPLACE.calc(100,func=cy_update,args=(0.01,0.01))','import pureLAPLACE;from laplace import cy_update')
    print "Cython time:",min(t.repeat(5,10))

    t = timeit.Timer('pureLAPLACE.calc(100,func=pureLAPLACE.weave_update)','import pureLAPLACE')
    print "scipy.weave time:",min(t.repeat(5,10))