Python 使用numpy进行乘法和累加

我在python3中使用numpy数组实现了以下循环：

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

我更愿意为循环将其矢量化。第一行可以简单地改写为：

dens = np.sum(f,2)

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

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

---

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

einsum

对于原始问题，您可以有两种方法，一种是使用

矩阵乘法

的方法，另一种是使用类似的方法-

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

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

momy

---

对于修改后的情况，您可能希望在有

f[：，：，：]

的地方使用

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

---

如果将索引切换到

f[s，：，：]

，该怎么办？@nluigi:在

ex

中添加一些额外的长度为1的轴来排列：

np.sum（ex[：，np.newaxis，np.newaxis]*f，0）

。您能否大致说明

重塑

和

einsum

方法的性能与我最初的

for loop

方法相比？目前，我发现矢量化的代码实际上比我原来的

for loop

@nluigi稍微慢一点。您正在使用的输入数组的形状是什么？

f.shape

目前是[9128128]，但我将使用更大的数组，例如[9512512]。@nluigi不确定您的基准测试过程，但我刚刚添加了我的基于。矢量化方法正以巨大的优势获胜！