python - 向量化 numpy.random.multinomial

Question

我正在尝试向量化以下代码:

for i in xrange(s.shape[0]):
            a[i] = np.argmax(np.random.multinomial(1,s[i,:]))

s.shape = 400 x 100 [给定]。

a.shape = 400 [预期]。

s 是一个二维矩阵，其中包含成对的概率。期望多项式从s矩阵的每一行中抽取一个随机样本，并将结果存储在向量a中。

Answer 1

在comments , 据说有人试图将这个 here 向量化然而，这不仅仅是一次尝试。也是这个问题的完整解决方案。

问题的目标是获取包含多项式事件的 1 的位置的索引。也就是说，以下实现 [0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0] 将产生 14。因此，它实际上等同于执行:

np.random.choice(np.arange(len(p)),p=p) # here, p is s[i,:]

因此，Warren Weckesser solution 对随机矩阵的所有行进行快速随机加权选择也是这个问题的解决方案。唯一的区别是概率向量是按行定义还是按列定义，这可以很容易地解决，可以将 s 转置为 prob_matrix 或定义自定义版本适用于 s 结构的 vectorized:

def vectorized(prob_matrix, items):
    s = prob_matrix.cumsum(axis=1)
    r = np.random.rand(prob_matrix.shape[0])
    k = (s < r).sum(axis=1)
    return items[k]

在这个问题中，尺寸为 400x400，加速大约为 10 倍:

%%timeit
a = np.empty(400)
for i in range(s.shape[0]):
    a[i] = np.argmax(np.random.multinomial(1,s[i,:]))
# 5.96 ms ± 46.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%%timeit 
vals = np.arange(400,dtype=int)
vectorized(s,vals)
# 544 µs ± 5.49 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)