Как эффективно работать с комбинациями по одному измерению в массиве np

Question

Дано S быть n x m матрица, как массив NumPy, я хочу, чтобы вызвать функцию f на пары (S[i], S[j]) рассчитать конкретное значение интереса, хранящееся в матрице с размерами n x n, В моем конкретном случае функция f коммутативно так f(x,y) = f(y,x),

Имея все это в виду, мне интересно, могу ли я сделать какие-то трюки, чтобы максимально ускорить это, n может быть довольно большим.

Когда я измеряю время функции f, она составляет пару микросекунд, как и ожидалось. Это довольно простой расчет. Ниже я покажу вам время, которое я получил, по сравнению с max() а также sum() для справки.

In [19]: %timeit sum(s[11])
4.68 µs ± 56.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [20]: %timeit max(s[11])
3.61 µs ± 64.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [21]: %timeit f(s[11], s[21], 50, 10, 1e-5)
1.23 µs ± 7.25 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

In [22]: %timeit f(s[121], s[321], 50, 10, 1e-5)
1.26 µs ± 31.1 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

Однако, когда я рассчитываю общее время обработки для данных выборки 500x50 (в результате получается сравнение 500 x 500 /2 = 125K), общее время значительно увеличивается (в минутах). Я бы ожидал что-то вроде 0,2-0,3 секунды (1,25 * 2E-6 с / расч.).

In [12]: @jit
    ...: def testf(s, n, m, p):
    ...:     tol = 1e-5
    ...:     sim = np.zeros((n,n))
    ...:     for i in range(n):
    ...:         for j in range(n):
    ...:             if i > j:
    ...:                 delta = p[i] - p[j]
    ...:                 if delta < 0:
    ...:                     res = f(s[i], s[j], m, abs(delta), tol) # <-- max(s[i])
    ...:                 else:
    ...:                     res = f(s[j], s[i], m, delta, tol) # <-- sum(s[j])
    ...:                 sim[i][j] = res
    ...:                 sim[j][i] = res
    ...:     return sim

В коде выше я изменил строки, где res назначен на max() а также sum() (закомментированные части) для тестирования, и код выполняется примерно в 100 раз быстрее, хотя сами функции медленнее по сравнению с моей функцией f()

Что подводит меня к моим вопросам:

1. Можно ли избежать двойной петли, чтобы ускорить это? В идеале я хочу иметь возможность запустить это для матриц, где размер n = 1E5. (Комментарий: поскольку функции max и sum работают значительно быстрее, я предполагаю, что циклы for здесь не являются узким местом, но все же полезно знать, есть ли лучший способ)

2. Что может вызвать серьезное замедление с моей функцией, если это не двойной цикл for?

РЕДАКТИРОВАТЬ

Специфика функции f была задана некоторыми комментариями. Он перебирает два массива и проверяет количество значений в двух массивах, которые "достаточно близки". Я удалил комментарии и изменил некоторые имена переменных, но логика такая, как показано ниже. Было интересно отметить, что math.isclose(x,y,rel_tol) что эквивалентно операторам if, которые у меня есть ниже, делает код значительно медленнее, возможно, из-за библиотечного вызова?

from numba import jit
@jit 
def f(arr1, arr2, n, d, rel_tol):

    counter = 0
    i,j,k = 0,0,0   
    while (i < n and j < n and k < n):
        val = arr1[j] + d
        if abs(arr1[i] - arr2[k]) < rel_tol * max(arr1[i], arr2[k]):
            counter += 1
            i += 1
            k += 1
        elif abs(val - arr2[k]) < rel_tol * max(val, arr2[k]):
            counter += 1
            j += 1
            k += 1
        else: 
            # incremenet the index corresponding to the lightest
            if arr1[i] <= arr2[k] and arr1[i] <= val:
                if i < n: 
                    i += 1
            elif val <= arr1[i] and val <= arr2[k]:
                if j < n:
                    j += 1
            else: 
                k += 1
    return counter