FR算法（Fruchterman-Reingold）Python实现

里克贝斯

发布于 2021-05-21 10:40:14

1.2K0

发布于 2021-05-21 10:40:14

文章被收录于专栏：图灵技术域图灵技术域

简介

FR算法将所有的结点看做是电子，每个结点收到两个力的作用：

其他结点的库伦力（斥力）

边对点的胡克力（引力）。

该算法遵循两个简单的原则：有边连接的节点应该互相靠近；节点间不能离得太近。FR算法建立在粒子物理理论的基础上，将图中的节点模拟成原子，通过模拟原子间的力场来计算节点间的位置关系。算法通过考虑原子间引力和斥力的互相作用，计算得到节点的速度和加速度。依照类似原子或者行星的运动规律，系统最终进入一种动态平衡状态。

Python实现

输入的A是稀疏邻接矩阵，用scipy的coo_matrix生成，k是上述引斥力的系数，fiexed是固定不变的点的序号，返回结果是点的坐标。

Python

import numpy as np
from scipy.sparse import coo_matrix


def sparse_fruchterman_reingold(A, k=None, pos=None, fixed=None, iterations=50, threshold=1e-5, dim=3):
    # np.random.seed(1)
    nodes_num = A.shape[0]
    A = A.tolil()
    A = A.astype('float')
    if pos is None:
        pos = np.asarray(np.random.rand(nodes_num, dim), dtype=A.dtype)
        print('Init pos', pos)
    else:
        pos = np.array(pos)
        pos = pos.astype(A.dtype)

    if fixed is None:
        fixed = []

    if k is None:
        k = np.sqrt(1.0 / nodes_num)
    t = max(max(pos.T[0]) - min(pos.T[0]), max(pos.T[1]) - min(pos.T[1])) * 0.1
    dt = t / float(iterations + 1)
    displacement = np.zeros((dim, nodes_num))
    for iteration in range(iterations):
        displacement *= 0
        for i in range(A.shape[0]):
            if i in fixed:
                continue
            delta = (pos[i] - pos).T
            distance = np.sqrt((delta ** 2).sum(axis=0))
            distance = np.where(distance < 0.01, 0.01, distance)
            Ai = np.asarray(A.getrowview(i).toarray())
            print('Ai', Ai)
            displacement[:, i] += \
                (delta * (k * k / distance ** 2 - Ai * distance / k)).sum(axis=1)
        # update positions
        length = np.sqrt((displacement ** 2).sum(axis=0))
        length = np.where(length < 0.01, 0.1, length)
        delta_pos = (displacement * t / length).T
        pos += delta_pos
        # cool temperature
        t -= dt
        err = np.linalg.norm(delta_pos) / nodes_num
        if err < threshold:
            break

    return pos


if __name__ == '__main__':

    row = np.array([0, 0, 1, 2, 3, 3])
    col = np.array([1, 3, 0, 3, 0, 2])
    data = np.array([1, 1, 1, 1, 1, 1])
    coo_matrix((data, (row, col)), shape=(4, 4)).toarray()
    print(coo_matrix((data, (row, col)), shape=(4, 4)))
    a = sparse_fruchterman_reingold(coo_matrix((data, (row, col)), shape=(4, 4)))
    print(a)

其他算法可见：