python - 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 进行 native 矢量化，这通常会很有用。这是一个示例(我假设索引是连续且已排序的，这也反射(reflect)在示例数据中):

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 倍(当然，与缓慢的循环代码相比)。另一个好处是(与使用 native numpy 的典型向量化不同)我们不需要额外的内存来实现这个速度，这完全是关于最终运行的优化和编译的低级代码。

上述函数假设 numpy int 数组为 int64默认情况下，在 Windows 上实际上并非如此。因此，另一种方法是从对 numba.njit 的调用中删除签名。 ，触发正确的即时编译。但这意味着函数将在第一次执行期间编译，这可能会干扰计时结果(我们可以手动执行一次函数，使用代表性数据类型，或者只是接受第一次计时执行会慢得多，这应该被忽略)。这正是我试图通过指定一个触发提前编译的签名来防止的。

无论如何，在正确的 JIT 情况下，我们需要的装饰器只是

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

请注意，我为 jit 编译的函数显示的上述时间仅适用于编译函数后。这要么发生在定义时(使用急切编译，当显式签名传递给 numba.njit 时)，或者发生在第一个函数调用期间(使用惰性编译，当没有签名传递给 numba.njit 时)。如果函数只执行一次，那么编译时间也应该考虑这个方法的速度。如果编译 + 执行的总时间小于未编译的运行时，通常只值得编译函数(在上述情况下实际上是这样， native python 函数非常慢)。这主要发生在您多次调用编译函数时。

如 max9111在评论中指出， numba 的一个重要特征是 cache keyword至 jit .路过 cache=True至 numba.jit将编译后的函数存储到磁盘，以便在给定 python 模块的下一次执行期间，该函数将从那里加载而不是重新编译，从长远来看，这再次可以节省您的运行时间。