交叉产品与einsums [英] cross products with einsums

问题描述

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

I'm trying to compute the cross-products of many 3x1 vector pairs as fast as possible. This

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

给出了正确的答案，但是出于对类似问题的回答，我认为einsum会得到我在某个地方.我发现两者都

gives the correct answer, but motivated by this answer to a similar question, I thought that einsum would get me somewhere. I found that both

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低得多:

compute the cross product, but their performance is disappointing: Both methods perform much worse than 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

关于如何改善einsum的任何想法?

Any ideas of how to improve the einsums?

推荐答案

einsum()的乘法运算计数大于cross()，并且在最新的NumPy版本中，cross()不会创建许多临时数组.因此，einsum()不能快于cross().

The count of multiply operation of einsum() is more then cross(), and in the newest NumPy version, cross() doesn't create many temporary arrays. So einsum() can't be faster than 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.

To speedup it, you need cython or numba.

这篇关于交叉产品与einsums的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！