Оптимизируйте идентификацию квантилей по столбцам массива
У меня есть массив A (размером m x n) и процент p в [0,1]. Мне нужно создать логический массив m x nB с True в записи (i,j), если A[i,j] находится в p^{th} квантиле столбца A[:,j].
Вот код, который я использовал до сих пор.
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)
Ответы 2
Простой способ ускорить этот код — запустить его параллельно с помощью 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))
На моей машине (6-ядерный процессор i5-9600KF) это решение занимает 0,23 мс вместо 29 мс для исходного кода Numpy. Это примерно в 130 раз быстрее.
Обратите внимание, что компиляция функции на моей машине занимает несколько секунд. Если это слишком дорого, альтернативой является написание кода с использованием Cython. Однако обратите внимание, что эта альтернатива требует написания функции np.quantile (поскольку Cython не может ускорить функции Numpy).
Более быстрое решение состоит в написании реализации квантиля, совместимой с SIMD (например, на основе сортировки Bitonic). Однако сделать это (легко) в Numba на самом деле невозможно. Таким образом, безусловно, необходимо реализовать это на родном языке (с поддержкой операций SIMD). Обратите внимание, что сделать это эффективно — непростая задача даже для опытных разработчиков.
Я не уверен, что это намного быстрее, но вы должны, по крайней мере, знать, что numpy.quantile имеет ключевой аргумент axis, поэтому вы можете вычислить все квантили с помощью одной команды:
quant_levels = np.quantile(A, p, axis=0)
B = (A >= quant_levels)
Быстрый тест показывает, что в вашем примере это происходит примерно в 28 раз быстрее.