2D滚动窗口分位数的最快方法？

Posted 2023-03-11

【中文标题】2D滚动窗口分位数的最快方法？

【英文标题】：Fastest way for 2D rolling window quantile?

【发布时间】：2020-05-22 19:30:15

【问题描述】：

我想按列计算尺寸为 (1e6, 1e5) 的大型二维矩阵的滚动分位数。我正在寻找最快的方法，因为我需要执行此操作数千次，而且计算量非常大。对于实验，使用 window=1000 和 q=0.1。

import numpy as np
import pandas as pd
import multiprocessing as mp
from functools import partial
import numba as nb
X = np.random.random((10000,1000)) # Original array has dimensions of about (1e6, 1e5)

我目前的做法：

熊猫：%timeit: 5.8 s ± 15.5 ms per loop

def pd_rolling_quantile(X, window, q):
    return pd.DataFrame(X).rolling(window).quantile(quantile=q)

Numpy 跨步：%timeit: 2min 42s ± 3.29 s per loop

def strided_app(a, L, S):
    nrows = ((a.size-L)//S)+1
    n = a.strides[0]
    return np.lib.stride_tricks.as_strided(a, shape=(nrows,L), strides=(S*n,n))
def np_1d(x, window, q):
    return np.pad(np.percentile(strided_app(x, window, 1), q*100, axis=-1), (window-1, 0) , mode='constant')
def np_rolling_quantile(X, window, q):
    results = []
    for i in np.arange(X.shape[1]):
        results.append(np_1d(X[:,i], window, q))
    return np.column_stack(results)

多处理：%timeit: 1.13 s ± 27.6 ms per loop

def mp_rolling_quantile(X, window, q):
    pool = mp.Pool(processes=12)
    results = pool.map(partial(pd_rolling_quantile, window=window, q=q), [X[:,i] for i in np.arange(X.shape[1])])
    pool.close()
    pool.join()
    return np.column_stack(results)

Numba：%timeit: 2min 28s ± 182 ms per loop

@nb.njit
def nb_1d(x, window, q):
    out = np.zeros(x.shape[0])
    for i in np.arange(x.shape[0]-window+1)+window:
        out[i-1] = np.quantile(x[i-window:i], q=q)
    return out
def nb_rolling_quantile(X, window, q):
    results = []
    for i in np.arange(X.shape[1]):
        results.append(nb_1d(X[:,i], window, q))
    return np.column_stack(results)

时机不是很好，理想情况下，我的目标是提高 10-50 倍的速度。我将不胜感激任何建议，如何加快速度。也许有人对使用较低级别的语言（Cython）或其他方式来使用基于 Numpy/Numba/Tensorflow 的方法加速它有想法。谢谢！

【参考方案1】：

我会推荐新的rolling-quantiles package。 为了证明，即使是为每列构造一个单独的过滤器的有点幼稚的方法也优于上述单线程pandas 实验：

pipes = [rq.Pipeline(rq.LowPass(window=1000, quantile=0.1)) for i in range(1000)]
%timeit [pipe.feed(X[:, i]) for i, pipe in enumerate(pipes)]
1.34 s ± 7.76 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

对

df = pd.DataFrame(X)
%timeit df.rolling(1000).quantile(0.1)
5.63 s ± 27 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

正如您所展示的，两者都可以通过multiprocessing 简单地并行化。