如何从分类分布中抽取样本

Question

我有一个 3D numpy 数组，其中最后一个维度中包含每个类别的概率。就像是：

import numpy as np
from scipy.special import softmax

array = np.random.normal(size=(10, 100, 5))
probabilities = softmax(array, axis=2)

如何从具有这些概率的分类分布中进行采样？

编辑：现在我正在这样做：

def categorical(x):
    return np.random.multinomial(1, pvals=x)

samples = np.apply_along_axis(categorical, axis=2, arr=probabilities)

但它非常慢，所以我想知道是否有一种方法可以向量化这个操作。

Answer 1

从给定概率分布中抽取样本是通过构建 0 到 1 范围内的随机数的逆累积分布的评估来完成的。对于少量离散类别（如问题中所示），您可以使用线性搜索找到逆累积分布：

## Alternative test dataset
probabilities[:, :, :] = np.array([0.1, 0.5, 0.15, 0.15, 0.1])

n1, n2, m = probabilities.shape

cum_prob = np.cumsum(probabilities, axis=-1) # shape (n1, n2, m)
r = np.random.uniform(size=(n1, n2, 1))

# argmax finds the index of the first True value in the last axis.
samples = np.argmax(cum_prob > r, axis=-1)

print('Statistics:')
print(np.histogram(samples, bins=np.arange(m+1)-0.5)[0]/(n1*n2))

对于测试数据集，典型的测试输出是：

Statistics:
[0.0998 0.4967 0.1513 0.1498 0.1024]

看起来不错。

如果您有很多很多类别（数千个），最好使用 numba 编译函数进行二分搜索。