如何在graph-tools中高效实现6万节点图的直线性中心性计算

问题背景 #

需要实现直线性中心性：节点与其他所有节点的欧氏距离除以图距离之和，再乘以1/(N-1)（N为总节点数）。当前处理的图约含60000个节点，使用Dask并行遍历节点未降低时间复杂度，已将节点x/y坐标存入以节点编号为键的字典，不清楚如何利用graph-tools的OpenMP加速，当前实现代码如下：

# too slow
import math
import graph_tool as gt
from dask.distributed import client

def euclidean_dist(x1, y1, x2, y2):
    return math.sqrt((x1 - x2)**2 + (y1 - y2)**2)

def bravo(target, n, vertID_dict_scattered, network_dist):
    euclidean_distance = euclidean_dist(vertID_dict_scattered[n][0], vertID_dict_scattered[n][1], vertID_dict_scattered[target][0], vertID_dict_scattered[target][1])
    return euclidean_distance / network_dist

@dask.delayed
def alpha(G, n, vertID_dict_scattered):
    straightness = 0
    sp = gt.shortest_distance(G, n, weights=G.edge_properties["mm_len"])
    sp_scattered = sp

    if len(sp.get_array()) > 0 and len(G) > 1:
        for target, value in enumerate(sp):
            if n != target:
                network_dist = sp_scattered[target]
                straightness += bravo(target, n, vertID_dict_scattered, network_dist)
        straightness_df = straightness * (1.0 / (len(vertID_dict_scattered.keys()) - 1.0))
    else:
        straightness_df = 0
    
    return n, straightness_df

def split_list(lst, chunk_size):
    return [lst[i:i+chunk_size] for i in range(0, len(lst), chunk_size)]
        
def straightness_centrality(G, vertID_dict):
    chunked_list = split_list(list(G.iter_vertices()), 32)
    G_scattered = client.scatter(G)
    vertID_dict_scattered = client.scatter(vertID_dict)
    results=[]

    for chunk in chunked_list:   
        delayed_objs = [alpha(G_scattered, n, vertID_dict_scattered) for n in chunk]
        new_results = client.compute(delayed_objs)
        gathered_results = client.gather(new_results)
        results.append(gathered_results)
        
    return results

straightness_df = straightness_centrality(gtG, vertID_dict)

优化方案 #

1. 启用graph-tools原生OpenMP加速 #

graph-tools的shortest_distance原生支持OpenMP，只需设置线程数即可利用多核CPU：

import graph_tool as gt
gt.openmp_set_num_threads(8)  # 根据CPU核心数调整，比如8核设为8

同时，批量计算多个源节点的最短路径，比单个节点循环调用效率高很多，减少函数调用开销。

2. 向量化计算替代Python循环 #

用NumPy向量化操作代替遍历每个目标节点的Python循环，利用CPU SIMD指令大幅提升计算速度：
首先将坐标字典转换为NumPy数组（索引对应节点编号）：

import numpy as np

# 假设节点编号连续，若不连续需调整索引映射
coords = np.array([vertID_dict[n] for n in range(len(G))])

然后批量处理源节点的核心计算逻辑：

def alpha_batch(G, sources, coords):
    # 批量计算多个源节点的最短路径矩阵
    sp_matrix = gt.shortest_distance(G, source=sources, weights=G.edge_properties["mm_len"])
    N = len(G)
    results = []
    for i, n in enumerate(sources):
        sp = sp_matrix[i]
        # 排除自身节点
        mask = np.arange(N) != n
        # 向量化计算欧氏距离
        euclid_dist = np.sqrt((coords[n, 0] - coords[mask, 0])**2 + (coords[n, 1] - coords[mask, 1])**2)
        # 过滤不可达节点（图距离为inf）
        valid_mask = sp[mask] != np.inf
        # 向量化求和比值
        ratio_sum = np.sum(euclid_dist[valid_mask] / sp[mask][valid_mask])
        # 计算最终中心性值
        straightness = ratio_sum * (1.0 / (N - 1)) if N > 1 else 0
        results.append((n, straightness))
    return results

3. 优化Dask并行策略 #

• 减少scatter/gather次数，一次性广播图对象和坐标数组到所有工作节点
• 增大chunk大小（比如1000个节点/ chunk），降低调度开销
• 使用client.map简化并行逻辑：

from dask.distributed import Client

def straightness_centrality(G, vertID_dict):
    # 启用OpenMP
    gt.openmp_set_num_threads(8)
    # 转换坐标为NumPy数组
    coords = np.array([vertID_dict[n] for n in range(len(G))])
    all_nodes = list(G.iter_vertices())
    # 设置合理的chunk大小
    chunk_size = 1000
    chunks = [all_nodes[i:i+chunk_size] for i in range(0, len(all_nodes), chunk_size)]
    
    # 广播共享数据到所有节点
    G_broadcast = client.scatter(G, broadcast=True)
    coords_broadcast = client.scatter(coords, broadcast=True)
    
    # 并行处理所有chunk
    futures = client.map(alpha_batch, G_broadcast, chunks, coords_broadcast)
    results = client.gather(futures)
    
    # 扁平化结果列表
    flat_results = [item for chunk_res in results for item in chunk_res]
    return flat_results

4. 额外优化建议 #

• 确保边属性mm_len为double类型，避免类型转换开销
• 若图存在多个连通分量，可单独处理每个分量，跳过不可达节点的无效计算
• 避免预先计算全量欧氏距离矩阵（6万节点会占用约22GB内存），按需计算更高效

内容的提问来源于stack exchange，提问作者Reuben