'How to draw a sample from a categorical distribution

I have a 3D numpy array with the probabilities of each category in the last dimension. Something like:

import numpy as np
from scipy.special import softmax

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

How can I sample from a categorical distribution with those probabilities?

EDIT: Right now I'm doing it like this:

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

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

But it's very slow so I want to know if there's a way to vectorize this operation.



Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source