[python] numpy.median.reduceat에 대한 빠른 대안

때때로 당신은 당신이 경우 비 관용적 NumPy와 코드를 작성할 필요가 정말 당신이 기본 NumPy와 함께 할 수없는 당신의 계산 속도를합니다.

numba파이썬 코드를 저수준 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

다음은 IPython의 %timeit마법을 사용한 몇 가지 타이밍입니다 .

>>> %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를 사용하는 일반적인 벡터화와 달리)이 속도를 달성하기 위해 추가 메모리가 필요하지 않다는 것입니다. 이것은 결국 실행되고 결국 최적화되고 컴파일 된 저수준 코드에 관한 것입니다.

위 함수는 numpy int 배열이 int64기본적으로 가정하지만 실제로는 Windows에서는 그렇지 않습니다. 대안은에 대한 호출에서 서명을 제거하여 numba.njit적절한 적시 컴파일을 트리거하는 것입니다. 그러나 이것은 첫 번째 실행 중에 함수가 컴파일되어 타이밍 결과와 충돌 할 수 있음을 의미합니다 (대표 데이터 유형을 사용하여 수동으로 한 번 함수를 한 번 실행하거나 첫 번째 타이밍 실행이 훨씬 느리다는 것을 받아 들일 수 있음) 무시하십시오). 이것은 미리 컴파일을 트리거하는 서명을 지정하여 방지하려고했던 것입니다.

어쨌든, 적절한 JIT 케이스에서 우리가 필요로하는 데코레이터는

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

jit 컴파일 된 함수에 대해 위에서 보여준 타이밍은 함수가 컴파일 된 후에 만 ​​적용됩니다. 이는 정의 (열심 한 컴파일, 명시 적 서명이 전달 될 때 numba.njit) 또는 첫 번째 함수 호출 (지연이없는 서명이 전달 될 때 numba.njit)에서 발생합니다. 함수가 한 번만 실행될 예정이라면이 방법의 속도에 대해 컴파일 시간도 고려해야합니다. 일반적으로 컴파일 + 실행의 총 시간이 컴파일되지 않은 런타임 (기본 파이썬 함수가 매우 느린 위의 경우에 해당)보다 짧은 경우에만 함수를 컴파일하는 것이 좋습니다. 컴파일 된 함수를 여러 번 호출 할 때 주로 발생합니다.

으로 max9111이 코멘트에 언급, 하나 개 중요한 기능 numba은 IS cache키워드 에 jit. 전달 cache=True하기 위해 numba.jit, 디스크에 컴파일 된 기능을 저장합니다 그래서 당신은 장기적으로이 아닌 재 컴파일을 다시 마련 할 수있는 런타임에서 주어진 파이썬 모듈의 다음 실행시에 기능이로드됩니다.

답변

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

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

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