与 einsum 的交叉产品答案

【问题标题】：cross products with einsums与 einsum 的交叉产品
【发布时间】：2017-02-01 10:02:42
【问题描述】：

我正在尝试尽快计算许多 3x1 向量对的叉积。这个

n = 10000
a = np.random.rand(n, 3)
b = np.random.rand(n, 3)
numpy.cross(a, b)

给出了正确的答案，但受到this answer to a similar question 的激励，我认为einsum 会带我到某个地方。我发现两者都有

eijk = np.zeros((3, 3, 3))
eijk[0, 1, 2] = eijk[1, 2, 0] = eijk[2, 0, 1] = 1
eijk[0, 2, 1] = eijk[2, 1, 0] = eijk[1, 0, 2] = -1

np.einsum('ijk,aj,ak->ai', eijk, a, b)
np.einsum('iak,ak->ai', np.einsum('ijk,aj->iak', eijk, a), b)

计算叉积，但它们的性能令人失望：两种方法的性能都比np.cross差得多：

%timeit np.cross(a, b)
1000 loops, best of 3: 628 µs per loop

%timeit np.einsum('ijk,aj,ak->ai', eijk, a, b)
100 loops, best of 3: 9.02 ms per loop

%timeit np.einsum('iak,ak->ai', np.einsum('ijk,aj->iak', eijk, a), b)
100 loops, best of 3: 10.6 ms per loop

关于如何改进einsums 的任何想法？

【问题讨论】：

标签： python performance numpy cross-product numpy-einsum

【解决方案1】：

einsum() 的乘法运算次数比cross() 多，在最新的 NumPy 版本中，cross() 不会创建很多临时数组。所以einsum() 不能比cross() 快。

这里是cross的旧代码：

x = a[1]*b[2] - a[2]*b[1]
y = a[2]*b[0] - a[0]*b[2]
z = a[0]*b[1] - a[1]*b[0]

这是cross的新代码：

multiply(a1, b2, out=cp0)
tmp = array(a2 * b1)
cp0 -= tmp
multiply(a2, b0, out=cp1)
multiply(a0, b2, out=tmp)
cp1 -= tmp
multiply(a0, b1, out=cp2)
multiply(a1, b0, out=tmp)
cp2 -= tmp

要加快速度，您需要 cython 或 numba。

【讨论】：

由于这些不涉及循环，cython 的改进可能并不显着。当这样表达时，cross 更像是一种代数运算而不是数组运算。

【解决方案2】：

您可以使用np.tensordot 引入矩阵乘法以丢失第一级的一个维度，然后使用np.einsum 丢失另一个维度，就像这样 -

np.einsum('aik,ak->ai',np.tensordot(a,eijk,axes=([1],[1])),b)

或者，我们可以使用np.einsum 在a 和b 之间执行广播元素乘法，然后使用np.tensordot 一次性丢失两个维度，就像这样 -

np.tensordot(np.einsum('aj,ak->ajk', a, b),eijk,axes=([1,2],[1,2]))

我们也可以通过引入新轴来执行元素乘法，例如 a[...,None]*b[:,None]，但它似乎减慢了速度。

不过，这些方法比建议的仅基于 np.einsum 的方法有很好的改进，但未能击败 np.cross。

运行时测试-

In [26]: # Setup input arrays
    ...: n = 10000
    ...: a = np.random.rand(n, 3)
    ...: b = np.random.rand(n, 3)
    ...: 

In [27]: # Time already posted approaches
    ...: %timeit np.cross(a, b)
    ...: %timeit np.einsum('ijk,aj,ak->ai', eijk, a, b)
    ...: %timeit np.einsum('iak,ak->ai', np.einsum('ijk,aj->iak', eijk, a), b)
    ...: 
1000 loops, best of 3: 298 µs per loop
100 loops, best of 3: 5.29 ms per loop
100 loops, best of 3: 9 ms per loop

In [28]: %timeit np.einsum('aik,ak->ai',np.tensordot(a,eijk,axes=([1],[1])),b)
1000 loops, best of 3: 838 µs per loop

In [30]: %timeit np.tensordot(np.einsum('aj,ak->ajk',a,b),eijk,axes=([1,2],[1,2]))
1000 loops, best of 3: 882 µs per loop

【讨论】：