Рисование случайной выборки из массива numpy с индексом

Question

У меня есть numpy матрица размером 12x12, содержащая вероятности. Цель состоит в том, чтобы взять случайную выборку с использованием этих вероятностей, а затем вернуть индекс выборки.

В настоящее время я использую следующий код, чтобы сделать это на основе np.random.choice, где grid = матрица numpy:

    rnd_choice = np.random.choice(grid.size, p=grid.ravel() / grid.ravel().sum())
    sample_index = np.unravel_index(rnd_choice, grid.shape)

Проблема в скорости, так как мне приходится делать это несколько тысяч раз за симуляцию. Снейквиз выделил это как область для улучшения, поэтому мне было интересно, есть ли у кого-нибудь какие-либо идеи относительно того, как увеличить скорость этого? .

Answer 1

Следующее выглядит немного (x4 на моем ноутбуке) быстрее:

c = grid.ravel().cumsum()                                              
out = np.unravel_index(c.searchsorted(rng.uniform(0,c[-1])),grid.shape)

Answer 2

Если grid.ravel необходимо вычислить в течение для l oop, поскольку вероятности меняются в течение каждого цикла для l oop, вы все равно можете потенциально уменьшить вычислительную нагрузку в два раза, вызвав .ravel() только один раз за цикл:

for cycle in loop:
    # grid gets renewed here...
    g_ravel = grid.ravel()         # do the ravel() process only once

    rnd_choice = np.random.choice(grid.size, p=g_ravel / g_ravel.sum())
    sample_index = np.unravel_index(rnd_choice, grid.shape)