从Python中的联合概率质量函数采样

Question

我有一个非负归一化向量

p

。我想从向量的索引集中采样一个索引。获得样本

k

的概率为

p[k]

。使用

np.random.choice

函数，我可以编写以下代码。

p = [0.2, 0.3, 0.1, 0.3, 0.1]
indices = np.arange(len(p))
k = np.random.choice(indices, p=p)

我的问题是，如何将这段代码推广到多维数组？例如，给定三维非负归一化 IxJxK 张量

p = np.random.rand(I,J,K)

如何以

p[i,j,k]

的概率对索引 (i,j,k) 进行采样？

Answer 1

假设

x

是一个概率矩阵：

x = np.random.random((3, 4, 5))
x /= np.sum(x)

您可以使用

flatten(x)

和

np.random.choice()

获取 0..3x4x5 范围内的值以及相关概率：

flat = x.flatten()
temp = np.random.choice(np.arange(len(flat)), 10, p=flat)
print(temp)

您现在可以将

temp

中的索引转换为原始数组中的索引：

np.unravel_index(temp, x.shape))

请注意，这将返回三个 numpy 数组的元组，第一个数组是第一个索引，第二个元素是第二个索引，依此类推。