如何使用 numpy 对带条件的屏蔽加权平均值进行矢量化？

问题描述

未向量化的代码如下：

import numpy as np
import numpy.ma as ma

np.random.seed(42)
H = np.random.uniform(0.1,1.0,size=(6,8))
r,c = H.shape

mask = H.max(axis=1) > 0.95


x = np.linspace(0,10,c)
weighted_averages = ma.masked_all((r,),dtype=H.dtype)

for i in range(r):
    if mask[i]:
        weighted_averages[i] = np.average(x,weights=H[i,:])

这是我对其进行矢量化的尝试：

_,xx = np.mgrid[0:10:r*1j,0:10:c*1j]
not_mask = np.logical_not(mask)


weighted_averages = np.average(xx,weights=H,axis=1)
mwa = ma.masked_array(weighted_averages,mask=not_mask)

它有效，从某种意义上说输出是相同的，但我在“作弊”，因为我首先计算所有平均值，然后屏蔽“不需要的”值。我怎样才能避免不必要的计算？我猜我必须以某种方式屏蔽 xx、H 或两者。

解决方法

这个怎么样 -

import numpy as np
import numpy.ma as ma

np.random.seed(42)
H = np.random.uniform(0.1,1.0,size=(6,8))
r,c = H.shape

mask = H.max(axis=1) > 0.95

x = np.linspace(0,10,c)

H_mask = H[mask]
wa = (np.sum(x * H_mask,axis=1))/np.sum(H_mask,axis=1)
weighted_averages = ma.masked_all((r,),dtype=H.dtype)

weighted_averages[mask] = wa

简单地先屏蔽数组，然后取平均值。我认为您不能在这里使用 np.average，因为它似乎不支持广播。因此，只需手动计算平均值即可。