A O(n^3) SimRank algorithm with caching

simrank.py

import networkx as nx
import numpy as np

def simrank_fast(G, c=0.9, max_iter=100, epsilon=1e-4):
    if type(G) == nx.MultiGraph or type(G) == nx.MultiDiGraph:
        raise Exception("simrank() not defined for graphs with multiedges.")

    if G.is_directed():
        raise Exception("simrank() not defined for directed graphs.")
        
    def _nnzcsr(mat, idx):
        return mat.indices[mat.indptr[idx]:mat.indptr[idx+1]] 
        
    gmat = nx.to_scipy_sparse_matrix(G)
    n = gmat.shape[0]
    degrees = np.array(np.sum(gmat, axis=0))[0]
    smat = np.eye(n)
    partials = np.zeros((n,n))
    for _ in range(max_iter):
        for i in range(n):
            partials[i, :] = np.sum(smat[_nnzcsr(gmat, i)], axis=0)
        smat1 = np.eye(n)
        for i in range(n):
            for j in range(n):
                if i == j: continue
                smat1[i, j] = c/(degrees[i]*degrees[j])*np.sum(partials[i, _nnzcsr(gmat, j)])
        if np.linalg.norm(smat1-smat, np.inf) < epsilon:
            break
        smat=np.copy(smat1)
    return smat