Швидка альтернатива numpy.median.reduceat

Що стосується цієї відповіді , чи існує швидкий спосіб обчислити медіани через масив, який містить групи з неоднаковою кількістю елементів?

Наприклад:

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

І тоді я хочу обчислити різницю між кількістю та медіаною на групу (наприклад, медіана групи 0- 1.025це перший результат 1.00 - 1.025 = -0.025). Отже, для масиву вище результати будуть виглядати як:

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

Оскільки np.median.reduceatне існує (поки що), чи існує інший швидкий спосіб досягти цього? Мій масив буде містити мільйони рядків, тому швидкість має вирішальне значення!

Індекси можна вважати суміжними та упорядкованими (їх легко перетворити, якщо їх немає).

Приклад даних для порівняння продуктивності:

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]

Іноді вам потрібно написати недіотичний код numpy, якщо ви дійсно хочете пришвидшити свій розрахунок, чого ви не можете зробити з рідною нумією.

numbaкомпілює ваш python-код на низькому рівні C. Оскільки велика кількість онімінь, як правило, так само швидко, як і C, це здебільшого стає корисним, якщо ваша проблема не піддається нативній векторизації з numpy. Це один із прикладів (де я припускав, що індекси є суміжними та сортованими, що також відображається в даних прикладу):

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

Ось кілька таймінгів із використанням %timeitмагії 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)

Використовуючи оновлені приклади даних у питанні, ці числа (тобто час виконання функції python порівняно з часом виконання функціонування, прискореного JIT):

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

Це становить 65-кратну швидкість у меншому випадку та 26-кратну прискорення у більшому випадку (порівняно з кодом повільного циклу), звичайно) з використанням прискореного коду. Інша перевага полягає в тому, що (на відміну від типової векторизації з власним нумером) нам не була потрібна додаткова пам'ять для досягнення цієї швидкості, це все про оптимізований і складений код низького рівня, який в кінцевому підсумку запускається.

Вищенаведена функція передбачає, що numpy int масиви int64за замовчуванням, що насправді не так у Windows. Отже, альтернативою є видалення підпису з виклику до numba.njit, запускаючи належну своєчасну компіляцію. Але це означає, що функція буде скомпільована під час першого виконання, що може поєднуватися з результатами часу (ми можемо виконати функцію один раз вручну, використовуючи представницькі типи даних, або просто прийняти, що перше виконання часу буде набагато повільніше, що повинно ігнорувати). Це саме те, що я намагався запобігти, вказавши підпис, що запускає дострокову компіляцію.

У будь-якому випадку, в належному випадку JIT декоратор, який нам потрібен, є просто

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

Зауважте, що вищезазначені таймінги, які я показав для функції, складеної jit, застосовуються лише після того, як функція була скомпільована. Це або відбувається при визначенні (з нетерплячим складанням, коли явний підпис передається доnumba.njit ), або під час першого виклику функції (при лінивій компіляції, коли підпис не передається numba.njit). Якщо функція буде виконуватися лише один раз, час компіляції також слід враховувати для швидкості цього методу. Як правило, варто лише компілювати функції, якщо загальний час компіляції + виконання менший, ніж некомпільований час виконання (що насправді вірно у вищенаведеному випадку, коли нативна функція пітона дуже повільна). В основному це відбувається, коли ви багато разів викликаєте свою компільовану функцію.

Як max9111 відзначена в коментарях, одна важлива особливістю numbaє cacheключовим словом , щоб jit. Перехід cache=Trueдо numba.jitзбереже складену функцію на диск, так що під час наступного виконання даного модуля python функція буде завантажена звідти, а не перекомпільована, що знову може позбавити вас часу виконання.

Одним із підходів було б використання Pandasтут виключно для використання groupby. Я трохи завищив розміри вхідних даних, щоб краще зрозуміти терміни (оскільки у створенні 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']

Дає наступне timeit :

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

За однакового розміру вибірки я отримую дікт-підхід Арієреса є таким:

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

Однак якщо ми збільшимо вхід на інший коефіцієнт 10, терміни стануть:

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

Однак, за рахунок певної реалістичності , відповідь Дівакара, що використовує чисту нумеру, надходить на:

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

У світлі нового набору даних (який дійсно повинен був бути встановлений на початку):

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

Можливо, ви вже це зробили, але якщо ні, то подивіться, чи це досить швидко:

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

Вихід:

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

Ось підхід, заснований на NumPy, щоб отримати медіану для позитивних значень бін / індексу -

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

Щоб вирішити наш конкретний випадок віднятих -

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