Python,成对 'distance',需要一种快速的方法来做到这一点

Python, Pairwise 'distance', need a fast way to do it

为了我博士期间的一个副项目,我在 Python 中参与了一些系统建模的任务。效率方面,我的程序在以下问题中遇到了瓶颈,我将在最小工作示例中公开。

我处理大量由 3D 起点和终点编码的片段,因此每个片段由 6 个标量表示。

我需要计算成对的最小段间距离。两段间最小距离的解析表达式见此source。致 MWE:

import numpy as np
N_segments = 1000
List_of_segments = np.random.rand(N_segments, 6)

Pairwise_minimal_distance_matrix = np.zeros( (N_segments,N_segments) )
for i in range(N_segments):
    for j in range(i+1,N_segments): 

        p0 = List_of_segments[i,0:3] #beginning point of segment i
        p1 = List_of_segments[i,3:6] #end point of segment i
        q0 = List_of_segments[j,0:3] #beginning point of segment j
        q1 = List_of_segments[j,3:6] #end point of segment j
        #for readability, some definitions
        a = np.dot( p1-p0, p1-p0)
        b = np.dot( p1-p0, q1-q0)
        c = np.dot( q1-q0, q1-q0)
        d = np.dot( p1-p0, p0-q0)
        e = np.dot( q1-q0, p0-q0)
        s = (b*e-c*d)/(a*c-b*b)
        t = (a*e-b*d)/(a*c-b*b)
        #the minimal distance between segment i and j
        Pairwise_minimal_distance_matrix[i,j] = sqrt(sum( (p0+(p1-p0)*s-(q0+(q1-q0)*t))**2)) #minimal distance

现在,我意识到这是非常低效的,这就是我来这里的原因。我已经广泛研究了如何避免循环,但我 运行 遇到了一些问题。显然,这种计算最好用 python 的 cdist 来完成。但是,它可以处理的自定义距离函数必须是二元函数。这在我的例子中是个问题,因为我的向量的长度特别是 6,并且必须按位拆分成它们的前 3 个分量和后 3 个分量。我不认为我可以将距离计算转化为二元函数。

欢迎任何意见。

这更像是一个元答案,至少对于初学者来说是这样。您的问题可能已经在 "my program hits a bottleneck" 和 "I realize this is extremely inefficient".

效率极低?通过什么措施?你有比较吗?您的代码是否太慢而无法在合理的时间内完成? 多长的时间对您来说是合理的?你能在这个问题上投入更多的计算能力吗?同样重要的是——您是否使用适当的基础设施来 运行 您的代码(numpy/scipy 使用供应商编译器编译,可能具有 OpenMP 支持)?

那么,如果您对上述所有问题都有答案并且需要进一步优化您的代码——您当前代码的瓶颈在哪里究竟?你有介绍过吗?循环体可能比循环本身的评估更重量级?如果是这样,那么 "the loop" 不是你的瓶颈,你首先不需要担心嵌套循环。首先优化主体,可能通过提出数据的非正统矩阵表示,以便您可以一步执行所有这些单一计算——例如,通过矩阵乘法。如果您的问题无法通过有效的线性代数运算解决,您可以开始编写 C 扩展或使用 Cython 或使用 PyPy(最近才获得一些基本的 numpy 支持!)。优化有无穷无尽的可能性——真正的问题是:你离实际解决方案有多近,你需要优化多少,你愿意付出多少努力投资。

免责声明:我也为我的博士学位做了 scipy/numpy 的非规范成对距离的东西 ;-)。对于一个特定的距离度量,我最终将 "pairwise" 部分编码为简单的 Python (即我还使用了双重嵌套循环),但花了一些努力使主体尽可能高效(使用i) 我的问题的加密矩阵乘法表示和 ii) 使用 bottleneck).

的组合

您可以使用 numpy 的向量化功能来加速计算。我的版本一次计算距离矩阵的所有元素,然后将对角线和下三角设置为零。

def pairwise_distance2(s):
    # we need this because we're gonna divide by zero
    old_settings = np.seterr(all="ignore")

    N = N_segments # just shorter, could also use len(s)

    # we repeat p0 and p1 along all columns
    p0 = np.repeat(s[:,0:3].reshape((N, 1, 3)), N, axis=1)
    p1 = np.repeat(s[:,3:6].reshape((N, 1, 3)), N, axis=1)
    # and q0, q1 along all rows
    q0 = np.repeat(s[:,0:3].reshape((1, N, 3)), N, axis=0)
    q1 = np.repeat(s[:,3:6].reshape((1, N, 3)), N, axis=0)

    # element-wise dot product over the last dimension,
    # while keeping the number of dimensions at 3
    # (so we can use them together with the p* and q*)
    a = np.sum((p1 - p0) * (p1 - p0), axis=-1).reshape((N, N, 1))
    b = np.sum((p1 - p0) * (q1 - q0), axis=-1).reshape((N, N, 1))
    c = np.sum((q1 - q0) * (q1 - q0), axis=-1).reshape((N, N, 1))
    d = np.sum((p1 - p0) * (p0 - q0), axis=-1).reshape((N, N, 1))
    e = np.sum((q1 - q0) * (p0 - q0), axis=-1).reshape((N, N, 1))

    # same as above
    s = (b*e-c*d)/(a*c-b*b)
    t = (a*e-b*d)/(a*c-b*b)

    # almost same as above
    pairwise = np.sqrt(np.sum( (p0 + (p1 - p0) * s - ( q0 + (q1 - q0) * t))**2, axis=-1))

    # turn the error reporting back on
    np.seterr(**old_settings)

    # set everything at or below the diagonal to 0
    pairwise[np.tril_indices(N)] = 0.0

    return pairwise

现在让我们试一试。以你的例子,N = 1000,我得到的时间是

%timeit pairwise_distance(List_of_segments)
1 loops, best of 3: 10.5 s per loop

%timeit pairwise_distance2(List_of_segments)
1 loops, best of 3: 398 ms per loop

当然,结果是一样的:

(pairwise_distance2(List_of_segments) == pairwise_distance(List_of_segments)).all()

returnsTrue。我也很确定算法中某处隐藏了一个矩阵乘法,因此应该有进一步加速(以及清理)的潜力。

顺便说一句:我试过先简单地使用 numba 但没有成功。不过不知道为什么。

你可以像这样使用它:

def distance3d (p, q):
    if (p == q).all ():
        return 0

    p0 = p[0:3]
    p1 = p[3:6]
    q0 = q[0:3]
    q1 = q[3:6]

    ...  # Distance computation using the formula above.

print (distance.cdist (List_of_segments, List_of_segments, distance3d))

虽然它似乎并没有更快,因为它在内部执行相同的循环。