Alternativa rápida para numpy.median.reduceat

Em relação a esta resposta , existe uma maneira rápida de calcular medianas em uma matriz que possui grupos com um número desigual de elementos?

Por exemplo:

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 então eu quero calcular a diferença entre o número e a mediana por grupo (por exemplo, mediana do grupo 0é 1.025o primeiro resultado 1.00 - 1.025 = -0.025). Portanto, para a matriz acima, os resultados apareceriam como:

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

Como np.median.reduceatainda não existe, existe outra maneira rápida de conseguir isso? Minha matriz conterá milhões de linhas, portanto a velocidade é crucial!

Pode-se supor que os índices sejam contíguos e ordenados (é fácil transformá-los se não forem).

Dados de exemplo para comparações de desempenho:

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]

Às vezes, você precisa escrever um código numpy não-idiomático, se realmente deseja acelerar seu cálculo, o que não pode ser feito com o numpy nativo.

numbacompila seu código python para o nível C. C. Como muitos numpy em si são geralmente tão rápidos quanto C, isso acaba sendo útil se o seu problema não se presta à vetorização nativa com numpy. Este é um exemplo (onde eu assumi que os índices são contíguos e classificados, o que também é refletido nos dados de exemplo):

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 aqui estão alguns horários usando a %timeitmagia do 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 os dados de exemplo atualizados na pergunta, esses números (ou seja, o tempo de execução da função python vs. o tempo de execução da função acelerada por JIT) são

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

Isso equivale a uma aceleração de 65x no caso menor e uma aceleração de 26x no caso maior (comparado ao código de loop lento, é claro) usando o código acelerado. Outra vantagem é que (ao contrário da vetorização típica com numpy nativo), não precisamos de memória adicional para atingir essa velocidade, trata-se de código de baixo nível otimizado e compilado que acaba sendo executado.

A função acima pressupõe que matrizes numpy int sejam int64por padrão, o que não é realmente o caso no Windows. Portanto, uma alternativa é remover a assinatura da chamada para numba.njit, acionando uma compilação just-in-time adequada. Mas isso significa que a função será compilada durante a primeira execução, o que pode interferir nos resultados de temporização (podemos executar a função uma vez manualmente, usando tipos de dados representativos, ou apenas aceitar que a primeira execução de temporização será muito mais lenta, o que deve ser ignorado). Isso é exatamente o que tentei impedir especificando uma assinatura, que aciona a compilação antecipada.

De qualquer forma, no caso JIT, o decorador de que precisamos é apenas

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

Observe que os tempos mostrados acima para a função jit-compiled somente se aplicam quando a função foi compilada. Isso acontece na definição (com compilação ansiosa, quando uma assinatura explícita é passada para numba.njit) ou durante a primeira chamada de função (com compilação lenta, quando nenhuma assinatura é passada para numba.njit). Se a função for executada apenas uma vez, o tempo de compilação também deve ser considerado para a velocidade desse método. Geralmente, vale a pena compilar funções se o tempo total de compilação + execução for menor que o tempo de execução não compilado (o que é realmente verdade no caso acima, onde a função python nativa é muito lenta). Isso acontece principalmente quando você está chamando sua função compilada várias vezes.

Como max9111 observou em um comentário, um recurso importante numbaé a cachepalavra - chave para jit. Passar cache=Truepara numba.jitarmazena a função compilada em disco, para que durante a próxima execução do módulo python fornecido, a função seja carregada a partir daí, em vez de recompilada, o que novamente pode poupar o tempo de execução a longo prazo.

Uma abordagem seria usar Pandasaqui puramente para fazer uso groupby. Aumentei um pouco os tamanhos de entrada para entender melhor os tempos (já que há sobrecarga na criação do 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']

Fornece o seguinte timeit:

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

Para o mesmo tamanho de amostra, obtenho a abordagem de ditado de Aryerez :

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

No entanto, se aumentarmos as entradas por outro fator de 10, os tempos se tornarão:

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

No entanto, às custas de alguma reatibilidade, a resposta de Divakar usando numpy puro chega em:

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

À luz do novo conjunto de dados (que realmente deveria ter sido definido no início):

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

Talvez você já tenha feito isso, mas se não, veja se isso é rápido o suficiente:

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

Resultado:

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

Aqui está uma abordagem baseada no NumPy para obter mediana binned para valores positivos de caixas / í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 nosso caso específico de subtraídos -

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