rapida Selezione ponderata casuale lungo tutte le righe di una matrice stocastica

Question

numpy.random.choice permette la selezione ponderata da un vettore, cioè

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

seleziona 1 con probabilità 0.2, 2 con probabilità 0.5 e 3 con probabilità 0.3.

E se volessimo farlo rapidamente in modo vettorizzato per un array 2D (matrice) per il quale ciascuna delle righe è un vettore di probabilità? Cioè, vogliamo un vettore di scelte da una matrice stocastica? Questo è il modo super lento:

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.]) 

This post suggerisce che cumsum e bisect potrebbe essere un potenziale approccio, ed è veloce. Ma mentre lo numpy.cumsum(arr, axis=1) può farlo lungo un asse di un array numpy, la funzione bisect.bisect funziona solo su un singolo array alla volta. Allo stesso modo, numpy.searchsorted funziona anche su array 1D.

C'è un modo rapido per farlo utilizzando solo le operazioni vettorializzate?

Answer 1

Ecco una versione completamente vettorizzati che è abbastanza veloce:

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] 

In teoria, searchsorted è la funzione diritto di utilizzare per la ricerca il valore casuale nelle probabilità cumulativamente sommati, ma con m essendo relativamente piccola , k = (s < r).sum(axis=0) finisce per essere molto più veloce. La sua complessità temporale è O (m), mentre il metodo searchsorted è O (log (m)), ma ciò avrà importanza solo per dimensioni molto maggiori m. Anche, cumsum è O (m), quindi sia che @ perimosocordiae improved sono O (m). (Se il tuo m è, infatti, molto più grande, dovrai eseguire alcuni test per vedere quanto può essere grande m prima che questo metodo sia più lento.)

Ecco i tempi che ottengo con m = 10 e n = 10000 (usando le funzioni original e improved da @ di perimosocordiae risposta):

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 

Lo script completo in cui vengono definite le funzioni è:

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) 

Answer 2

Non penso sia possibile completamente vettorizzare questo, ma è comunque possibile ottenere una buona velocità associando il più possibile. Ecco quello che mi si avvicinò con:

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] 

Test contro la versione nella domanda:

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 

Ecco l'aumento di velocità (utilizzando il codice di impostazione data alla questione):

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 

Non sono sicuro del motivo per cui c'è una grande discrepanza nei tempi per la mia versione, ma anche la corsa più lenta (~ 650 μs) è ancora quasi 5 volte più veloce.