python - x 两行点之间的距离

Question

我有两个一维 numpy 数组 A 和 B，大小分别为 (n, ) 和 (m, )，它们对应于直线上点的 x 位置。我想计算 A 中每个点到 B 中每个点之间的距离。然后我需要使用这些距离设置的 y 距离 d，来计算 A 中每个点的电势。

我目前使用的是:

V = numpy.zeros(n)
for i in range(n):
    xdist = A[i] - B
    r = numpy.sqrt(xdist**2 + d**2)
    dV = 1/r
    V[i] = numpy.sum(dV)

这可行，但对于大型数据集可能需要一段时间，所以我想使用类似于 scipy.spatial.distance.cdist 的函数，它不适用于一维数组，我不想添加另一个维度到阵列，因为它们变得太大。

Answer 1

矢量化方法

将 A 扩展到 2D 并使用 np.newaxis/None 引入新轴后的一种矢量化方法从而利用broadcasting会是——

(1/(np.sqrt((A[:,None] - B)**2 + d**2))).sum(1)

大型阵列的混合方法

现在，对于大型数组，我们可能必须将数据分成 block 。

因此，使用 BSZ 作为 block 大小，我们将采用混合方法，就像这样 -

dsq = d**2   
V = np.zeros((n//BSZ,BSZ))
for i in range(n//BSZ):
    V[i] = (1/(np.sqrt((A[i*BSZ:(i+1)*BSZ,None] - B)**2 + dsq))).sum(1)

---

运行时测试

方法-

def original_app(A,B,d):
    V = np.zeros(n)
    for i in range(n):
        xdist = A[i] - B
        r = np.sqrt(xdist**2 + d**2)
        dV = 1/r
        V[i] = np.sum(dV)
    return V

def vectorized_app1(A,B,d):
    return (1/(np.sqrt((A[:,None] - B)**2 + d**2))).sum(1)
    
def vectorized_app2(A,B,d, BSZ = 100):
    dsq = d**2   
    V = np.zeros((n//BSZ,BSZ))
    for i in range(n//BSZ):
        V[i] = (1/(np.sqrt((A[i*BSZ:(i+1)*BSZ,None] - B)**2 + dsq))).sum(1)
    return V.ravel()

时间和验证-

In [203]: # Setup inputs
     ...: n,m = 10000,2000
     ...: A = np.random.rand(n)
     ...: B = np.random.rand(m)
     ...: d = 10
     ...: 

In [204]: out1 = original_app(A,B,d)
     ...: out2 = vectorized_app1(A,B,d)
     ...: out3 = vectorized_app2(A,B,d, BSZ = 100)
     ...: 
     ...: print np.allclose(out1, out2)
     ...: print np.allclose(out1, out3)
     ...: 
True
True

In [205]: %timeit original_app(A,B,d)
10 loops, best of 3: 133 ms per loop

In [206]: %timeit vectorized_app1(A,B,d)
10 loops, best of 3: 138 ms per loop

In [207]: %timeit vectorized_app2(A,B,d, BSZ = 100)
10 loops, best of 3: 65.2 ms per loop

我们可以使用参数 block 大小 BSZ -

In [208]: %timeit vectorized_app2(A,B,d, BSZ = 200)
10 loops, best of 3: 74.5 ms per loop

In [209]: %timeit vectorized_app2(A,B,d, BSZ = 50)
10 loops, best of 3: 67.4 ms per loop

因此，在我这边，最好的方法似乎是在 block 大小为 100 的情况下提供 2x 加速。