Question

我尝试将两个对象（一个数据集包含大约五百万个元素，另一个包含大约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')

Answer 1

您需要问自己的基本问题是：您能否减少数据集？

如果不是，我有一些坏消息：500000 * 2000000是1e12。这意味着你正在尝试进行一万亿次操作。

角度分离涉及一些三角函数（我认为这里涉及cos，sin和sqrt）所以它将大致在几百纳秒到几微秒的数量级操作。假设每个操作需要1us，您仍需要12天才能完成此操作。这假设您没有任何Python循环或IO开销，我认为1us对于这类操作是合理的。

但是有一些方法可以优化它：SkyCoord允许矢量化但只有1D：

# 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 {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，所以我无法在那里提出建议。

但仍存在一个问题，即在普通计算机上进行一万亿次三角运算根本不可行。

如何减少时间的其他一些想法：

使用多处理，以便您的计算机的每个核心并行工作，理论上这可以加快核心数量。在实践中，这是不可达的，我建议只有当你有超过＆gt; = 8个核心或集群可用时才这样做。否则，使多处理正常工作所花费的时间可能超过单核运行时间。特别是因为多处理可能无法正常工作，然后您必须重新运行计算。
预处理元素：删除仅RA或DEC差异使得无法找到匹配项的项目。但是，如果这不能消除很大一部分元素，那么额外的减法和比较实际上可能会减慢这一点。

Answer 2

这是一种在内存中使用缓冲区来减少I / O的实现。注意：我更喜欢使用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())

注意：在Python 2中，我使用xrange而不是range来减少内存使用。

buf.truncate(0)可以替换为这样的新实例：buf = io.BytesIO()。它可能更有效......

Answer 3

加速的第一种方法：c2 =在ra2，dec2 len（RA1）次计算每对的SkyCoord。您可以通过制作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)))

Answer 4

假设您希望将数据集减少到<2度差异（根据您的评论），您可以通过广播制作掩码（可能需要以块为单位，但方法相同）

aMask=(abs(RA1[:,None]-RA2[None,:])<2)&(abs(DEC1[:,None]-DEC2[None,:])<2)

在一些较小规模的测试中，这会将数据集减少大约1/5000。然后创建一个掩码的位置数组。

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]]创建1d数组，按照@ MSeifert的回答，你可以通过SkyCoord推送。

当您获得输出并与1.5进行比较时，您将得到一个布尔bmask，然后您可以outlocs=locs[0][bmask]并输出RA1[outlocs]等。

我做过类似的尝试在壳上进行空间导数的FEM分析，其中在所有数据点之间进行全面的比较同样效率低下。

Answer 5

以这种方式使用的

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)

即只是将格式化的行写入文件

如何为百万元素加速Python嵌套循环

5 个答案: