Alternativa rápida para numpy.median.reduceat

En relación con esta respuesta , ¿hay una manera rápida de calcular medianas sobre una matriz que tiene grupos con un número desigual de elementos?

P.ej:

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

Y luego quiero calcular la diferencia entre el número y la mediana por grupo (por ejemplo, la mediana del grupo 0es 1.025el primer resultado 1.00 - 1.025 = -0.025). Entonces, para la matriz anterior, los resultados aparecerían como:

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

Dado np.median.reduceatque no existe (todavía), ¿hay otra forma rápida de lograr esto? ¡Mi matriz contendrá millones de filas, por lo que la velocidad es crucial!

Se puede suponer que los índices son contiguos y ordenados (es fácil transformarlos si no lo son).

Datos de ejemplo para comparaciones de rendimiento:

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 veces necesita escribir código numpy no idiomático si realmente quiere acelerar su cálculo, lo que no puede hacer con numpy nativo.

numbacompila su código de Python a bajo nivel C. Dado que una gran cantidad de numpy suele ser tan rápido como C, esto generalmente resulta útil si su problema no se presta a la vectorización nativa con numpy. Este es un ejemplo (donde supuse que los índices son contiguos y ordenados, lo que también se refleja en los datos del ejemplo):

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

Y aquí hay algunos tiempos usando la %timeitmagia de 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)

Usando los datos de ejemplo actualizados en la pregunta, estos números (es decir, el tiempo de ejecución de la función python frente al tiempo de ejecución de la función acelerada por JIT) son

>>> %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)

Esto equivale a una aceleración de 65x en el caso más pequeño y una aceleración de 26x en el caso más grande (en comparación con el código de bucle lento, por supuesto) usando el código acelerado. Otra ventaja es que (a diferencia de la vectorización típica con numpy nativo) no necesitábamos memoria adicional para lograr esta velocidad, se trata de código de bajo nivel optimizado y compilado que termina ejecutándose.

La función anterior supone que las matrices numpy int son int64por defecto, lo cual no es el caso en Windows. Así que una alternativa es eliminar la firma de la llamada a numba.njit, lo que provocó adecuada compilación justo a tiempo. Pero esto significa que la función se compilará durante la primera ejecución, lo que puede interferir con los resultados de temporización (podemos ejecutar la función una vez manualmente, utilizando tipos de datos representativos, o simplemente aceptar que la primera ejecución de temporización será mucho más lenta, lo que debería ser ignorado). Esto es exactamente lo que intenté evitar al especificar una firma, que desencadena una compilación anticipada.

De todos modos, en el caso adecuado de JIT, el decorador que necesitamos es solo

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

Tenga en cuenta que los tiempos anteriores que mostré para la función compilada jit solo se aplican una vez que la función se ha compilado. Esto ocurre ya sea en la definición (con la compilación ansiosos, cuando una firma explícita se pasa a numba.njit), o durante la primera llamada de función (con la compilación perezoso, cuando no hay ninguna firma se pasa a numba.njit). Si la función solo se ejecutará una vez, el tiempo de compilación también debe considerarse para la velocidad de este método. Por lo general, solo vale la pena compilar funciones si el tiempo total de compilación + ejecución es menor que el tiempo de ejecución sin compilar (que en realidad es cierto en el caso anterior, donde la función nativa de Python es muy lenta). Esto ocurre principalmente cuando llama a su función compilada muchas veces.

Como max9111 señaló en un comentario, una característica importante de numbaes la cachepalabra clave to jit. Pasar cache=Truea numba.jitalmacenará la función compilada en el disco, de modo que durante la próxima ejecución del módulo python dado, la función se cargará desde allí en lugar de volver a compilarse, lo que de nuevo puede ahorrarle tiempo de ejecución a largo plazo.

Un enfoque sería usar Pandasaquí simplemente para usar groupby. He inflado un poco los tamaños de entrada para dar una mejor comprensión de los tiempos (ya que hay una sobrecarga en la creación 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']

Da lo siguiente timeit:

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

Para el mismo tamaño de muestra, obtengo el enfoque dict de Aryerez :

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

Sin embargo, si aumentamos las entradas por otro factor de 10, los tiempos se convierten en:

%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)

Sin embargo, a expensas de cierta fiabilidad, la respuesta de Divakar usando numpy puro llega a:

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

A la luz del nuevo conjunto de datos (que realmente debería haberse configurado al principio):

%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)

Tal vez ya hiciste esto, pero si no, mira si es lo suficientemente rápido:

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

Salida:

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

Aquí hay un enfoque basado en NumPy para obtener la mediana de binned para bins positivos / valores de índice:

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

Para resolver nuestro caso específico de los restados:

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