与 einsum 的叉积

Question

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

\n\n\n\n

n = 10000\na = np.random.rand(n, 3)\nb = np.random.rand(n, 3)\nnumpy.cross(a, b)\n

\n\n

给出了正确的答案，但受到类似问题的答案的启发，我认为这einsum会让我有所收获。我发现两者

\n\n\n\n

eijk = np.zeros((3, 3, 3))\neijk[0, 1, 2] = eijk[1, 2, 0] = eijk[2, 0, 1] = 1\neijk[0, 2, 1] = eijk[2, 1, 0] = eijk[1, 0, 2] = -1\n\nnp.einsum(\'ijk,aj,ak->ai\', eijk, a, b)\nnp.einsum(\'iak,ak->ai\', np.einsum(\'ijk,aj->iak\', eijk, a), b)\n

\n\n

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

\n\n

%timeit np.cross(a, b)\n1000 loops, best of 3: 628 \xc2\xb5s per loop\n

\n\n\n\n

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

\n\n\n\n

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

\n\n

关于如何改进einsums 有什么想法吗？

\n

Answer 1

您可以引入矩阵乘法，使用np.tensordot丢失第一级的一个维度，然后使用np.einsum丢失另一个维度，如下所示 -

\n\n

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

\n\n

或者，我们可以使用a和之间执行广播元素乘法，然后一次性丢失两个维度bnp.einsumnp.tensordot，如下所示 -

\n\n

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

\n\n

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

\n\n\n\n

虽然，这些方法比所提出的np.einsum仅基于方法显示出良好的改进，但未能击败np.cross。

\n\n

运行时测试 -

\n\n

In [26]: # Setup input arrays\n    ...: n = 10000\n    ...: a = np.random.rand(n, 3)\n    ...: b = np.random.rand(n, 3)\n    ...: \n\nIn [27]: # Time already posted approaches\n    ...: %timeit np.cross(a, b)\n    ...: %timeit np.einsum(\'ijk,aj,ak->ai\', eijk, a, b)\n    ...: %timeit np.einsum(\'iak,ak->ai\', np.einsum(\'ijk,aj->iak\', eijk, a), b)\n    ...: \n1000 loops, best of 3: 298 \xc2\xb5s per loop\n100 loops, best of 3: 5.29 ms per loop\n100 loops, best of 3: 9 ms per loop\n\nIn [28]: %timeit np.einsum(\'aik,ak->ai\',np.tensordot(a,eijk,axes=([1],[1])),b)\n1000 loops, best of 3: 838 \xc2\xb5s per loop\n\nIn [30]: %timeit np.tensordot(np.einsum(\'aj,ak->ajk\',a,b),eijk,axes=([1,2],[1,2]))\n1000 loops, best of 3: 882 \xc2\xb5s per loop\n

\n

Answer 2

乘法运算的次数einsum()更多cross()，并且在最新的 NumPy 版本中，cross()不会创建很多临时数组。所以einsum()不能比 更快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]

这是新的交叉代码：

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。