如果您查看
源代码 中的
np.cross
,它基本上将所有数组的
xyz
维度移到形状元组的前面,然后像这样拼写出每个分量的计算方式:
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]
在您的情况下,每个产品都需要分配大型数组,因此总体性能不是非常高效。
让我们设置一些测试数据:
u = np.random.rand(1000, 3)
v = np.random.rand(2000, 3)
In [13]: %timeit s1 = np.cross(u[:, None, :], v[None, :, :])
1 loops, best of 3: 591 ms per loop
我们可以尝试使用Levi-Civita符号和np.einsum
来计算它,具体方法如下:
使用Levi-Civita符号
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
In [14]: %timeit s2 = np.einsum('ijk,uj,vk->uvi', eijk, u, v)
1 loops, best of 3: 706 ms per loop
In [15]: np.allclose(s1, s2)
Out[15]: True
所以它确实可行,但性能较差。问题在于np.einsum
处理两个以上的操作数时会出现问题,但对于两个或更少的操作数有优化的路径。因此我们可以尝试将其重写为两个步骤,看看是否有帮助:
In [16]: %timeit s3 = np.einsum('iuk,vk->uvi', np.einsum('ijk,uj->iuk', eijk, u), v)
10 loops, best of 3: 63.4 ms per loop
In [17]: np.allclose(s1, s3)
Out[17]: True
太棒了!接近一个数量级的提升...
以下是NumPy 1.11.0的一些性能数据,使用a=numpy.random.rand(n,3)
和b=numpy.random.rand(n,3)
:
![enter image description here](https://istack.dev59.com/uVc2X.webp)
对于测试过的最大的n
值,嵌套的einsum
比cross
快约两倍。