Numpy: применить маску к значениям, затем принять среднее значение, но параллельно
У меня есть массив значений 1d:
v = np.array([0, 1, 4, 0, 5])
Кроме того, у меня есть двумерный массив логических масок (в производстве существуют миллионы масок):
m = np.array([
[True, True, False, False, False],
[True, False, True, False, True],
[True, True, True, True, True],
])
Я хочу применить каждую строку маски к массиву v, а затем вычислить среднее значение замаскированных значений.
Ожидаемое поведение:
results = []
for mask in m:
results.append(np.mean(v[mask]))
print(results) # [0.5, 3.0, 2.0]
Легко сделать последовательно, но я уверен, что параллельно есть красивый вариант? Одно решение, которое я нашел:
mask = np.ones(m.shape)
mask[~m] = np.nan
np.nanmean(v * mask, axis=1) # [0.5, 3.0, 2.0]
Есть ли другое решение, возможно, с использованием модуля np.ma? Я ищу решение, которое будет быстрее, чем два моих текущих решения.
Содержит ли ваш реальный массив по 5 элементов в строке? Numpy действительно не оптимизирован для этого. Это 32-битное целое число, как в вашем примере?
Обратите внимание, что маскированные массивы обычно не очень эффективны. В первую очередь они предназначены для удобства. Создание временных массивов в этом случае очень затратно, поскольку операция ограничена памятью. В этом случае могут помочь Numba/Cython, но они могут быть явно неоптимальны по другим причинам (все равно намного быстрее, чем наивный код Numpy). Также обратите внимание, что функции Numpy не являются многопоточными (только функции BLAS). Однако функции Numpy могут выиграть от параллелизма SIMD.
Формы реального мира могут иметь вид v = (1 000 000) и m = (10 000, 1 000 000). Память не проблема. И dtype имеет значение float
«Память не проблема» Я имею в виду, что проблемой является пропускная способность памяти, а не ее использование. Пропускная способность памяти довольно ограничена на большинстве основных платформ.
Ответы 2
Я думаю, что самый чистый векторизованный подход будет примерно таким:
result = np.broadcast_to(v, m.shape).mean(axis=1, where=m)
Однако это предполагает явную трансляцию v в форму m, поэтому в зависимости от ограничений памяти это может быть неоптимально.
Отлично, это самый быстрый на данный момент! Память не проблема
трансляция создает просмотры, поэтому не мешает использованию памяти.
Более быстрый код Numpy
Более быстрый способ сделать это с помощью Numpy — выполнить умножение матриц:
(m @ v) / m.sum(axis=1)
Мы можем оптимизировать это дальше, избегая неявного преобразования и выполняя суммирование с 8-битным целым числом (это безопасно только потому, что v.shape[1] мало — т. е. меньше 127):
(m @ v) / m.view(np.int8).sum(dtype=np.int8, axis=1)
Еще более быстрый код с многопоточностью Numba
Мы можем использовать Numba для написания аналогичной функции и даже использовать несколько потоков:
import numba as nb
@nb.njit('(float32[::1], bool_[:,::1])', parallel=True)
def compute(v, m):
si, sj = m.shape
res = np.empty(si, dtype=np.float32)
for i in nb.prange(si):
s = np.float32(0)
count = 0
for j in range(sj):
if m[i, j]:
s += v[j]
count += 1
if count > 0:
res[i] = s / count
else:
res[i] = np.nan
return res
Полученные результаты
Вот результаты на моем процессоре i5-9600KF (с 6 ядрами):
Initial vectorized code: 136 ms
jakevdp's solution: 74 ms
Numpy matmul: 28 ms
Numba sequential code: 21 ms
Optimized Numpy matmul: 20 ms <-----
Numba parallel code: 4 ms <-----
Последовательная реализация Numba, к сожалению, не намного лучше, чем реализация Numpy (в этом случае Numba не генерирует очень хороший код), но она хорошо масштабируется, поэтому параллельная версия работает значительно быстрее.
Параллельная версия Numba в 34 раза быстрее, чем исходный векторизованный код, и в 18 раз быстрее, чем решение jakevdp. Оптимизированное Numpy решение на основе умножения матриц в 3,7 раза быстрее решения jakevdp.
Обратите внимание, что код Numba также немного лучше, чем решение jakevdp и оптимизированное решение Numpy, поскольку он поддерживает случай, когда среднее значение недопустимо (NaN), без вывода предупреждения о делении на 0.
Нашел ответ, используя np.ma: np.ma.array(v * m, Mask=~m).mean(axis=1).data Но все равно решение медленнее, чем я надеялся (в моем случае значения и маски очень большой).