为给定的2D轴概率数组矢量化`numpy.random.choice`

时间:2017-12-08 20:55:46

标签: python numpy random vectorization

Numpy具有random.choice功能,允许您从分类分布中进行采样。你会如何在轴上重复这个?为了说明我的意思,这是我目前的代码:

categorical_distributions = np.array([
    [.1, .3, .6],
    [.2, .4, .4],
])
_, n = categorical_distributions.shape
np.array([np.random.choice(n, p=row)
          for row in categorical_distributions])

理想情况下,我想消除for循环。

1 个答案:

答案 0 :(得分:1)

这是获取每行随机索引的一种矢量化方式,a2D概率数组 -

(a.cumsum(1) > np.random.rand(a.shape[0])[:,None]).argmax(1)

概括以涵盖2D数组 -

的行和列
def random_choice_prob_index(a, axis=1):
    r = np.expand_dims(np.random.rand(a.shape[1-axis]), axis=axis)
    return (a.cumsum(axis=axis) > r).argmax(axis=axis)

让我们通过运行它超过一百万次验证给定的样本 -

In [589]: a = np.array([
     ...:     [.1, .3, .6],
     ...:     [.2, .4, .4],
     ...: ])

In [590]: choices = [random_choice_prob_index(a)[0] for i in range(1000000)]

# This should be close to first row of given sample
In [591]: np.bincount(choices)/float(len(choices))
Out[591]: array([ 0.099781,  0.299436,  0.600783])

运行时测试

原始循环方式 -

def loopy_app(categorical_distributions):
    m, n = categorical_distributions.shape
    out = np.empty(m, dtype=int)
    for i,row in enumerate(categorical_distributions):
        out[i] = np.random.choice(n, p=row)
    return out

更大阵列上的计时 -

In [593]: a = np.array([
     ...:     [.1, .3, .6],
     ...:     [.2, .4, .4],
     ...: ])

In [594]: a_big = np.repeat(a,100000,axis=0)

In [595]: %timeit loopy_app(a_big)
1 loop, best of 3: 2.54 s per loop

In [596]: %timeit random_choice_prob_index(a_big)
100 loops, best of 3: 6.44 ms per loop