Testing numba's autojit decorator on simple code for matrix factorization.

numba_test.py

# -*- coding: utf-8 -*-
#!/usr/bin/python
#
# Modified example of Albert Au Yeung's on matrix factorization:
# http://www.quuxlabs.com/blog/2010/09/matrix-factorization-a-simple-tutorial-and-implementation-in-python/#source-code

import numpy as np
from numba.decorators import autojit, jit


def run():
    R = [
         [5,3,0,1],
         [4,0,0,1],
         [1,1,0,5],
         [1,0,0,4],
         [0,1,5,4],
        ]

    R = np.array(R, dtype=np.float)
    N, M = R.shape
    K = 2
    P = np.random.rand(N,K)
    Q = np.random.rand(M,K)
    nP, nQ, steps, error = matrix_factorization(R, P, Q, K)
    print 'Factorization complete in %i steps with a total erorr of %.3f' \
           % (steps, error)
    nR = np.dot(nP, nQ.T)
    return nR


#@autojit
def matrix_factorization(R, P, Q, K):
    steps = 5000
    alpha = 0.0002
    beta = 0.02 / 2.0
    Q = Q.T
    n, m = R.shape
    step = 0.0
    for step in xrange(steps):
        for i in xrange(n):
            for j in xrange(m):
                if R[i][j] > 0:
                    eij = R[i][j] - np.dot(P[i,:], Q[:,j])
                    for k in xrange(K):
                        P[i][k] += alpha * (2 * eij * Q[k][j] - beta * P[i][k])
                        Q[k][j] += alpha * (2 * eij * P[i][k] - beta * Q[k][j])
        e = 0.0
        for i in xrange(n):
            for j in xrange(m):
                if R[i][j] > 0:
                    e += (R[i][j] - np.dot(P[i,:], Q[:,j]))**2
                    for k in range(K):
                        e += beta * (P[i][k]**2 + Q[k][j]**2)
    
        if e < 0.001:
            break

    return P, Q.T, step, e