我有一个大的矢量场,其中场很大(例如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,:] )
但这很慢,有什么更快的吗?
答案 0 :(得分:6)
最快的可能是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