軸に沿った確率の与えられた2D配列に対してnumpy.random.choiceをベクトル化する

Question

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ループを削除したいと思います。

Answer 1

ここ確率の2Dアレイとしてaで、行ごとにランダムインデックスを取得する一つのベクトル化方法です -

(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) 

100万回以上実行してサンプルを検証しましょう -

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