Szybka alternatywa dla numpy.median.reduceat

Czy w związku z tą odpowiedzią istnieje szybki sposób na obliczenie median na podstawie tablicy zawierającej grupy o nierównej liczbie elementów?

Na przykład:

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

Następnie chcę obliczyć różnicę między liczbą a medianą na grupę (np. Mediana grupy 0jest 1.025pierwszym wynikiem 1.00 - 1.025 = -0.025). Tak więc dla powyższej tablicy wyniki wyglądałyby następująco:

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

Ponieważ np.median.reduceatnie istnieje (jeszcze), czy istnieje inny szybki sposób na osiągnięcie tego? Moja tablica będzie zawierać miliony wierszy, więc szybkość ma kluczowe znaczenie!

Można założyć, że indeksy są ciągłe i uporządkowane (łatwo je przekształcić, jeśli nie są).

Przykładowe dane do porównań wydajności:

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]

Czasami musisz napisać nieidiomatyczny kod numpy, jeśli naprawdę chcesz przyspieszyć swoje obliczenia, czego nie możesz zrobić z natywną numpy.

numbakompiluje kod Pythona do niskiego poziomu C. Ponieważ sama liczba numpy jest zwykle tak szybka jak C, zwykle jest to przydatne, jeśli twój problem nie nadaje się do natywnej wektoryzacji za pomocą numpy. To jest jeden przykład (gdzie założyłem, że indeksy są ciągłe i posortowane, co znajduje również odzwierciedlenie w przykładowych danych):

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

A oto niektóre czasy z wykorzystaniem %timeitmagii 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)

Wykorzystując zaktualizowane przykładowe dane w pytaniu, te liczby (tj. Środowisko wykonawcze funkcji python vs. środowisko uruchomieniowe funkcji przyspieszonej przez JIT) są

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

Odpowiada to przyspieszeniu 65-krotnemu w mniejszym przypadku i 26-krotnym przyspieszeniu w większym przypadku (oczywiście w porównaniu z kodem wolnej pętli) przy użyciu kodu przyspieszonego. Kolejną zaletą jest to, że (w przeciwieństwie do typowej wektoryzacji z natywną numpy) nie potrzebowaliśmy dodatkowej pamięci, aby osiągnąć tę prędkość, chodzi o zoptymalizowany i skompilowany kod niskiego poziomu, który ostatecznie jest uruchamiany.

Powyższa funkcja zakłada, że ​​tablice int numpy są int64domyślnie, co w rzeczywistości nie ma miejsca w systemie Windows. Alternatywą jest więc usunięcie podpisu z wywołania do numba.njit, uruchamiając odpowiednią kompilację just-in-time. Oznacza to jednak, że funkcja zostanie skompilowana podczas pierwszego wykonania, co może wtrącać się w wyniki pomiaru czasu (możemy wykonać funkcję raz ręcznie, używając reprezentatywnych typów danych, lub po prostu zaakceptować, że pierwsze wykonanie pomiaru czasu będzie znacznie wolniejsze, co powinno być zignorowanym). Właśnie tego próbowałem zapobiec, określając podpis, który uruchamia kompilację z wyprzedzeniem.

W każdym razie, w przypadku właściwego JIT, dekorator, którego potrzebujemy, jest po prostu

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

Zauważ, że powyższe czasy, które pokazałem dla funkcji skompilowanej z użyciem Jit, mają zastosowanie tylko po skompilowaniu funkcji. Dzieje się tak w momencie definicji (przy gorącej kompilacji, gdy przekazywany jest wyraźny podpis numba.njit) lub podczas pierwszego wywołania funkcji (przy leniwej kompilacji, gdy nie jest przekazywany żaden podpis numba.njit). Jeśli funkcja ma zostać wykonana tylko raz, należy również wziąć pod uwagę czas kompilacji dla szybkości tej metody. Zazwyczaj warto kompilować funkcje tylko wtedy, gdy całkowity czas kompilacji + wykonania jest krótszy niż nieskompilowany środowisko wykonawcze (co w rzeczywistości jest prawdą w powyższym przypadku, gdy natywna funkcja python jest bardzo wolna). Dzieje się tak głównie wtedy, gdy wywołujesz skompilowaną funkcję wiele razy.

Jak zauważono w komentarzu max9111 , jedną ważną cechą numbajest cachesłowo kluczowe to jit. Przekazanie cache=Truedo numba.jitspowoduje zapisanie skompilowanej funkcji na dysku, dzięki czemu podczas następnego wykonania danego modułu Pythona funkcja zostanie stamtąd załadowana, a nie ponownie skompilowana, co w dłuższej perspektywie może zaoszczędzić ci czasu działania.

Jednym z podejść byłoby użycie Pandastutaj wyłącznie w celu wykorzystania groupby. Podniosłem nieco rozmiary wejściowe, aby lepiej zrozumieć czasy (ponieważ tworzenie DF wiąże się z dodatkowymi kosztami).

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']

Daje następujące timeit:

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

W przypadku tej samej wielkości próbki podejście Aryerez brzmi :

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

Jeśli jednak zwiększymy nakłady o kolejny współczynnik 10, czasy stają się:

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

Jednak kosztem pewnego reability, odpowiedź przez Divakar użyciu czystego numpy znalazł się na:

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

W świetle nowego zestawu danych (który naprawdę powinien był zostać ustawiony na początku):

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

Może już to zrobiłeś, ale jeśli nie, sprawdź, czy to wystarczająco szybko:

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

Wynik:

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

Oto podejście oparte na NumPy, aby uzyskać medianę binned dla dodatnich wartości bin / wartości indeksu -

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

Aby rozwiązać nasz konkretny przypadek odejmowanych -

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