如何在一个简单的函数上 Cythonize / 允许 numba.jit ：在网络中查找三角形

如何解决如何在一个简单的函数上 Cythonize / 允许 numba.jit ：在网络中查找三角形

背景故事：

我一直在寻找一种高性能的方法来查找网络中低于给定维度的派系（例如，所有 k

Networkx 非常慢；然而，networkit 有一个性能更高的解决方案，带有 Cython 后端。

不幸的是，networkit 没有列出所有派系 所有可能维度运行（据我所知）。它也只计数个三角形，但没有列出构成每个三角形的节点。因此，我正在编写自己的函数，该函数现在在下面实现了一种相当有效的方法。

问题：

我有下面的函数nk_triangles；然而，它抵抗了对 numba 或 Cython 的简单干扰。因此，我想看看是否有人在这些领域拥有更多专业知识，可以将其推向更快的速度。

我在这里制作了一个简单但完全可行的代码片段，其中包含感兴趣的功能：

import networkit as nk
import numba
from itertools import combinations
from urllib.request import urlopen
import tempfile

graph_url="https://raw.githubusercontent.com/networkit/networkit/master/input/tiny_02.graph"
big_graph_url="https://raw.githubusercontent.com/networkit/networkit/master/input/caidaRouterLevel.graph"
with tempfile.NamedTemporaryFile() as f:
    with urlopen(graph_url) as r:
        f.write(r.read())
    f.read()
    G = nk.readGraph(f.name,nk.Format.metis)

#@numba.jit
def nk_triangles(g):
    # Source:
    # https://cs.stanford.edu/~rishig/courses/ref/l1.pdf
    triangles = set()
    for node in g.iterNodes():
        ndeg = g.degree(node)

        neighbors = [neigh for neigh in g.iterNeighbors(node)
                     if (ndeg < g.degree(neigh)) or
                        ((ndeg == g.degree(neigh))
                          and node < neigh)]

        node_triangles = set({(node,*c): max(g.weight(u,v)
                                              for u,v in combinations([node,*c],2))
                              for c in combinations(neighbors,2)
                              if g.hasEdge(*c)})
        triangles = triangles.union(node_triangles)
    return triangles


tris = nk_triangles(G)
tris

可以切换 big_graph_url 以查看算法实际上是否表现得相当好。 （我的图表仍然比这大几个数量级）

就目前而言，这需要大约 40 分钟来计算我的机器（单线程 python 循环调用 networkit 和 itertools 中的 C 后端代码）。大网络中的三角形数为455,062。

解决方法

这是您的代码的 numpy 版本，为您的大图花费了大约 1 分钟。

%%time

graph_url="https://raw.githubusercontent.com/networkit/networkit/master/input/tiny_02.graph"
big_graph_url="https://raw.githubusercontent.com/networkit/networkit/master/input/caidaRouterLevel.graph"
with tempfile.NamedTemporaryFile() as f:
    with urlopen(big_graph_url) as r:
        f.write(r.read())
    f.read()
    G = nk.readGraph(f.name,nk.Format.METIS)

nodes = np.array(tuple(G.iterNodes()))
adjacency_matrix = nk.algebraic.adjacencyMatrix(G,matrixType='sparse').astype('bool')
degrees = np.sum(adjacency_matrix,axis=0)
degrees = np.array(degrees).reshape(-1)



def get_triangles(node,neighbors):
    buffer = neighbors[np.argwhere(triangle_condition(*np.meshgrid(neighbors,neighbors)))]
    triangles = np.empty((buffer.shape[0],buffer.shape[1]+1),dtype='int')
    triangles[:,0] = node
    triangles[:,1:] = buffer
    return triangles

def triangle_condition(v,w):
    upper = np.tri(*v.shape,-1,dtype='bool').T
    upper[np.where(upper)] = adjacency_matrix[v[upper],w[upper]]
    return upper

def nk_triangles():
    triangles = list()
    for node in nodes:
        ndeg = degrees[node]
        neighbors = nodes[adjacency_matrix[node].toarray().reshape(-1)]
        neighbor_degs = degrees[neighbors]
        neighbors = neighbors[(ndeg < neighbor_degs) | ((ndeg == neighbor_degs) & (node < neighbors))]
        if len(neighbors) >= 2:
            triangles.append(get_triangles(node,neighbors))
    return triangles

tris = np.concatenate(nk_triangles())
print('triangles:',len(tris))

给我

triangles: 455062
CPU times: user 50.6 s,sys: 375 ms,total: 51 s
Wall time: 52 s