从多维数组中向量化采样

Question

我有一个形状为 D x N x K.

的 numpy 数组

我需要一个结果 D x N 随机元素数组 类，其中对于每个索引 [d, n] 类，有一个不同的概率向量，由第三轴表示。

np.random.choice 的 numpy documentation 只允许概率的一维数组。

我可以向量化这种类型的采样，还是必须使用如下的 for 循环：

# input_array of shape (D, N, K)
# output_array of shape (D, N)

for d in range(input_array.shape[0]):
    for n in range(input_array.shape[1]):
        probabilities = input_array[d, n]
        element = np.random.choice(K, p=probabilities)
        output_array[d, n] = element

如果有

这样的功能就好了

output_array = np.random.choice(input_array, K, probability_axis=-1)

编辑：设法找到“手工设计”的解决方案 。

Answer 1

都没有np.random.choice nor np.random.default_rng().choice support broadcasting of probabilities in the way that you intend. However, you can cobble together something that works similarly using np.cumsum:

sprob = input_array.cumsum(axis=-1, dtype=float)
sprob /= sprob[:, :, -1:]
output_array = (np.random.rand(D, N, 1) > sprob).argmin(-1)

很遗憾，np.searchsorted也不支持多维查找（可能是出于相关原因）。