python - 在 python 中使用重复索引对循环进行矢量化

Question

我正在尝试优化一个被调用很多(数百万次)的片段，因此任何类型的速度改进(希望删除 for 循环)都会很棒。

我正在计算第 j 个粒子与所有其他粒子的相关函数

C_j(|r-r'|) = sqrt(E((s_j(r')-s_k(r))^2)) 在 k 上取平均值。

我的想法是有一个变量 corrfun 将数据分箱到一些分箱(r，在别处定义)。我找到每个 s_k 属于 r 的哪个 bin，并将其存储在 ind 中。所以 ind[0] 是 j=0 点对应的 r 的索引(因此是 corrfun)。多个点可以落入同一个箱子(事实上我希望箱子足够大以包含多个点)所以我将所有的 (s_j(r')-s_k(r))^2 加在一起然后除以指向那个 bin(存储在变量 rw 中)。我最终为此制作的代码如下(np 用于 numpy):

for k, v in enumerate(ind):
        if j==k:
            continue
        corrfun[v] += (s[k]-s[j])**2
        rw[v] += 1
rw2 = rw
rw2[rw < 1] = 1
corrfun = np.sqrt(np.divide(corrfun, rw2))

请注意，rw2 业务是因为我想避免除以 0 的问题，但我确实返回了 rw 数组并且我希望能够区分 rw=0 和 rw=1 元素。也许对此也有更优雅的解决方案。

有没有办法让 for 循环更快？虽然我不想添加 self 交互 (j==k)，但我什至可以接受 self 交互，如果这意味着我可以获得更快的计算速度(ind ~ 1E6 的长度，所以无论如何 self 交互可能都是微不足道的)。

谢谢!

伊利亚

编辑:

这是完整的代码。请注意，在完整代码中，我也在对 j 进行平均。

import numpy as np

def twopointcorr(x,y,s,dr):

    width = np.max(x)-np.min(x)
    height = np.max(y)-np.min(y)

    n = len(x)

    maxR = np.sqrt((width/2)**2 + (height/2)**2)

    r = np.arange(0, maxR, dr)
    print(r)
    corrfun = r*0
    rw = r*0
    print(maxR)
    ''' go through all points'''
    for j in range(0, n-1):
        hypot = np.sqrt((x[j]-x)**2+(y[j]-y)**2)
        ind = [np.abs(r-h).argmin() for h in hypot]

        for k, v in enumerate(ind):
            if j==k:
                continue
            corrfun[v] += (s[k]-s[j])**2
            rw[v] += 1

    rw2 = rw
    rw2[rw < 1] = 1
    corrfun = np.sqrt(np.divide(corrfun, rw2))
    return r, corrfun, rw

我用下面的方式调试测试

from twopointcorr import twopointcorr
import numpy as np
import matplotlib.pyplot as plt
import time

n=1000
x = np.random.rand(n)
y = np.random.rand(n)
s = np.random.rand(n)

print('running two point corr functinon')

start_time = time.time()
r,corrfun,rw = twopointcorr(x,y,s,0.1)
print("--- Execution time is %s seconds ---" % (time.time() - start_time))

fig1=plt.figure()
plt.plot(r, corrfun,'-x')

fig2=plt.figure()
plt.plot(r, rw,'-x')
plt.show()

同样，主要问题是在真实数据集 n~1E6 中。当然，我可以重新采样以使其更小，但我很乐意实际浏览数据集。

Answer 1

下面是使用 broadcast、hypot、round、bincount 去除所有循环的代码:

def twopointcorr2(x, y, s, dr):
    width = np.max(x)-np.min(x)
    height = np.max(y)-np.min(y)
    n = len(x)
    maxR = np.sqrt((width/2)**2 + (height/2)**2)
    r = np.arange(0, maxR, dr)    
    osub = lambda x:np.subtract.outer(x, x)

    ind = np.clip(np.round(np.hypot(osub(x), osub(y)) / dr), 0, len(r)-1).astype(int)
    rw = np.bincount(ind.ravel())
    rw[0] -= len(x)
    corrfun = np.bincount(ind.ravel(), (osub(s)**2).ravel())
    return r, corrfun, rw

为了比较，我修改了你的代码如下:

def twopointcorr(x,y,s,dr):
    width = np.max(x)-np.min(x)
    height = np.max(y)-np.min(y)

    n = len(x)

    maxR = np.sqrt((width/2)**2 + (height/2)**2)

    r = np.arange(0, maxR, dr)
    corrfun = r*0
    rw = r*0
    for j in range(0, n):
        hypot = np.sqrt((x[j]-x)**2+(y[j]-y)**2)
        ind = [np.abs(r-h).argmin() for h in hypot]
        for k, v in enumerate(ind):
            if j==k:
                continue
            corrfun[v] += (s[k]-s[j])**2
            rw[v] += 1

    return r, corrfun, rw        

下面是检查结果的代码:

import numpy as np

n=1000
x = np.random.rand(n)
y = np.random.rand(n)
s = np.random.rand(n)

r1, corrfun1, rw1 = twopointcorr(x,y,s,0.1)
r2, corrfun2, rw2 = twopointcorr2(x,y,s,0.1)

assert np.allclose(r1, r2)
assert np.allclose(corrfun1, corrfun2)
assert np.allclose(rw1, rw2)        

和 %timeit 结果:

%timeit twopointcorr(x,y,s,0.1)
%timeit twopointcorr2(x,y,s,0.1)    

输出:

1 loop, best of 3: 5.16 s per loop
10 loops, best of 3: 134 ms per loop