python - 从多维数组矢量化采样

Question

我有一个形状的 numpy 数组D x N x K。

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

numpy 文档仅np.random.choice允许概率的一维数组。

我可以对这种类型的采样进行矢量化，还是必须使用 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)

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

标签： pythonnumpyrandomvectorizationprobability

Answer 1

既不np.random.choice也不np.random.default_rng().choice支持以您想要的方式广播概率。但是，您可以使用以下方法拼凑一些类似的东西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也不支持多维查找（可能出于相关原因）。