找到矢量场的幅度(长度)平方的最快方法

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.您可以解决通过指定返回D型,为无论是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