如何加速 Python 百万元素的嵌套循环
how to speed up Python nested loops for million elements
我尝试将两个满足一定条件的对象（一个数据集包含约50万个元素，另一个数据集包含约200万个元素）配对，然后将两个对象的信息保存到一个文件中。许多变量不参与配对计算，但它们对我接下来的分析很重要，因此我需要跟踪这些变量并保存它们。如果有办法对整个分析进行矢量化，速度会快得多。下面我以随机数为例：
import numpy as np
from astropy import units as u
from astropy.coordinates import SkyCoord
from PyAstronomy import pyasl
RA1 = np.random.uniform(0,360,500000)
DEC1 = np.random.uniform(-90,90,500000)
d = np.random.uniform(55,2000,500000)
z = np.random.uniform(0.05,0.2,500000)
e = np.random.uniform(0.05,1.0,500000)
s = np.random.uniform(0.05,5.0,500000)
RA2 = np.random.uniform(0,360,2000000)
DEC2 = np.random.uniform(-90,90,2000000)
n = np.random.randint(10,10000,2000000)
m = np.random.randint(10,10000,2000000)
f = open('results.txt','a')
for i in range(len(RA1)):
if i % 50000 == 0:
print i
ra1 = RA1[i]
dec1 = DEC1[i]
c1 = SkyCoord(ra=ra1*u.degree, dec=dec1*u.degree)
for j in range(len(RA2)):
ra2 = RA2[j]
dec2 = DEC2[j]
c2 = SkyCoord(ra=ra2*u.degree, dec=dec2*u.degree)
ang = c1.separation(c2)
sep = d[i] * ang.radian
pa = pyasl.positionAngle(ra1, dec1, ra2, dec2)
if sep < 1.5:
np.savetxt(f,np.c_[ra1,dec1,sep,z[i],e[i],s[i],n[j],m[j]], fmt = '%1.4f %1.4f %1.4f %1.4f %1.4f %1.4f %i %i')
这是一个使用内存中的缓冲区来减少 I/O 的实现。注意：我更喜欢对文件 input/output 使用 io 模块，以便与 Python 更加兼容 3. 我认为这是最佳实践。你不会有较低的性能。
import io
with io.open('results.txt', 'a') as f:
buf = io.BytesIO()
for i in xrange(len(RA1)):
if i % 50000 == 0:
print(i)
f.write(buf.getvalue())
buf.truncate(0)
ra1 = RA1[i]
dec1 = DEC1[i]
c1 = SkyCoord(ra=ra1 * u.degree, dec=dec1 * u.degree)
for j in xrange(len(RA2)):
ra2 = RA2[j]
dec2 = DEC2[j]
c2 = SkyCoord(ra=ra2 * u.degree, dec=dec2 * u.degree)
ang = c1.separation(c2)
sep = d[i] * ang.radian
pa = pyasl.positionAngle(ra1, dec1, ra2, dec2)
if sep < 1.5:
np.savetxt(buf, np.c_[ra1, dec1, sep, z[i], e[i], s[i], n[j], m[j]],
fmt='%1.4f %1.4f %1.4f %1.4f %1.4f %1.4f %i %i')
f.write(buf.getvalue())
注意：在Python2中，我使用xrange代替range来减少内存使用。
buf.truncate(0) 可以用这样的新实例替换：buf = io.BytesIO()。它可以更有效率......
savetxt这样使用本质上是
astr = fmt % (ra1,dec1,sep,z[i],e[i],s[i],n[j],m[j])
astr += '\n' # or include in fmt
f.write(astr)
也就是说，只是将格式化的行写入文件
第一种加速方法：c2 = 为 ra2 中的每一对计算 SkyCoord，dec2 len(RA1) 次。您可以通过制作 SkyCoord 的缓冲区数组来加速：
f = open('results.txt','a')
C1 = [SkyCoord(ra=ra1*u.degree, dec=DEC1[i]*u.degree)
for i, ra1 in enumerate(RA1)] )
C2 = [SkyCoord(ra=ra2*u.degree, dec=DEC2[i]*u.degree)
for i, ra2 in enumerate(RA2)] ) # buffer coords
for i, c1 in enumerate(C1): # we only need enumerate() to get i
for j, c2 in enumerate(C2):
ang = c1.separation(c2) # note we don't have to calculate c2
if d[i] < 1.5 / ang.radian:
# now we don't have to multiply every iteration.
# The right part is a constant
# the next line is only executed if objects are close enough
pa = pyasl.positionAngle(RA1[i], DEC1[i], RA2[j], DEC2[j])
np.savetxt('...whatever')
您可以通过阅读 SkyCoord.separation 代码并将其矢量化以替换 SkyCoord 来进一步加快速度，但我懒得自己做。我假设如果我们有两个 2xN 坐标矩阵 x1, x2 它看起来类似于 (Matlab/Octave):
cos = pdist2(x1, x2) / (sqrt(dot(x1, x1)) * sqrt(dot(x2, x2)))
你需要问自己的基本问题是：你能减少数据集吗？
如果不是，我有一些坏消息：500000 * 2000000 是 1e12。这意味着您要尝试执行一万亿次操作。
angular分离涉及到一些三角函数（我认为这里涉及到cos、sin和sqrt）所以大概是几百个每个操作的纳秒到微秒。假设每个操作需要 1us，您仍然需要 12 天才能完成。这假设您没有任何 Python 循环或 IO 开销，我认为 1us 对于这类操作是合理的。
但肯定有优化它的方法：SkyCoord 允许矢量化，但只允许一维：
# Create the SkyCoord for the longer array once
c2 = SkyCoord(ra=RA2*u.degree, dec=DEC2*u.degree)
# and calculate the seperation from each coordinate of the shorter list
for idx, (ra, dec) in enumerate(zip(RA1, DEC1)):
c1 = SkyCoord(ra=ra*u.degree, dec=dec*u.degree)
# x will be the angular seperation with a length of your RA2 and DEC2 arrays
x = c1.separation(c2)
这将产生几个数量级的加速：
# note that I made these MUCH shorter
RA1 = np.random.uniform(0,360,5)
DEC1 = np.random.uniform(-90,90,5)
RA2 = np.random.uniform(0,360,10)
DEC2 = np.random.uniform(-90,90,10)
def test(RA1, DEC1, RA2, DEC2):
"""Version with vectorized inner loop."""
c2 = SkyCoord(ra=RA2*u.degree, dec=DEC2*u.degree)
for idx, (ra, dec) in enumerate(zip(RA1, DEC1)):
c1 = SkyCoord(ra=ra*u.degree, dec=dec*u.degree)
x = c1.separation(c2)
def test2(RA1, DEC1, RA2, DEC2):
"""Double loop."""
for ra, dec in zip(RA1, DEC1):
c1 = SkyCoord(ra=ra*u.degree, dec=dec*u.degree)
for ra, dec in zip(RA2, DEC2):
c2 = SkyCoord(ra=ra*u.degree, dec=dec*u.degree)
x = c1.separation(c2)
%timeit test(RA1, DEC1, RA2, DEC2) # 1 loop, best of 3: 225 ms per loop
%timeit test2(RA1, DEC1, RA2, DEC2) # 1 loop, best of 3: 2.71 s per loop
这已经快了 10 倍，而且扩展性更好：
RA1 = np.random.uniform(0,360,5)
DEC1 = np.random.uniform(-90,90,5)
RA2 = np.random.uniform(0,360,2000000)
DEC2 = np.random.uniform(-90,90,2000000)
%timeit test(RA1, DEC1, RA2, DEC2) # 1 loop, best of 3: 2.8 s per loop
# test2 scales so bad I only use 50 elements here
RA2 = np.random.uniform(0,360,50)
DEC2 = np.random.uniform(-90,90,50)
%timeit test2(RA1, DEC1, RA2, DEC2) # 1 loop, best of 3: 11.4 s per loop
请注意，通过矢量化内部循环，我能够在 1/4 的时间内计算出 40000 倍以上的元素。因此，通过矢量化内部循环，您应该快大约 200k 倍。
此处我们在 3 秒内计算了 5 次 200 万次分离，因此每次操作大约需要 300 ns。以这种速度，您需要 3 天才能完成此任务。
即使您也可以将剩余的循环向量化，我认为这不会产生任何大的加速，因为在那个级别，循环开销比每个循环中的计算时间少得多。使用 line-profiler 支持此语句：
Line # Hits Time Per Hit % Time Line Contents
==============================================================
11 def test(RA1, DEC1, RA2, DEC2):
12 1 216723 216723.0 2.6 c2 = SkyCoord(ra=RA2*u.degree, dec=DEC2*u.degree)
13 6 222 37.0 0.0 for idx, (ra, dec) in enumerate(zip(RA1, DEC1)):
14 5 206796 41359.2 2.5 c1 = SkyCoord(ra=ra*u.degree, dec=dec*u.degree)
15 5 7847321 1569464.2 94.9 x = c1.separation(c2)
如果从 5 x 2,000,000 运行 的 Hits 看不明显，这里是 test2 上 5x20 运行 的比较：
Line # Hits Time Per Hit % Time Line Contents
==============================================================
17 def test2(RA1, DEC1, RA2, DEC2):
18 6 80 13.3 0.0 for ra, dec in zip(RA1, DEC1):
19 5 195030 39006.0 0.6 c1 = SkyCoord(ra=ra*u.degree, dec=dec*u.degree)
20 105 1737 16.5 0.0 for ra, dec in zip(RA2, DEC2):
21 100 3871427 38714.3 11.8 c2 = SkyCoord(ra=ra*u.degree, dec=dec*u.degree)
22 100 28870724 288707.2 87.6 x = c1.separation(c2)
test2 规模变差的原因是 c2 = SkyCoord 部分占用了总时间的 12%，而不是仅仅 2.5%，并且每次调用 seperation 都有一些重要的高架。因此，真正使它变慢的并不是 Python 循环开销，而是 SkyCoord 构造函数和 seperation.
的静态部分
您显然需要向量化 pa 计算并保存到文件（我没有使用过 PyAstronomy 和 numpy.savetext，所以我不能在那里提供建议）。
但还有一个问题就是，在普通计算机上进行一万亿次三角运算根本行不通
关于如何减少时间的一些额外想法：
使用多处理，让您的计算机的每个核心并行工作，理论上这可以通过您的核心数量来加快速度。在实践中，这将无法实现，我建议仅当您有超过 >= 8 个核心或可用集群时才这样做。否则，使多处理正常工作所花费的时间可能会超过单核 运行ning 时间。特别是因为 multiprocessing 可能无法正常工作，然后你必须重新运行计算。
预处理元素：删除仅靠 RA 或 DEC 差异就无法找到匹配项的项目。但是，如果这不能删除很大一部分元素，则额外的减法和比较实际上可能会减慢速度。
假设你想将你的数据集减少到 <2 度差异（根据你的评论），你可以通过广播制作一个面具（可能需要分块做，但方法是一样的）
aMask=(abs(RA1[:,None]-RA2[None,:])<2)&(abs(DEC1[:,None]-DEC2[None,:])<2)
在一些较小规模的测试中，这将数据集减少了大约 1/5000。然后做一个mask的location数组。
locs=np.where(aMask)
(array([ 0, 2, 4, ..., 4998, 4999, 4999], dtype=int32),
array([3575, 1523, 1698, ..., 4869, 1801, 2792], dtype=int32))
（来自我的 5k x 5k 测试）。通过例如 d[locs[0]] 转储所有其他变量以创建一维数组，您可以根据@MSeifert 的回答将其推送到 SkyCoord。
当您获得输出并与 1.5 进行比较时，您将得到一个布尔值 bmask，然后您可以 outlocs=locs[0][bmask] 并输出 RA1[outlocs] 等
我做过类似的事情，尝试对壳进行空间导数以进行 FEM 分析，其中对所有数据点进行全秩比较同样效率低下。
我尝试将两个满足一定条件的对象（一个数据集包含约50万个元素，另一个数据集包含约200万个元素）配对，然后将两个对象的信息保存到一个文件中。许多变量不参与配对计算，但它们对我接下来的分析很重要，因此我需要跟踪这些变量并保存它们。如果有办法对整个分析进行矢量化，速度会快得多。下面我以随机数为例：
import numpy as np
from astropy import units as u
from astropy.coordinates import SkyCoord
from PyAstronomy import pyasl
RA1 = np.random.uniform(0,360,500000)
DEC1 = np.random.uniform(-90,90,500000)
d = np.random.uniform(55,2000,500000)
z = np.random.uniform(0.05,0.2,500000)
e = np.random.uniform(0.05,1.0,500000)
s = np.random.uniform(0.05,5.0,500000)
RA2 = np.random.uniform(0,360,2000000)
DEC2 = np.random.uniform(-90,90,2000000)
n = np.random.randint(10,10000,2000000)
m = np.random.randint(10,10000,2000000)
f = open('results.txt','a')
for i in range(len(RA1)):
if i % 50000 == 0:
print i
ra1 = RA1[i]
dec1 = DEC1[i]
c1 = SkyCoord(ra=ra1*u.degree, dec=dec1*u.degree)
for j in range(len(RA2)):
ra2 = RA2[j]
dec2 = DEC2[j]
c2 = SkyCoord(ra=ra2*u.degree, dec=dec2*u.degree)
ang = c1.separation(c2)
sep = d[i] * ang.radian
pa = pyasl.positionAngle(ra1, dec1, ra2, dec2)
if sep < 1.5:
np.savetxt(f,np.c_[ra1,dec1,sep,z[i],e[i],s[i],n[j],m[j]], fmt = '%1.4f %1.4f %1.4f %1.4f %1.4f %1.4f %i %i')
这是一个使用内存中的缓冲区来减少 I/O 的实现。注意：我更喜欢对文件 input/output 使用 io 模块，以便与 Python 更加兼容 3. 我认为这是最佳实践。你不会有较低的性能。
import io
with io.open('results.txt', 'a') as f:
buf = io.BytesIO()
for i in xrange(len(RA1)):
if i % 50000 == 0:
print(i)
f.write(buf.getvalue())
buf.truncate(0)
ra1 = RA1[i]
dec1 = DEC1[i]
c1 = SkyCoord(ra=ra1 * u.degree, dec=dec1 * u.degree)
for j in xrange(len(RA2)):
ra2 = RA2[j]
dec2 = DEC2[j]
c2 = SkyCoord(ra=ra2 * u.degree, dec=dec2 * u.degree)
ang = c1.separation(c2)
sep = d[i] * ang.radian
pa = pyasl.positionAngle(ra1, dec1, ra2, dec2)
if sep < 1.5:
np.savetxt(buf, np.c_[ra1, dec1, sep, z[i], e[i], s[i], n[j], m[j]],
fmt='%1.4f %1.4f %1.4f %1.4f %1.4f %1.4f %i %i')
f.write(buf.getvalue())
注意：在Python2中，我使用xrange代替range来减少内存使用。
buf.truncate(0) 可以用这样的新实例替换：buf = io.BytesIO()。它可以更有效率......
savetxt这样使用本质上是
astr = fmt % (ra1,dec1,sep,z[i],e[i],s[i],n[j],m[j])
astr += '\n' # or include in fmt
f.write(astr)
也就是说，只是将格式化的行写入文件
第一种加速方法：c2 = 为 ra2 中的每一对计算 SkyCoord，dec2 len(RA1) 次。您可以通过制作 SkyCoord 的缓冲区数组来加速：
f = open('results.txt','a')
C1 = [SkyCoord(ra=ra1*u.degree, dec=DEC1[i]*u.degree)
for i, ra1 in enumerate(RA1)] )
C2 = [SkyCoord(ra=ra2*u.degree, dec=DEC2[i]*u.degree)
for i, ra2 in enumerate(RA2)] ) # buffer coords
for i, c1 in enumerate(C1): # we only need enumerate() to get i
for j, c2 in enumerate(C2):
ang = c1.separation(c2) # note we don't have to calculate c2
if d[i] < 1.5 / ang.radian:
# now we don't have to multiply every iteration.
# The right part is a constant
# the next line is only executed if objects are close enough
pa = pyasl.positionAngle(RA1[i], DEC1[i], RA2[j], DEC2[j])
np.savetxt('...whatever')
您可以通过阅读 SkyCoord.separation 代码并将其矢量化以替换 SkyCoord 来进一步加快速度，但我懒得自己做。我假设如果我们有两个 2xN 坐标矩阵 x1, x2 它看起来类似于 (Matlab/Octave):
cos = pdist2(x1, x2) / (sqrt(dot(x1, x1)) * sqrt(dot(x2, x2)))
你需要问自己的基本问题是：你能减少数据集吗？
如果不是，我有一些坏消息：500000 * 2000000 是 1e12。这意味着您要尝试执行一万亿次操作。
angular分离涉及到一些三角函数（我认为这里涉及到cos、sin和sqrt）所以大概是几百个每个操作的纳秒到微秒。假设每个操作需要 1us，您仍然需要 12 天才能完成。这假设您没有任何 Python 循环或 IO 开销，我认为 1us 对于这类操作是合理的。
但肯定有优化它的方法：SkyCoord 允许矢量化，但只允许一维：
# Create the SkyCoord for the longer array once
c2 = SkyCoord(ra=RA2*u.degree, dec=DEC2*u.degree)
# and calculate the seperation from each coordinate of the shorter list
for idx, (ra, dec) in enumerate(zip(RA1, DEC1)):
c1 = SkyCoord(ra=ra*u.degree, dec=dec*u.degree)
# x will be the angular seperation with a length of your RA2 and DEC2 arrays
x = c1.separation(c2)
这将产生几个数量级的加速：
# note that I made these MUCH shorter
RA1 = np.random.uniform(0,360,5)
DEC1 = np.random.uniform(-90,90,5)
RA2 = np.random.uniform(0,360,10)
DEC2 = np.random.uniform(-90,90,10)
def test(RA1, DEC1, RA2, DEC2):
"""Version with vectorized inner loop."""
c2 = SkyCoord(ra=RA2*u.degree, dec=DEC2*u.degree)
for idx, (ra, dec) in enumerate(zip(RA1, DEC1)):
c1 = SkyCoord(ra=ra*u.degree, dec=dec*u.degree)
x = c1.separation(c2)
def test2(RA1, DEC1, RA2, DEC2):
"""Double loop."""
for ra, dec in zip(RA1, DEC1):
c1 = SkyCoord(ra=ra*u.degree, dec=dec*u.degree)
for ra, dec in zip(RA2, DEC2):
c2 = SkyCoord(ra=ra*u.degree, dec=dec*u.degree)
x = c1.separation(c2)
%timeit test(RA1, DEC1, RA2, DEC2) # 1 loop, best of 3: 225 ms per loop
%timeit test2(RA1, DEC1, RA2, DEC2) # 1 loop, best of 3: 2.71 s per loop
这已经快了 10 倍，而且扩展性更好：
RA1 = np.random.uniform(0,360,5)
DEC1 = np.random.uniform(-90,90,5)
RA2 = np.random.uniform(0,360,2000000)
DEC2 = np.random.uniform(-90,90,2000000)
%timeit test(RA1, DEC1, RA2, DEC2) # 1 loop, best of 3: 2.8 s per loop
# test2 scales so bad I only use 50 elements here
RA2 = np.random.uniform(0,360,50)
DEC2 = np.random.uniform(-90,90,50)
%timeit test2(RA1, DEC1, RA2, DEC2) # 1 loop, best of 3: 11.4 s per loop
请注意，通过矢量化内部循环，我能够在 1/4 的时间内计算出 40000 倍以上的元素。因此，通过矢量化内部循环，您应该快大约 200k 倍。
此处我们在 3 秒内计算了 5 次 200 万次分离，因此每次操作大约需要 300 ns。以这种速度，您需要 3 天才能完成此任务。
即使您也可以将剩余的循环向量化，我认为这不会产生任何大的加速，因为在那个级别，循环开销比每个循环中的计算时间少得多。使用 line-profiler 支持此语句：
Line # Hits Time Per Hit % Time Line Contents
==============================================================
11 def test(RA1, DEC1, RA2, DEC2):
12 1 216723 216723.0 2.6 c2 = SkyCoord(ra=RA2*u.degree, dec=DEC2*u.degree)
13 6 222 37.0 0.0 for idx, (ra, dec) in enumerate(zip(RA1, DEC1)):
14 5 206796 41359.2 2.5 c1 = SkyCoord(ra=ra*u.degree, dec=dec*u.degree)
15 5 7847321 1569464.2 94.9 x = c1.separation(c2)
如果从 5 x 2,000,000 运行 的 Hits 看不明显，这里是 test2 上 5x20 运行 的比较：
Line # Hits Time Per Hit % Time Line Contents
==============================================================
17 def test2(RA1, DEC1, RA2, DEC2):
18 6 80 13.3 0.0 for ra, dec in zip(RA1, DEC1):
19 5 195030 39006.0 0.6 c1 = SkyCoord(ra=ra*u.degree, dec=dec*u.degree)
20 105 1737 16.5 0.0 for ra, dec in zip(RA2, DEC2):
21 100 3871427 38714.3 11.8 c2 = SkyCoord(ra=ra*u.degree, dec=dec*u.degree)
22 100 28870724 288707.2 87.6 x = c1.separation(c2)
test2 规模变差的原因是 c2 = SkyCoord 部分占用了总时间的 12%，而不是仅仅 2.5%，并且每次调用 seperation 都有一些重要的高架。因此，真正使它变慢的并不是 Python 循环开销，而是 SkyCoord 构造函数和 seperation.
您显然需要向量化 pa 计算并保存到文件（我没有使用过 PyAstronomy 和 numpy.savetext，所以我不能在那里提供建议）。
但还有一个问题就是，在普通计算机上进行一万亿次三角运算根本行不通
关于如何减少时间的一些额外想法：
使用多处理，让您的计算机的每个核心并行工作，理论上这可以通过您的核心数量来加快速度。在实践中，这将无法实现，我建议仅当您有超过 >= 8 个核心或可用集群时才这样做。否则，使多处理正常工作所花费的时间可能会超过单核 运行ning 时间。特别是因为 multiprocessing 可能无法正常工作，然后你必须重新运行计算。
预处理元素：删除仅靠 RA 或 DEC 差异就无法找到匹配项的项目。但是，如果这不能删除很大一部分元素，则额外的减法和比较实际上可能会减慢速度。
假设你想将你的数据集减少到 <2 度差异（根据你的评论），你可以通过广播制作一个面具（可能需要分块做，但方法是一样的）
aMask=(abs(RA1[:,None]-RA2[None,:])<2)&(abs(DEC1[:,None]-DEC2[None,:])<2)
在一些较小规模的测试中，这将数据集减少了大约 1/5000。然后做一个mask的location数组。
locs=np.where(aMask)
(array([ 0, 2, 4, ..., 4998, 4999, 4999], dtype=int32),
array([3575, 1523, 1698, ..., 4869, 1801, 2792], dtype=int32))
（来自我的 5k x 5k 测试）。通过例如 d[locs[0]] 转储所有其他变量以创建一维数组，您可以根据@MSeifert 的回答将其推送到 SkyCoord。
当您获得输出并与 1.5 进行比较时，您将得到一个布尔值 bmask，然后您可以 outlocs=locs[0][bmask] 并输出 RA1[outlocs] 等
我做过类似的事情，尝试对壳进行空间导数以进行 FEM 分析，其中对所有数据点进行全秩比较同样效率低下。