时间戳

首页 > 解决方案 > 在 Numpy 中有效地计算欧几里德分布矩阵？

python - 在 Numpy 中有效地计算欧几里德分布矩阵？

问题描述

我有一个二维数据的大数组（约 20k 个条目），我想计算所有条目之间的成对欧几里得距离。我需要输出具有标准的正方形。已经提出了针对这个问题的多种解决方案，但它们似乎都不能有效地用于大型阵列。

对于大型数组，使用复杂转置的方法失败。

Scipy pdist似乎是使用 numpy 最有效的方法。但是，在结果上使用squareform来获得方阵会导致效率非常低。

所以我能想到的最好的方法是使用Scipy cdist，这有点尴尬，因为它确实计算了每对距离两次。提供的时间测量显示了 pdist 在原始距离计算方面的优势。

复杂：49.605 秒

Cdist：4.820 秒

Pdist 1.785 秒

方形 Pdist 10.212 s

标签： pythonnumpynumpy-ndarrayeuclidean-distance

解决方案

既然你已经暗示你不需要完整的结果方阵，注意到 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 上，这种特定的大规模数据分布似乎特别/出乎意料地缓慢，但无论如何，这一点都是值得的。

推荐阅读