在numpy数组中对飞机进行加权的最好方法是什么？

Question

我有以下代码，其中w是兼容维度的1D numpy数组，M是4D数组，

i = 0
for weight in w:
    M[:, :, i, :] *= weight
    i += 1

有没有更好的方法来达到同样的效果？

Answer 1

您正在使用来自w的元素（即1D数组）沿axis=2缩放M 因此，您需要将w扩展为具有np.newaxis/None的2D数组，这将使extended version of w与M之间的轴对齐。 然后，在这两个数组之间执行逐元素乘法，为矢量化解决方案带来broadcasting ，如下所示 -

M *= w[:,None]

如果M axis=2的长度大于w的元素数，则需要在乘以之前选择M沿axis=2的范围，如下所示 -

M[...,np.arange(w.size),:] *= w[:,None]

Answer 2

这个答案是基于我试图理解@Divakar的答案。 什么帮助我了解正在发生的事情是写他们的建议

M *= w[:,None]

如

M *= w[None,None,:,None]

现在，M和扩展w的尺寸明显相同。 当然，@ Divakar的版本更短，更优雅，但不太直观。

所以，一个完整的工作示例是：

import numpy as np
M = np.ones((1,4,3,2))
w = np.arange(3)
M *= w[None,None,:,None]
print M