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

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)

如果Hits从5 x 2,000,000运行的那个不明显,并且为了比较,这是从5x20运行的那个test2:

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个核心或集群可用时才这样做.否则,使多处理正常工作所花费的时间可能超过单核运行时间.特别是因为多处理可能无法正常工作,然后您必须重新运行计算.

• 预处理元素:删除仅RA或DEC差异使得无法找到匹配项的项目.然而,如果这不能消除大部分元素,则额外的减法和比较可能实际上减慢了这一点.