优化沿数组列的分位数识别

Question

我有一个数组 A（大小为 m x n），以及 (0,1) 中的百分比 p。我需要生成一个 m x n 布尔数组 B，如果 A[i,j] 位于 A[:,j] 列的 p^{th} 分位数中，则 (i,j) 条目中包含“True”。

这是我迄今为止使用过的代码。

import numpy as np

m = 200
n = 300

A = np.random.rand(m, n)

p = 0.3

quant_levels = np.zeros(n)
 
for i in range(n):
    quant_levels[i] = np.quantile(A[:,i],p)
    
B = np.array(A >= quant_levels)

非常感谢！

Answer 1

使此代码更快的一个简单方法是使用 Numba 并行运行它。这也减少了很多 Numpy 开销，这似乎是这里的瓶颈。

import numba as nb

@nb.njit('(float64[:,:], float64)', parallel=True)
def compute_quantiles(A, p):
    quant_levels = np.empty(n)
 
    for i in nb.prange(n):
        quant_levels[i] = np.quantile(A[:,i],p)
    
    return quant_levels

B = np.array(A >= compute_quantiles(A, p))

在我的机器上，Numpy 解决方案需要 0.23 毫秒，而不是 29 毫秒。在我的 i5-9600KF CPU（6 核）上，这大约快了

130 倍。

---

更快的方法是使用 SIMD 友好的分位数实现（例如基于 Bitonic 排序）。然而，在 Numba 中这实际上不可能（容易）做到这一点。因此，我们当然需要用一种本地语言来实现（支持 SIMD 操作）。请注意，即使对于静止的开发人员来说，高效地做到这一点也不是一项简单的任务。