python - 与 einsums 的交叉产品

Question

我正在尝试尽可能快地计算许多 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

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

Answer 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。