寻找向量场长度平方的最快方法

Question

我有一个大向量场，其中场很大（例如512^3；但不一定是正方形），而向量可以是2D或3D（例如形状为[512, 512, 512, 2]或[512, 512, 512, 3]）。

最快的计算向量平方幅值的标量场的方法是什么？

我可以只循环每个方向，即：

import numpy as np
shp = [256,256,256,3]                       # Shape of vector field
vf = np.arange(3*(256**3)).reshape(shp)     # Create vector field
sf = np.zeros(shp[:3])                      # Create scalar field for result

for ii in range(shp[0]):
    for jj in range(shp[1]):
        for kk in range(shp[2]):
            sf[ii,jj,kk] = np.dot( vf[ii,jj,kk,:] , vf[ii,jj,kk,:] )

但那个相对较慢，有更快的选择吗？

- DilithiumMatrix

---

你的问题是否类似于这个？https://dev59.com/LWw15IYBdhLWcg3wG4AJ - Garth5689

Answer 1

最快的方法可能是使用np.einsum函数：

np.einsum('...j,...j->...', vf, vf)

以上代码告诉numpy获取其两个输入，并通过乘以相应的值并将它们相加来减少每个输入的最后一个维度。由于您的数据集存在溢出问题，因为大小不适合32位整数，这是 np.arange 的默认返回值。您可以通过指定返回dtype（ np.int64 或 np.double ）来解决此问题：

>>> np.einsum('...j,...j->...', vf,vf)[-1, -1, -1]
-603979762
>>> np.einsum('...j,...j->...', vf,vf).dtype
dtype('int32')

>>> np.einsum('...j,...j->...', vf,vf, dtype=np.int64)[-1, -1, -1]
7599823767207950
>>> np.einsum('...j,...j->...', vf,vf, dtype=np.double)[-1, -1, -1]
7599823767207950.0

- Jaime

---

谢谢@Jaime，你能否解释一下这意味着什么，或者为什么它是最优的？ - DilithiumMatrix

哇，einsum 真的让我这个星期过得非常愉快。 - DilithiumMatrix

另一种可能更明确的方法是使用np.sum(vf * vf, axis=-1)。然而，快速测试（针对给定的数组大小）显示einsum版本的速度是原来的两倍以上（我猜测这是因为它避免了使用临时的vf * vf数组）。 - bogatron

@bogatron Einsum 可以使用 SSE2，而 numpy 的 ufuncs 在 1.7 版本中不支持，详见这里。 - Daniel

@Ophion 感谢提供链接。这非常有启发性。如果我理解有误，请纠正我，但是您在链接中展示的结果表明，由于此处提出的问题最接近您的最后一个示例，因此我的上面的评论仍然有效，对于这种情况，einsum 仍比 ufunc 解决方案快得多（即使使用 SSE2）。 - bogatron

@bogatron Einsum避免了大多数中间变量，这使得它在内存方面表现良好，但速度来自SSE。这回答了你的问题吗？ - Daniel