Question

我在python3中使用numpy数组进行以下循环：

for s in range(ns):
    dens += f[:,:,s]
    momx += ex[s]*f[:,:,s]
    momy += ey[s]*f[:,:,s]

我更喜欢对这个for循环进行矢量化。第一行可以简单地重写为：

dens = np.sum(f,2)

最后两行是乘法累加运算，但我不知道如何以矢量化方式实现它。我一直在看np.ufunc.accumulate，但这似乎做了别的事情，而不是我想要的。谁知道如何实现这个？

Answer 1

momx = np.sum(ex*f, 2)

您还可以使用其他一些工具，例如einsum。

Answer 2

对于您的原始问题，您可以采用两种方法，一种使用matrix-multiplication使用ndarray.dot，另一种使用np.einsum，就像这样 -

momx = f.reshape(-1,f.shape[2]).dot(ex).reshape(f.shape[:2])
momx = np.einsum('ijk,k->ij',f,ex)

您可以使用类似的策略来查找momy。

对于修改过的案例，您可能希望在f[s,:,:]的地方使用f[:,:,s]，修改后的解决方案将是 -

momx = ex.dot(f.reshape(f.shape[0],-1)).reshape(f.shape[1:])
momx = np.einsum('ijk,i->jk',f,ex)

运行时测试并验证输出 -

In [573]: def org_app(f,ex):
     ...:     ns = f.shape[2]
     ...:     momx = np.zeros(f.shape[:2])
     ...:     for s in range(ns):
     ...:         momx += ex[s]*f[:,:,s]
     ...:     return momx
     ...: 

In [574]: f = np.random.rand(9,512,512)
     ...: ex = np.random.rand(f.shape[2])
     ...: 

In [575]: np.allclose(org_app(f,ex),f.reshape(-1,f.shape[2]).dot(ex).reshape(f.shape[:2]))
Out[575]: True

In [576]: np.allclose(org_app(f,ex),np.einsum('ijk,k->ij',f,ex))
Out[576]: True

In [581]: %timeit org_app(f,ex)
10 loops, best of 3: 44.8 ms per loop

In [582]: %timeit f.reshape(-1,f.shape[2]).dot(ex).reshape(f.shape[:2])
100 loops, best of 3: 4.8 ms per loop

In [583]: %timeit np.einsum('ijk,k->ij',f,ex)
100 loops, best of 3: 3.84 ms per loop

使用numpy进行乘法累加

2 个答案: