使用numpy进行多次累积

Question

我在python3中有以下循环和numpy数组：

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

我更愿意将这个for循环矢量化。第一行可以改写为：

^{pr2}$

最后两行是乘法累加运算，但我不知道如何以矢量化的方式实现它。我一直在看^{}，但这似乎做了一些我不想做的事情。谁知道如何实现这一点？在

Answer 1

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

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

Answer 2

对于最初的问题，您可以有两种方法，一种是使用matrix-multiplication使用^{}，另一种使用{a2}，如下-

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,:,:]，修改后的解决方案将是-

^{pr2}$

---

运行时测试和验证输出-

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