如何将轨迹段之间的点积向量化

Question

下面是一个用于计算轨迹（xy坐标）连续线段点积的函数。结果符合预期，但由于“for循环”的存在，运行速度非常慢。

In [94]:
def func1(xy, s):
    size = xy.shape[0]-2*s
    out = np.zeros(size)
    for i in range(size):
        p1, p2 = xy[i], xy[i+s]     #segment 1
        p3, p4 = xy[i+s], xy[i+2*s] #segment 2
        out[i] = np.dot(p1-p2, p4-p3)
    return out

xy = np.array([[1,2],[2,3],[3,4],[5,6],[7,8],[2,4],[5,2],[9,9],[1,1]])
func1(xy, 2)

Out[94]:
array([-16.,  15.,  32.,  31., -14.])

我寻找一种向量化上述内容的方法，希望能够使它运行更快。以下是我想出的方法：

In [95]:
def func2(xy, s):
    size = xy.shape[0]-2*s
    p1 = xy[0:size]
    p2 = xy[s:size+s]
    p3 = p2
    p4 = xy[2*s:size+2*s]
    return np.diagonal(np.dot((p1-p2), (p4-p3).T))

func2(xy, 2)

Out[95]:
array([-16,  15,  32,  31, -14])

不幸的是，点积会产生一个方阵，我需要从中取出对角线：

In [96]:
print np.dot((p1-p2), (p4-p3).T)
np.diagonal(np.dot((p1-p2), (p4-p3).T))

[[-16  10  16 -24  10]
 [-24  15  24 -36  15]
 [-32  20  32 -48  20]
 [ 20 -13 -18  31 -14]
 [ 32 -18 -40  44 -14]]

Out[96]:
array([-16,  15,  32,  31, -14])

我的解决方案真的很糟糕。它只能将速度提高2倍，更重要的是它现在不可扩展。我的平均轨迹有几万个点，这意味着我必须处理巨大的矩阵。

你们知道更好的方法吗？ 谢谢

编辑： 太棒了！einsum绝对是解决方案。在我沮丧的时候，我自己写了点积。我知道，不太可读，并且它违背了使用优化库的目的，但是这里还是有（func4）。速度与einsum相当。

def func4(xy, s):
    size = xy.shape[0]-2*s
    tmp1 = xy[0:size] - xy[s:size+s]
    tmp2 = xy[2*s:size+2*s] - xy[s:size+s]
    return tmp1[:, 0] * tmp2[:, 0] + tmp1[:, 1] * tmp2[:, 1]

- user3329302

---

对于更大的数组，你的func4更快。 - hpaulj

Answer 1

您在func2中的想法自然而然地导致使用np.einsum。

func2的好处是，它只计算一次p1、p2、p3、p4，并将它们作为较大的数组而不是像func1中那样分成小块。

func2的不好之处在于它进行了很多你不关心的点积运算。

这就是einsum的用处。它是np.dot的更灵活版本。每当你计算一组乘积的和时，请考虑使用np.einsum。使用NumPy计算该量通常会是最快的（如果不是最快的）方法。

def func3(xy, s):
    size = xy.shape[0]-2*s
    p1 = xy[0:size]
    p2 = xy[s:size+s]
    p3 = p2
    p4 = xy[2*s:size+2*s]
    return np.einsum('ij,ij->i', p1-p2, p4-p3)

下标字符串'ij,ij->i'的意思如下：

下标字符串'ij,ij->i'有两个部分：箭头（->）之前的左边是ij,ij，箭头之后的右边是i。

在左边，逗号前面的ij是p1-p2的下标，逗号后面的ij是p4-p3的下标。

爱因斯坦求和符号对于不出现在箭头之后的重复下标进行求和。在这种情况下，j被重复使用且不会出现在箭头之后。

因此，对于每个i，计算求和(p1-p2)[i,j]*(p4-p3)[i,j]，其中求和遍历所有的j。结果是一个由i索引的数组。

---

检查无误：

In [90]: np.allclose(func1(xy, 2), func3(xy, 2))
Out[90]: True

---

下面是一个基准测试：在一个形状为(9000,2)的数组xy上，使用einsum比使用func1快450倍，比使用func2快7470倍:

In [13]: xy = np.tile(xy, (1000,1))

In [14]: %timeit func1(xy, 2)
10 loops, best of 3: 42.1 ms per loop

In [15]: %timeit func2(xy, 2)
1 loops, best of 3: 686 ms per loop

In [16]: %timeit func3(xy, 2)
10000 loops, best of 3: 91.8 µs per loop

原帖作者的func4比func3表现更出色！

In [92]: %timeit func4(xy, 2)
10000 loops, best of 3: 74.1 µs per loop

我认为 func4 之所以比 einsum 更快是因为在只有两次迭代的情况下，einsum 中循环的设置成本太高了，而手动编写求和操作则更加经济。

- unutbu

Answer 2

einsum是一个很好的工具，用于推广dot乘积。玩弄它，我可以通过以下方式重现您的数字：

np.einsum('ij,ij->i',p1-p2,p4-p3)

'ij,kj'会产生dot(p1-p2, (p4-p3).T)；而'i...,i...->i'则一步完成了对角线的计算。

---

作为你之前提出的向量叉积问题的一个分支，我尝试着

tmp11,tmp21),tmp11[:,0]*tmp21[:,0]+tmp11[:,1]*tmp21[:,1])

对于5000行的数组，它的速度几乎是einsum计算速度的两倍。

- hpaulj