Как сделать выборку из нескольких категориальных дистрибутивов, используя Python

Question

Пусть P - массив, в котором каждая строка суммирует до 1. Как я могу сгенерировать матрицу A, где

• A имеет те же размеры, что и P, и имеет A_ {ij}равно 1 с вероятностью P_ {ij}

• A имеет ровно одну запись, равную 1 в каждой строке, а все остальные записи равны нулю

Как я могу сделать это в Numpy или Scipy?

Я могу сделать это, используя циклы for, но это, очевидно, медленно.Есть ли способ использовать Numpy, чтобы сделать его эффективным?Или Нумба?

Answer 1

нормально, используя выбор с расширением 2d

import numpy as np

def f(P):
    a = np.zeros(4, dtype=np.int64)
    q = np.random.choice(4, size=1, replace=True, p=P)
    a[q] = 1
    return a

P = np.array([[1/4, 1/4, 1/4, 1/4],
              [1/3,1/3,1/6,1/6]])

r = np.apply_along_axis(f, 1, P)
print(r)

[[0 0 0 1] [0 0 1 0]]

[[1 0 0 0] [0 1 00]]

Answer 2

Это следует из Википедии.

import numpy.random as rnd
import numpy as np

A_as_numbers = np.argmax(np.log(P) + rnd.gumbel(size=P.shape), axis=1)
A_one_hot = np.eye(P.shape[1])[A_as_numbers].reshape(P.shape)

Проверено на:

P = np.matrix([[1/4, 1/4, 1/4, 1/4], [1/3,1/3,1/6,1/6]])

Получил:

array([[ 1.,  0.,  0.,  0.],
       [ 0.,  1.,  0.,  0.]])