Быстрая альтернатива для 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, если вы действительно хотите ускорить вычисления, чего вы не можете сделать с помощью Numpy.

numbaкомпилирует ваш код Python в низкоуровневый C. Поскольку большая часть numpy обычно работает так же быстро, как 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-кратное ускорение в большем случае (по сравнению с медленным зацикливанием кода, конечно) с использованием ускоренного кода. Другим преимуществом является то, что (в отличие от типичной векторизации с нативной нативой) нам не требовалась дополнительная память для достижения этой скорости, все дело в оптимизированном и скомпилированном низкоуровневом коде, который в конечном итоге запускается.

Вышеприведенная функция предполагает, что int64по умолчанию используются пустые массивы int , чего нет в Windows. Таким образом, альтернативой является удаление подписи из вызова to numba.njit, запускающей правильную сборку точно в срок. Но это означает, что функция будет скомпилирована во время первого выполнения, что может повлиять на результаты синхронизации (мы можем либо выполнить функцию один раз вручную, используя репрезентативные типы данных, либо просто принять, что первое выполнение синхронизации будет намного медленнее, что должно быть проигнорированным). Это именно то, что я пытался предотвратить, указав сигнатуру, которая запускает преждевременную компиляцию.

Во всяком случае, в надлежащем случае JIT декоратор нам нужен просто

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

Обратите внимание, что приведенные выше моменты времени, которые я показал для функции jit-compiled, применяются только после того, как функция была скомпилирована. Это может происходить либо при определении (с готовой компиляцией, когда передается явная подпись numba.njit), либо во время первого вызова функции (с отложенной компиляцией, когда подпись не передается numba.njit). Если функция будет выполняться только один раз, то время компиляции также должно учитываться для скорости этого метода. Как правило, компиляция функций имеет смысл только в том случае, если общее время выполнения + компиляции меньше, чем некомпилированное время выполнения (что на самом деле верно в вышеупомянутом случае, когда собственная функция python очень медленная). В основном это происходит, когда вы часто вызываете скомпилированную функцию.

Как max9111 отмечена в комментариях, одна важная особенностью numbaявляется cacheключевым словом , чтобы jit. Передача cache=Trueto 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)

Тем не менее, за счет некоторой надежности, ответ Дивакара с использованием чистого numpy приходит в:

%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, чтобы получить усредненное значение для положительных значений bin / index -

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