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

Question

У меня есть 3D numpy массив с вероятностями каждой категории в последнем измерении. Что-то вроде:

import numpy as np
from scipy.special import softmax

array = np.random.normal(size=(10, 100, 5))
probabilities = softmax(array, axis=2)

Как я могу выбрать из категориального распределения с этими вероятностями?

РЕДАКТИРОВАТЬ: Сейчас я делаю это так:

def categorical(x):
    return np.random.multinomial(1, pvals=x)

samples = np.apply_along_axis(categorical, axis=2, arr=probabilities)

Но это очень медленно, поэтому я хочу знать, есть ли способ векторизовать эту операцию.

Answer 1

Выборка из заданного распределения вероятностей выполняется путем построения обратного кумулятивного распределения для случайного числа в диапазоне от 0 до 1. Для небольшого числа дискретных категорий - как в вопросе - вы можете найти обратное, используя линейный поиск:

## Alternative test dataset
probabilities[:, :, :] = np.array([0.1, 0.5, 0.15, 0.15, 0.1])

n1, n2, m = probabilities.shape

cum_prob = np.cumsum(probabilities, axis=-1) # shape (n1, n2, m)
r = np.random.uniform(size=(n1, n2, 1))

# argmax finds the index of the first True value in the last axis.
samples = np.argmax(cum_prob > r, axis=-1)

print('Statistics:')
print(np.histogram(samples, bins=np.arange(m+1)-0.5)[0]/(n1*n2))

Для тестового набора данных типичный тестовый результат был:

Statistics:
[0.0998 0.4967 0.1513 0.1498 0.1024]

, что выглядит нормально.

Если у вас много, много категорий (тысячи), вероятно, лучше выполнить поиск по делению пополам, используя скомпилированную функцию numba.