Numpy выборка из 2d numpy массива вероятностей

Question

У меня есть двумерный массив вероятностей, p, сумма которого равна 1,0:

p = np.array([[0.05, 0.05, 0.1 , 0.  ],
              [0.4 , 0.1 , 0.1 , 0.2 ]])

Я хочу иметь функцию, которая выбирает индекс из двумерного массива p на основе вероятностей. Например, я должен вернуть [1,0] с вероятностью 0.4 или [0,2] с вероятностью 0.1.

Если я вызову функцию sample(p), которая принимает массив двумерных вероятностей p. Затем с массивом 2d p, как я определил выше, я мог ожидать индексы в диапазоне от [0,0] до [1,3], но не включая [0,3], так как вероятность появления 0 равна:

sample(p)
>> [1,0]
sample(p)
>> [0,2]
sample(p)
>> [1,0]
sample(p)
>> [1,3]

Обычно я бы использовал numpy.random.choice, однако он принимает только вероятности 1D массива. Есть ли способ, которым я могу расширить это, чтобы работать для двумерных массивов вероятности любого размера .

Answer 1

Вы можете использовать np.unravel_index() вместе с np.random.choice() для результата np.arange() с вероятностью .ravel() (или .flatten()).

Например:

p = np.array([[0.05, 0.05, 0.1 , 0.  ],
              [0.4 , 0.1 , 0.1 , 0.2 ]])


def sample_index(p):
    i = np.random.choice(np.arange(p.size), p=p.ravel())
    return np.unravel_index(i, p.shape)

np.random.seed(0)
x = [sample_index(p) for _ in range(10)]
print(x)
# [(1, 0), (1, 2), (1, 1), (1, 0), (1, 0), (1, 1), (1, 0), (1, 3), (1, 3), (1, 0)]

Answer 2

Просто сначала сделайте ваш массив плоским:)

p = np.array([[0.05, 0.05, 0.1 , 0.  ],
          [0.4 , 0.1 , 0.1 , 0.2 ]])
r = np.random.choice(np.arange(8), p=p.ravel())

Равель возвращает представление вашего массива в одномерной форме.

p.ravel().shape
> (8,)

Просто проверяйте, работает ли он правильно:

r = np.random.choice(np.arange(8), size=10000, p=p.ravel())
np.unique(r, return_counts=True)
> (array([0, 1, 2, 4, 5, 6, 7]), array([ 530,  523,  968, 4014, 1012,  944, 2009]))