matrix-factorization.rb

require 'nmatrix'
require 'pp'

###############################################################################

=begin
@INPUT:
    r     : a matrix to be factorized, dimension n x m
    p     : an initial matrix of dimension n x k
    q     : an initial matrix of dimension m x k
    ak    : the number of latent features
    steps : the maximum number of steps to perform the optimisation
    alpha : the learning rate
    beta  : the regularization parameter
@OUTPUT:
    the final matrices p and q
=end

def matrix_factorization(r, p, q, ak, steps=5000, alpha=0.0002, beta=0.02)
    q = q.transpose
    0.upto steps-1 do |step|
        0.upto r.length-1 do |i|
            0.upto r[i].length-1 do |j|
                if r[i][j] > 0
                    eij = r[i][j] - (p.row(i).dot q.column(j))[0]
                    0.upto ak-1 do |k|
                        p[i,k] = p[i,k] + alpha * (2 * eij * q[k,j] - beta * p[i,k])
                        q[k,j] = q[k,j] + alpha * (2 * eij * p[i,k] - beta * q[k,j])
                    end
                end
            end
        end
        e = 0
        0.upto r.length-1 do |i|
            0.upto r[i].length-1 do |j|
                if r[i][j] > 0
                    e = e + ((r[i][j] - (p.row(i).dot q.column(j))[0]) ** 2)
                    0.upto ak-1 do |k|
                        e = e + (beta/2) * ((p[i,k] ** 2) + (q[k,j] ** 2))
                    end
                end
            end
        end
        if e < 0.001
            break
        end
    end
    return p, q.transpose
end    

r = NMatrix[
            [5,3,0,1],
            [4,0,0,1],
            [1,1,0,5],
            [1,0,0,4],
            [0,1,5,4],
           ]

r = r.to_a

n = r.length
m = r[0].length
k = 2

p = NMatrix.random([n, k])
q = NMatrix.random([m, k])

np, nq = matrix_factorization(r, p, q, k)
nr = np.dot nq.transpose
pp r
pp nr