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

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.int64or来解决这个问题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