numpy.median.reduceat的快速替代

Question

与this answer相关的是，是否有一种快速方法来计算数组上的中间值，该数组中的组具有不相等的元素数？

例如：

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]

Answer 1

有时，如果真的想加快计算速度，就需要编写非惯用的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

下面是使用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加速的functio的运行时)是

>>> %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时)。如果函数只执行一次，那么编译时间也应该考虑到该方法的速度。通常只有在编译+执行的总时间小于未编译的运行时(在上述情况下，本机python函数非常慢的情况下)，才值得编译函数。这通常发生在您多次调用已编译的函数时。

正如max9111在一篇评论中指出的，numba的一个重要特性是cache keyword to jit。将cache=True传递给numba.jit将将编译后的函数存储到磁盘，因此在下一次执行给定的python模块时，函数将从那里加载，而不是重新编译，这将再次节省您的运行时时间。

Answer 2

一种方法是在这里使用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)

对于相同的样本大小，我得到的dict approach of Aryerez是：

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

然而，以牺牲一些可重用性为代价，Divakar使用纯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)

Answer 3

也许你已经这么做了，但如果没有，看看这是否足够快：

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