simrank in python

simrank.py

import numpy as np


def read_matrix(fname):
    x,y,z = [],[],[]
    with open(fname) as f:
        max_v, max_u = -1, -1
        for l in open(fname):
            u,v,c = l.strip().split()
            u,v,c = (int(u), int(v), float(c))
            max_u = max(u, max_u)
            max_v = max(v, max_v)
            x.append(u)
            y.append(v)
            z.append(c)
    
    m = np.zeros([max_v + 1, max_u + 1])
    m[y,x] = z
    return m

def print_matrix(title, m):
    print title
    print m
    print 


def load_graph(fname):
    P = read_matrix(fname)

    # set P[i,j] = Spread[i] * NormalizedWeight(P[i,j]) 
    col_sums, row_sums = P.sum(axis=0), P.sum(axis=1)
    Q = P / P.sum(axis=0)[np.newaxis, :]
    A = np.transpose(P / P.sum(axis=1)[:, np.newaxis])
    
    W = np.vstack([
        np.column_stack([np.zeros([A.shape[0], A.shape[0]]), A]),
        np.column_stack([Q, np.zeros([Q.shape[0], Q.shape[0]])]),
        ])

    print_matrix("P =", P)
    print_matrix("A =", A)
    print_matrix("Q =", Q)
    print_matrix("W =", W)

    return (np.matrix(Q), np.matrix(A), np.matrix(W))


def boosting(S):
    for k in xrange(S.shape[0]):
        S[k,k] = 1.0
    return S

def simrank(W, C=1.0, eps=1e-8, loop=1000):
    Wt = np.transpose(W)
    S = boosting(Wt * W * C)

    for i in xrange(loop):
        S0 = np.copy(S)
        S = boosting(Wt * S0 * W * C)
            
        #print_matrix("S" + str(i) + " = ", S)
        if np.allclose(S, S0, atol=eps):
            break
    return S


def simrank_qa(Q, A, Cq=.9, Ca=.8, eps=1e-3, loop=20):
    Qt = np.transpose(Q)
    At = np.transpose(A)
    Sq = boosting(Qt * Q * Cq)
    Sa = boosting(At * A * Ca)

    for i in xrange(loop):
        S0 = np.copy(Sq)
        Sq = boosting(Qt * Sa * Q * Cq)
        Sa = boosting(At * Sq * A * Ca)

        #print_matrix("S" + str(i) + " = ", S)
        if np.allclose(Sq, S0, atol=eps):
            break
    return Sq, Sa


if __name__ == "__main__":
    from tornado.options import options,define
    from tornado.options import parse_command_line
    define("c", default=20)
    define("e", default=1e-4)
    define("da", default=.9)
    define("dq", default=.9)
    define("dw", default=.9)

    args = parse_command_line()

    np.set_printoptions(linewidth=200, precision=5, suppress=True)
    Q, A, W = load_graph(args[0])
    sq, sa = simrank_qa(Q, A, options.dq, options.da, options.e, options.c)
    print_matrix("final Sq:", sq)
    print_matrix("final Sa:", sa)
    #print_matrix("final :", simrank(W, options.dw, options.e, options.c))