Эффективное вычисление евклидовой матрицы распределения в Numpy?

Поскольку вы подразумевали, что вам не нужна полная квадратная матрица результатов, отметив, что cdist неудобен, потому что он дважды вычисляет попарные расстояния, вы можете использовать Numba для написания UDF, который рассчитывает только для нижнего или верхнего треугольника квадратной матрицы .

Обратите внимание, что при первом запуске есть накладные расходы из-за JIT-компиляции.

from scipy.spatial import distance
import pandas as pd
from numba import njit, prange
import numpy as np

@njit(parallel=True)
def euclidean_distance(coords1, coords2):
    # allocate output array
    c1_length, c2_length = len(coords1), len(coords2)
    out = np.empty(shape=(c1_length, c2_length), dtype=np.float64)

    # fill the lower triangle with euclidean distance formula
    # assuming coordiantes are (lat, lon) based on the example https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.cdist.html
    for lat_ix in prange(c1_length):
        for lon_ix in prange(c2_length):
            if lat_ix >= lon_ix: # do the reverse for the upper triangle
                out[lat_ix, lon_ix] = (
                    (coords1[lat_ix, 0] - coords2[lon_ix, 0]) ** 2
                    + (coords1[lat_ix, 1] - coords2[lon_ix, 1]) ** 2
                ) ** 0.5
            else:
                out[lat_ix, lon_ix] = 0
    return out


for n in [10, 100, 5000, 20000]:
    arr = np.random.normal(0, 100, (n, 2))
    print(n, arr.shape)

    %time out = euclidean_distance(arr, arr)
    %time out_cdist = distance.cdist(arr, arr, 'euclidean')

    if n < 1000:
        np.testing.assert_array_almost_equal(out, np.tril(out_cdist))
    print()

Выход:

10 (10, 2)
CPU times: user 987 ms, sys: 19.3 ms, total: 1.01 s
Wall time: 1.01 s
CPU times: user 79 µs, sys: 12 µs, total: 91 µs
Wall time: 95.1 µs

100 (100, 2)
CPU times: user 1.05 ms, sys: 404 µs, total: 1.45 ms
Wall time: 1.16 ms
CPU times: user 926 µs, sys: 254 µs, total: 1.18 ms
Wall time: 946 µs

5000 (5000, 2)
CPU times: user 125 ms, sys: 128 ms, total: 253 ms
Wall time: 75 ms
CPU times: user 184 ms, sys: 92.6 ms, total: 277 ms
Wall time: 287 ms

20000 (20000, 2)
CPU times: user 2.21 s, sys: 2.15 s, total: 4.36 s
Wall time: 2.55 s
CPU times: user 3.1 s, sys: 2.71 s, total: 5.81 s
Wall time: 31.9 s

С массивом из 20 000 элементов UDF работает немного быстрее, поскольку позволяет сэкономить половину вычислений. cdist кажется особенно / неожиданно медленным для этого конкретного распределения данных в масштабе на моем Macbook Air, но суть в том, что все равно.

Ограничение в этой проблеме - пропускная способность памяти.

Сначала попробуйте выполнить несколько простых операций с памятью, чтобы получить эталонное время.

import numba as nb
import numpy as np
from scipy.spatial import distance

#Should be at least 0.47 (SVML-Bug)
print(nb.__version__)

@nb.njit(fastmath=True,parallel=True)
def dist_simply_write(res):
    for i in nb.prange(A.shape[0]):
        for j in range(A.shape[0]):
            res[i,j]=1.
    return res

res_1=np.empty((A.shape[0],A.shape[0]))
res_2=np.empty((A.shape[0],A.shape[0]))

#Copying the array to a new array, which has to be allocated
%timeit res_2=np.copy(res_1)
#1.32 s ± 118 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

#Copying the array to a new array, which is already allocated
%timeit np.copyto(res_1,res_2)
#328 ms ± 14.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

#fill an array with 1., without calculating anything
%timeit out=dist_simply_write(A,res)
#246 ms ± 707 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Требуется ли больше времени для вычисления евклидова расстояния вместо того, чтобы записать 1.?

@nb.njit(fastmath=True,parallel=True)
def dist_arr_1(A):
    res=np.empty((A.shape[0],A.shape[0]))
    for i in nb.prange(A.shape[0]):
        for j in range(A.shape[0]):
            acc=0
            for k in range(A.shape[1]):
                acc+=(A[i,k]-A[j,k])**2
            res[i,j]=np.sqrt(acc)
    return res

@nb.njit(fastmath=True,parallel=True)
def dist_arr_2(A,res):
    for i in nb.prange(A.shape[0]):
        for j in range(A.shape[0]):
            acc=0
            for k in range(A.shape[1]):
                acc+=(A[i,k]-A[j,k])**2
            res[i,j]=np.sqrt(acc)
    return res

%timeit out=dist_arr_1(A)
#559 ms ± 85.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
res=np.empty((A.shape[0],A.shape[0]))

#If we can reuse the output memory
%timeit out=dist_arr_2(A,res)
#238 ms ± 4.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Как видите, совершенно не имеет значения, выполняем ли мы простые вычисления (евклидово расстояние) или записываем в массив только число. Вычисление только половины значений и их последующее копирование на самом деле медленнее (без непрерывной итерации в памяти и перезагрузки данных).

Я пробовал и numpy трансляцию, и scipy.spatial.distance.cdist, и оба они кажутся похожими, когда дело касается эффективности времени:

import numpy as np
from scipy.spatial.distance import cdist
import time

def dist_numpy(a, b):
    d = np.linalg.norm(a[:, None, :] - b[None, :, :], axis=2)
    d = np.transpose(d)
    sorted_d = np.sort(d)
    sorted_ind = np.argsort(d)
    return sorted_d, sorted_ind

def dist_scipy(a, b):
    d = cdist(a, b, 'euclidean')
    d = np.transpose(d)
    sorted_d = np.sort(d)
    sorted_ind = np.argsort(d)
    return sorted_d, sorted_ind

def get_a_b(r=10**4,c=10** 1):
    a = np.random.uniform(-1, 1, (r, c)).astype('f')
    b = np.random.uniform(-1, 1, (r, c)).astype('f')
    return a,b

if __name__ == "__main__":
    a, b = get_a_b()
    st_t = time.time()
    #dist_numpy(a,b) # comment/ uncomment to execute the code! 
    dist_scipy(a,b) # comment/ uncomment to execute the code!
    print('it took {} s'.format(time.time()-st_t))