Selección ponderada aleatoria rápida en todas las filas de una matriz estocástica

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 en forma vectorial para una matriz 2D (matriz) para la cual cada una de las filas es un vector de probabilidades? Es decir, ¿queremos un vector de elecciones desde una matriz estocástica? Esta es la forma super 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 que el cumsum y la cumsum podrían ser un enfoque potencial y es rápido. Pero mientras numpy.cumsum(arr, axis=1) puede hacer esto a lo largo de un eje de una matriz numpy, la función bisect.bisect solo funciona en una única matriz a la vez. Del mismo modo, numpy.searchsorted solo funciona en arreglos 1D también.

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

Aquí hay una versión completamente vectorizada que es bastante rápida:

 def vectorized(prob_matrix, items): s = prob_matrix.cumsum(axis=0) r = np.random.rand(prob_matrix.shape[1]) k = (s < r).sum(axis=0) return items[k] 

En teoría , searchsorted es la función correcta a usar para buscar el valor aleatorio en las probabilidades acumuladas acumulativamente, pero con m siendo relativamente pequeño, k = (s < r).sum(axis=0) termina siendo mucho más rápido. Su complejidad temporal es O (m), mientras que el método de searchsorted es O (log (m)), pero eso solo importará para m mucho más grande. Además , el cumsum es O (m), por lo tanto, tanto vectorized como @ perimosocordiae improved son O (m). (Si su m es, de hecho, mucho más grande, tendrá que realizar algunas pruebas para ver qué tan grande puede ser m antes de que este método sea más lento).

Aquí está el tiempo que obtengo con m = 10 y n = 10000 (usando las funciones original y improved de la respuesta de @ perimosocordiae):

 In [115]: %timeit original(prob_matrix, items) 1 loops, best of 3: 270 ms per loop In [116]: %timeit improved(prob_matrix, items) 10 loops, best of 3: 24.9 ms per loop In [117]: %timeit vectorized(prob_matrix, items) 1000 loops, best of 3: 1 ms per loop 

El script completo donde se definen las funciones es:

 import numpy as np def improved(prob_matrix, items): # transpose here for better data locality later cdf = np.cumsum(prob_matrix.T, axis=1) # random numbers are expensive, so we'll get all of them at once ridx = np.random.random(size=n) # the one loop we can't avoid, made as simple as possible idx = np.zeros(n, dtype=int) for i, r in enumerate(ridx): idx[i] = np.searchsorted(cdf[i], r) # fancy indexing all at once is faster than indexing in a loop return items[idx] def original(prob_matrix, items): 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]) return choices def vectorized(prob_matrix, items): s = prob_matrix.cumsum(axis=0) r = np.random.rand(prob_matrix.shape[1]) k = (s < r).sum(axis=0) return items[k] m = 10 n = 10000 # 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) 

No creo que sea posible vectorizar completamente esto, pero aún puedes obtener una aceleración decente vectorizando todo lo que puedas. Esto es lo que se me ocurrió:

 def improved(prob_matrix, items): # transpose here for better data locality later cdf = np.cumsum(prob_matrix.T, axis=1) # random numbers are expensive, so we'll get all of them at once ridx = np.random.random(size=n) # the one loop we can't avoid, made as simple as possible idx = np.zeros(n, dtype=int) for i, r in enumerate(ridx): idx[i] = np.searchsorted(cdf[i], r) # fancy indexing all at once is faster than indexing in a loop return items[idx] 

Probando contra la versión en la pregunta:

 def original(prob_matrix, items): 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]) return choices 

Aquí está la aceleración (usando el código de configuración dado en la pregunta):

 In [45]: %timeit original(prob_matrix, items) 100 loops, best of 3: 2.86 ms per loop In [46]: %timeit improved(prob_matrix, items) The slowest run took 4.15 times longer than the fastest. This could mean that an intermediate result is being cached 10000 loops, best of 3: 157 µs per loop 

No estoy seguro de por qué hay una gran discrepancia en los tiempos de mi versión, pero incluso la ejecución más lenta (~ 650 µs) es casi 5 veces más rápida.