Блог

Главная - Вопросы по программированию - Короткое замыкание np.all с вложенным np.less для сравнения больших массивов в numpy

Короткое замыкание np.all с вложенным np.less для сравнения больших массивов в numpy

#python-3.x #numpy #numba

#python-3.x #numpy #numba

Вопрос:

В моем текущем коде (см. MWE) у меня есть узкое место, где я выполняю np.all с вложенным np.less для больших 2D-массивов. Я знаю, что если в файле есть одно false значение, np.less мы можем прекратить проверку, потому что остальные значения в индексе будут вычисляться false кодом (поскольку я объединяю все значения в одном индексе заданного измерения вместе).

Есть ли способ с помощью numba или numpy, где я могу использовать это условие «раннего выхода / короткого замыкания» для создания значимого ускорения в этом вычислении?

Предпоследняя строка в MWE — это то, что я пытаюсь ускорить. Пожалуйста, обратите N внимание, что и M может быть очень большим, но только очень немногие сравнения будут фактически оцениваться true .

 import numpy as np

N = 10000
M = 10 # Reduced to small value to show that sometimes the comparisons evaluate to 'True'

array = np.random.uniform(low=0.0, high=10.0, size=(N, M))
comparison_array = np.random.uniform(low=0.0, high=10.0, size=(M))

# Can we apply an early exit condition on this?
mask = np.all(np.less(array, comparison_array), axis=-1)

print(f"Number of 'True' comparisons: {np.sum(mask)}")

Ответ №1:

Вот numba версия, достаточно разработанная для работы, не обязательно оптимизированная:

 @numba.njit
def foo(arr, carr):
    N, M = arr.shape
    mask = np.ones(N, dtype=np.bool_)
    for i in range(N):
        for j in range(M):
            if arr[i,j]>=carr[j]:
                mask[i]=False
                break
    return mask

Тестирование:

 In [178]: np.sum(foo(array, comparison_array))
Out[178]: 2
In [179]: np.sum(np.all(np.less(array, comparison_array), axis=1))
Out[179]: 2

время:

 In [180]: timeit np.sum(foo(array, comparison_array))
155 µs ± 6.36 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [181]: timeit np.sum(np.all(np.less(array, comparison_array), axis=1))
451 µs ± 5.19 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Это приличное улучшение.

Комментарии:

1. Вау, это отличное улучшение! На чем мне следует сосредоточиться, чтобы использовать ваше решение для его дальнейшей оптимизации?