Alternativa veloce per numpy.median.reduceat

In relazione a questa risposta , esiste un modo rapido per calcolare i mediani su un array che ha gruppi con un numero diseguale di elementi?

Per esempio:

data =  [1.00, 1.05, 1.30, 1.20, 1.06, 1.54, 1.33, 1.87, 1.67, ... ]
index = [0,    0,    1,    1,    1,    1,    2,    3,    3,    ... ]

E poi voglio calcolare la differenza tra il numero e la mediana per gruppo (es. La mediana del gruppo 0è 1.025quindi il primo risultato è 1.00 - 1.025 = -0.025). Quindi, per l'array sopra, i risultati appariranno come:

result = [-0.025, 0.025, 0.05, -0.05, -0.19, 0.29, 0.00, 0.10, -0.10, ...]

Dal momento np.median.reduceatche non esiste (ancora), c'è un altro modo rapido per raggiungere questo obiettivo? Il mio array conterrà milioni di righe, quindi la velocità è cruciale!

Si può presumere che gli indici siano contigui e ordinati (è facile trasformarli se non lo sono).

Dati di esempio per confronti di prestazioni:

import numpy as np

np.random.seed(0)
rows = 10000
cols = 500
ngroup = 100

# Create random data and groups (unique per column)
data = np.random.rand(rows,cols)
groups = np.random.randint(ngroup, size=(rows,cols)) + 10*np.tile(np.arange(cols),(rows,1))

# Flatten
data = data.ravel()
groups = groups.ravel()

# Sort by group
idx_sort = groups.argsort()
data = data[idx_sort]
groups = groups[idx_sort]

A volte è necessario scrivere codice numpy non idiomatico se si desidera davvero accelerare il calcolo che non si può fare con numpy nativo.

numbacompila il tuo codice Python a basso livello C. Dato che un sacco di numpy stesso è di solito veloce quanto C, questo finisce per essere utile se il tuo problema non si presta alla vettorizzazione nativa con numpy. Questo è un esempio (dove ho assunto che gli indici siano contigui e ordinati, che si riflette anche nei dati di esempio):

import numpy as np
import numba

# use the inflated example of roganjosh https://stackoverflow.com/a/58788534
data =  [1.00, 1.05, 1.30, 1.20, 1.06, 1.54, 1.33, 1.87, 1.67]
index = [0,    0,    1,    1,    1,    1,    2,    3,    3] 

data = np.array(data * 500) # using arrays is important for numba!
index = np.sort(np.random.randint(0, 30, 4500))               

# jit-decorate; original is available as .py_func attribute
@numba.njit('f8[:](f8[:], i8[:])') # explicit signature implies ahead-of-time compile
def diffmedian_jit(data, index): 
    res = np.empty_like(data) 
    i_start = 0 
    for i in range(1, index.size): 
        if index[i] == index[i_start]: 
            continue 

        # here: i is the first _next_ index 
        inds = slice(i_start, i)  # i_start:i slice 
        res[inds] = data[inds] - np.median(data[inds]) 

        i_start = i 

    # also fix last label 
    res[i_start:] = data[i_start:] - np.median(data[i_start:])

    return res

E qui ci sono alcuni tempi usando la %timeitmagia di IPython :

>>> %timeit diffmedian_jit.py_func(data, index)  # non-jitted function
... %timeit diffmedian_jit(data, index)  # jitted function
...
4.27 ms ± 109 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
65.2 µs ± 1.01 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Utilizzando i dati di esempio aggiornati nella domanda, questi numeri (ovvero il runtime della funzione python rispetto al runtime della funzione accelerata JIT) sono

>>> %timeit diffmedian_jit.py_func(data, groups) 
... %timeit diffmedian_jit(data, groups)
2.45 s ± 34.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
93.6 ms ± 518 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Ciò equivale a uno speedup 65x nel caso più piccolo e uno speedup 26x nel caso più grande (rispetto al codice a ciclo lento, ovviamente) usando il codice accelerato. Un altro aspetto positivo è che (a differenza della vettorializzazione tipica con intorpidimento nativo) non abbiamo avuto bisogno di memoria aggiuntiva per raggiungere questa velocità, è tutto basato sul codice di basso livello ottimizzato e compilato che finisce per essere eseguito.

La funzione sopra presuppone che gli array int intorpiditi siano int64predefiniti, il che non è effettivamente il caso su Windows. Quindi un'alternativa è rimuovere la firma dalla chiamata a numba.njit, innescando una corretta compilazione just-in-time. Ciò significa che la funzione verrà compilata durante la prima esecuzione, che può intromettersi con i risultati di temporizzazione (possiamo eseguire la funzione una volta manualmente, utilizzando tipi di dati rappresentativi, o semplicemente accettare che la prima esecuzione di temporizzazione sarà molto più lenta, il che dovrebbe essere ignorato). Questo è esattamente ciò che ho cercato di prevenire specificando una firma, che attiva la compilazione anticipata.

Ad ogni modo, nel caso correttamente JIT il decoratore di cui abbiamo bisogno è giusto

@numba.njit
def diffmedian_jit(...):

Si noti che i tempi sopra indicati per la funzione compilata jit si applicano solo dopo che la funzione è stata compilata. Ciò accade sia in fase di definizione (con compilation ansiosa, quando viene passata una firma esplicita numba.njit), sia durante la prima chiamata di funzione (con compilation lazy, quando non viene passata alcuna firma numba.njit). Se la funzione verrà eseguita solo una volta, anche il tempo di compilazione dovrebbe essere considerato per la velocità di questo metodo. In genere vale la pena di compilare le funzioni solo se il tempo totale di compilazione + esecuzione è inferiore al runtime non compilato (che è effettivamente vero nel caso precedente, in cui la funzione nativa di Python è molto lenta). Ciò accade soprattutto quando si chiama la funzione compilata molte volte.

Come max9111 ha osservato in un commento, una caratteristica importante del numbaè la cacheparola chiave per jit. Passando cache=Truea numba.jit, la funzione compilata verrà archiviata su disco, in modo che durante la successiva esecuzione del modulo python dato la funzione verrà caricata da lì anziché ricompilata, il che può di nuovo risparmiare tempo di esecuzione nel lungo periodo.

Un approccio sarebbe quello di utilizzare Pandasqui puramente per utilizzare groupby. Ho gonfiato un po 'le dimensioni dell'input per dare una migliore comprensione dei tempi (poiché c'è un sovraccarico nella creazione del DF).

import numpy as np
import pandas as pd

data =  [1.00, 1.05, 1.30, 1.20, 1.06, 1.54, 1.33, 1.87, 1.67]
index = [0,    0,    1,    1,    1,    1,    2,    3,    3]

data = data * 500
index = np.sort(np.random.randint(0, 30, 4500))

def df_approach(data, index):
    df = pd.DataFrame({'data': data, 'label': index})
    df['median'] = df.groupby('label')['data'].transform('median')
    df['result'] = df['data'] - df['median']

Dà quanto segue timeit:

%timeit df_approach(data, index)
5.38 ms ± 50.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Per le stesse dimensioni del campione, ottengo che l' approccio dict di Aryerez sia:

%timeit dict_approach(data, index)
8.12 ms ± 3.47 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Tuttavia, se aumentiamo gli input di un altro fattore di 10, i tempi diventano:

%timeit df_approach(data, index)
7.72 ms ± 85 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit dict_approach(data, index)
30.2 ms ± 10.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Tuttavia, a scapito della ragionevolezza, la risposta di Divakar che utilizza puro intorpidimento arriva a:

%timeit bin_median_subtract(data, index)
573 µs ± 7.48 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Alla luce del nuovo set di dati (che avrebbe dovuto essere impostato all'inizio):

%timeit df_approach(data, groups)
472 ms ± 2.52 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit bin_median_subtract(data, groups) #https://stackoverflow.com/a/58788623/4799172
3.02 s ± 31.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit dict_approach(data, groups) #https://stackoverflow.com/a/58788199/4799172
<I gave up after 1 minute>

# jitted (using @numba.njit('f8[:](f8[:], i4[:]') on Windows) from  https://stackoverflow.com/a/58788635/4799172
%timeit diffmedian_jit(data, groups)
132 ms ± 3.12 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Forse l'hai già fatto, ma in caso contrario, vedi se è abbastanza veloce:

median_dict = {i: np.median(data[index == i]) for i in np.unique(index)}
def myFunc(my_dict, a): 
    return my_dict[a]
vect_func = np.vectorize(myFunc)
median_diff = data - vect_func(median_dict, index)
median_diff

Produzione:

array([-0.025,  0.025,  0.05 , -0.05 , -0.19 ,  0.29 ,  0.   ,  0.1  ,
   -0.1  ])

Ecco un approccio basato su NumPy per ottenere mediana binata per valori bin / indice positivi -

def bin_median(a, i):
    sidx = np.lexsort((a,i))

    a = a[sidx]
    i = i[sidx]

    c = np.bincount(i)
    c = c[c!=0]

    s1 = c//2

    e = c.cumsum()
    s1[1:] += e[:-1]

    firstval = a[s1-1]
    secondval = a[s1]
    out = np.where(c%2,secondval,(firstval+secondval)/2.0)
    return out

Per risolvere il nostro caso specifico di quelli sottratti -

def bin_median_subtract(a, i):
    sidx = np.lexsort((a,i))

    c = np.bincount(i)

    valid_mask = c!=0
    c = c[valid_mask]    

    e = c.cumsum()
    s1 = c//2
    s1[1:] += e[:-1]
    ssidx = sidx.argsort()
    starts = c%2+s1-1
    ends = s1

    starts_orgindx = sidx[np.searchsorted(sidx,starts,sorter=ssidx)]
    ends_orgindx  = sidx[np.searchsorted(sidx,ends,sorter=ssidx)]
    val = (a[starts_orgindx] + a[ends_orgindx])/2.
    out = a-np.repeat(val,c)
    return out