问题描述
我正在尝试尽可能快地计算许多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 einsum
s?
推荐答案
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的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持!