找到向量场的大小（长度）平方的最快方法

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,:] )

但是速度很慢，有什么更快的吗？

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