numpy.random.choice permite la selección ponderada de un vector, es decir

arr = numpy.array([1, 2, 3])
weights = numpy.array([0.2, 0.5, 0.3])
choice = numpy.random.choice(arr, p=weights) 

selecciona 1 con probabilidad 0.2, 2 con probabilidad 0.5 y 3 con probabilidad 0.3.

¿Qué pasaría si quisiéramos hacer esto rápidamente de forma vectorizada para una matriz 2D (matriz) para la cual cada una de las filas es un vector de probabilidades? Es decir, ¿queremos un vector de opciones de una matriz estocástica? Esta es la forma súper lenta:

import numpy as np

m = 10
n = 100 # Or some very large number

items = np.arange(m)
prob_weights = np.random.rand(m, n)
prob_matrix = prob_weights / prob_weights.sum(axis=0, keepdims=True)

choices = np.zeros((n,))
# This is slow, because of the loop in Python
for i in range(n):
    choices[i] = np.random.choice(items, p=prob_matrix[:,i])

print(choices):

array([ 4.,  7.,  8.,  1.,  0.,  4.,  3.,  7.,  1.,  5.,  7.,  5.,  3.,
        1.,  9.,  1.,  1.,  5.,  9.,  8.,  2.,  3.,  2.,  6.,  4.,  3.,
        8.,  4.,  1.,  1.,  4.,  0.,  1.,  8.,  5.,  3.,  9.,  9.,  6.,
        5.,  4.,  8.,  4.,  2.,  4.,  0.,  3.,  1.,  2.,  5.,  9.,  3.,
        9.,  9.,  7.,  9.,  3.,  9.,  4.,  8.,  8.,  7.,  6.,  4.,  6.,
        7.,  9.,  5.,  0.,  6.,  1.,  3.,  3.,  2.,  4.,  7.,  0.,  6.,
        3.,  5.,  8.,  0.,  8.,  3.,  4.,  5.,  2.,  2.,  1.,  1.,  9.,
        9.,  4.,  3.,  3.,  2.,  8.,  0.,  6.,  1.])

Esta publicación sugiere quecumsum ybisect podría ser un enfoque potencial, y es rápido. Pero mientrasnumpy.cumsum(arr, axis=1) puede hacer esto a lo largo de un eje de una matriz numpy, elbisect.bisect La función solo funciona en una sola matriz a la vez. Similar,numpy.searchsorted solo funciona en matrices 1D también.

¿Hay una manera rápida de hacer esto usando solo operaciones vectorizadas?