À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.
numba
compila 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 %timeit
magia 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 int64
por 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 cache
palavra - chave para jit
. Passar cache=True
para numba.jit
armazena 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.
scipy.ndimage.median
sugestão na resposta vinculada? Não me parece que ele precise de um número igual de elementos por rótulo. Ou eu perdi alguma coisa?