用 python 中的重复索引向量化 for 循环
Vectorizing for loop with repeated indices in python
我正在尝试优化一个被调用很多（数百万次）的片段，因此任何类型的速度改进（希望删除 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 的问题，但我 return rw 数组并且我希望能够区分 rw=0 和 rw=1 元素。也许对此也有更优雅的解决方案。
有没有办法让 for 循环更快？虽然我不想添加自我交互 (j==k) 我什至可以接受自我交互，如果这意味着我可以显着加快计算速度（ind ~ 1E6 的长度，所以自我交互可能无关紧要）。
谢谢！
伊利亚
编辑：
这是完整的代码。请注意，在完整代码中，我也在对 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中。当然，我可以重新采样以使其更小，但我很乐意实际浏览数据集。
你的原始代码在我的系统上 运行s 大约需要 5.7 秒。我完全矢量化了内部循环并在 0.39 秒内将其设置为 运行。只需将 "go through all points" 循环替换为：
points = np.column_stack((x,y))
hypots = scipy.spatial.distance.cdist(points, points)
inds = np.rint(hypots.clip(max=maxR) / dr).astype(np.int)
# go through all points
for j in range(n): # n.b. previously n-1, not sure why
ind = inds[j]
np.add.at(corrfun, ind, (s - s[j])**2)
np.add.at(rw, ind, 1)
rw[ind[j]] -= 1 # subtract self
第一个观察结果是您的 hypot 代码正在计算 2D 距离，因此我将其替换为 SciPy 中的 cdist 以在一次调用中完成所有操作。第二个是内部 for 循环很慢，多亏了@hpaulj 的有见地的评论，我也使用 np.add.at().
对其进行了矢量化
既然你也问了如何向量化内循环，我后来就做了。现在 运行 需要 0.25 秒，总加速超过 20 倍。这是最终代码：
points = np.column_stack((x,y))
hypots = scipy.spatial.distance.cdist(points, points)
inds = np.rint(hypots.clip(max=maxR) / dr).astype(np.int)
sn = np.tile(s, (n,1)) # n copies of s
diffs = (sn - sn.T)**2 # squares of pairwise differences
np.add.at(corrfun, inds, diffs)
rw = np.bincount(inds.flatten(), minlength=len(r))
np.subtract.at(rw, inds.diagonal(), 1) # subtract self
这使用了更多内存，但与上面的单循环版本相比确实产生了显着的加速。
下面是使用 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
好的，事实证明外部产品的内存消耗非常大，但是，使用@HYRY 和@JohnZwinck 的答案，我能够编写出在内存中仍然大致呈线性的代码，并且计算速度很快（0.5 秒）对于测试用例）
import numpy as np
def twopointcorr(x,y,s,dr,maxR=-1):
width = np.max(x)-np.min(x)
height = np.max(y)-np.min(y)
n = len(x)
if maxR < dr:
maxR = np.sqrt((width/2)**2 + (height/2)**2)
r = np.arange(0, maxR+dr, dr)
corrfun = r*0
rw = r*0
for j in range(0, n):
ind = np.clip(np.round(np.hypot(x[j]-x,y[j]-y) / dr), 0, len(r)-1).astype(int)
np.add.at(corrfun, ind, (s - s[j])**2)
np.add.at(rw, ind, 1)
rw[0] -= n
corrfun = np.sqrt(np.divide(corrfun, np.maximum(rw,1)))
r=np.delete(r,-1)
rw=np.delete(rw,-1)
corrfun=np.delete(corrfun,-1)
return r, corrfun, rw
我正在尝试优化一个被调用很多（数百万次）的片段，因此任何类型的速度改进（希望删除 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 的问题，但我 return rw 数组并且我希望能够区分 rw=0 和 rw=1 元素。也许对此也有更优雅的解决方案。
有没有办法让 for 循环更快？虽然我不想添加自我交互 (j==k) 我什至可以接受自我交互，如果这意味着我可以显着加快计算速度（ind ~ 1E6 的长度，所以自我交互可能无关紧要）。
谢谢！
伊利亚
编辑：
这是完整的代码。请注意，在完整代码中，我也在对 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中。当然，我可以重新采样以使其更小，但我很乐意实际浏览数据集。
你的原始代码在我的系统上 运行s 大约需要 5.7 秒。我完全矢量化了内部循环并在 0.39 秒内将其设置为 运行。只需将 "go through all points" 循环替换为：
points = np.column_stack((x,y))
hypots = scipy.spatial.distance.cdist(points, points)
inds = np.rint(hypots.clip(max=maxR) / dr).astype(np.int)
# go through all points
for j in range(n): # n.b. previously n-1, not sure why
ind = inds[j]
np.add.at(corrfun, ind, (s - s[j])**2)
np.add.at(rw, ind, 1)
rw[ind[j]] -= 1 # subtract self
第一个观察结果是您的 hypot 代码正在计算 2D 距离，因此我将其替换为 SciPy 中的 cdist 以在一次调用中完成所有操作。第二个是内部 for 循环很慢，多亏了@hpaulj 的有见地的评论，我也使用 np.add.at().
既然你也问了如何向量化内循环，我后来就做了。现在 运行 需要 0.25 秒，总加速超过 20 倍。这是最终代码：
points = np.column_stack((x,y))
hypots = scipy.spatial.distance.cdist(points, points)
inds = np.rint(hypots.clip(max=maxR) / dr).astype(np.int)
sn = np.tile(s, (n,1)) # n copies of s
diffs = (sn - sn.T)**2 # squares of pairwise differences
np.add.at(corrfun, inds, diffs)
rw = np.bincount(inds.flatten(), minlength=len(r))
np.subtract.at(rw, inds.diagonal(), 1) # subtract self
这使用了更多内存，但与上面的单循环版本相比确实产生了显着的加速。
下面是使用 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
好的，事实证明外部产品的内存消耗非常大，但是，使用@HYRY 和@JohnZwinck 的答案，我能够编写出在内存中仍然大致呈线性的代码，并且计算速度很快（0.5 秒）对于测试用例）
import numpy as np
def twopointcorr(x,y,s,dr,maxR=-1):
width = np.max(x)-np.min(x)
height = np.max(y)-np.min(y)
n = len(x)
if maxR < dr:
maxR = np.sqrt((width/2)**2 + (height/2)**2)
r = np.arange(0, maxR+dr, dr)
corrfun = r*0
rw = r*0
for j in range(0, n):
ind = np.clip(np.round(np.hypot(x[j]-x,y[j]-y) / dr), 0, len(r)-1).astype(int)
np.add.at(corrfun, ind, (s - s[j])**2)
np.add.at(rw, ind, 1)
rw[0] -= n
corrfun = np.sqrt(np.divide(corrfun, np.maximum(rw,1)))
r=np.delete(r,-1)
rw=np.delete(rw,-1)
corrfun=np.delete(corrfun,-1)
return r, corrfun, rw