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

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 编译函数进行二分搜索。

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

How to draw a sample from a categorical distribution

numpy

probability

vectorization