加速 Astropy 中的计算

Speed up calculation in Astropy

我正在尝试使用 astropy 计算点列表之间的距离总和。但是我的实现速度太慢,无法用我的数据实现,这是我的代码示例之一:

import pandas as pd
import numpy as np  

# synthetic data
lst2 = list(range(50))
lst2 = np.array(lst2)/50
lst3 = np.array(lst2)/51

df = pd.DataFrame(list(zip(lst2, lst3)),
               columns =['A', 'B'])

# Sum of the distance between different points
def Sum(df_frame):
    length = len(df_frame) #Size of "for" loops
    Sum = 0 
    for i in range(length - 1): 
        for j in range(i+1,length):
            c1 = SkyCoord(df_frame['A'].iloc[i]*u.deg, df_frame['A'].iloc[i]*u.deg, frame='icrs')
            c2 = SkyCoord(df_frame['B'].iloc[j]*u.deg, df_frame['B'].iloc[j]*u.deg, frame='icrs') 
            angle = c1.separation(c2).deg
            Sum += angle
    return  Sum

Sum(df)

有谁知道如何提高这段代码的计算速度?

我的真实数据有几百万个点

您应该知道,有时使用现成的产品会更快,因为所有工具都可用。但是在某些情况下,就像您的情况一样,使用现成的产品会使您的执行时间变慢。

在您正在创建的代码中

  1. 一个单位对象,就是你的角度。
  2. 一个 SkyCoord 对象,它是您天体的坐标

然后您只需使用 separation 计算它们之间的距离即可。这些对象比您使用的对象更强大,这就是它们速度较慢的原因。

现在我们知道可以使用以下方法计算 angular 分离:

arccos(sin(delta1) * sin(delta2) + cos(delta1) * cos(delta2) * sin(alpha1 - alpha2))

参见:https://en.wikipedia.org/wiki/Angular_distance

现在你可以实现了。只是不要忘记您的角度在 degrees 中,三角函数接受 radians

中的角度
def my_sum(df_frame):
    length = len(df_frame)  # Size of "for" loops
    Sum = 0
    df_frame_rad = np.deg2rad(df_frame)
    for i in range(length - 1):
        for j in range(i + 1, length):
            # print(a2, d2)
            dist = np.rad2deg(
                np.arccos(
                    np.sin(df_frame_rad['A'].iloc[i]) * np.sin(df_frame_rad['B'].iloc[j]) + \
                    np.cos(df_frame_rad['A'].iloc[i]) * np.cos(df_frame_rad['B'].iloc[j]) * \
                    np.cos(df_frame_rad['A'].iloc[i] - df_frame_rad['B'].iloc[j])
                )
            )
            Sum += dist
    return Sum

对于相同的数据集,结果为:

Astropy 函数:533.3069727968582

纯数学函数:533.3069727982754

不错。

Astropy 函数已完成,2.932075262069702 sec 完成

纯数学函数用了:0.07899618148803711 sec 完成

仍然会非常慢,尤其是在大型数据帧上,因为对于每个 O(n^2) 对元素,您都有一个双循环索引数据帧,如 df['A'].loc[i]

我尝试使用每列仅包含 1000 个元素的数据框进行此操作,这花了很长时间。对于更大的数字,我只是放弃了等待。如果您改为将列作为普通 numpy 数组传递给函数,然后在执行距离计算之前还分配 A_i = A[i]; B_j = B[j],它会大大加快速度,即像:

使用纯 NumPy

def my_sum2(A, B):
    length = len(A)  # Size of "for" loops
    assert length == len(B)
    Sum = 0
    A = np.deg2rad(np.asarray(A))
    B = np.deg2rad(np.asarray(B))
    for i in range(length - 1):
        for j in range(i + 1, length):
            # print(a2, d2)
            A_i = A[i]
            B_j = B[j]
            dist = np.rad2deg(
                np.arccos(
                    np.sin(A_i) * np.sin(B_j) + \
                    np.cos(A_i) * np.cos(B_j) * \
                    np.cos(A_i - B_j)
                )
            )
            Sum += dist
    return Sum

对于 100 个元素,我得到:

>>> %timeit my_sum(df)
229 ms ± 3.06 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
>>> %timeit my_sum2(df['A'], df['B'])
41.1 ms ± 2.88 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

但是您可以通过使用矢量化运算预先计算正弦和余弦来做得更好。这会导致更大的内存使用量,以牺牲速度为代价(我们也可以为 cos(A[i] - B[j]) 因子构建矩阵 cos_A_B = np.cos(A[:, np.newaxis] - B),但如果 A 和 B 很大,这将非常耗尽内存) :

def my_sum3(A, B):
    length = len(A)  # Size of "for" loops
    assert length == len(B)
    Sum = 0
    A = np.deg2rad(np.asarray(A))
    B = np.deg2rad(np.asarray(B))
    cos_A = np.cos(A)
    sin_A = np.sin(A)
    cos_B = np.cos(B)
    sin_B = np.sin(B)

    for i in range(length - 1):
        for j in range(i + 1, length):
            # print(a2, d2)
            dist = np.rad2deg(
                np.arccos(
                    sin_A[i] * sin_B[j] + \
                    cos_A[i] * cos_B[j] * \
                    np.cos(A[i] - B[j])
                )
            )
            Sum += dist
    return Sum
>>> %timeit my_sum3(df['A'], df['B'])
20.2 ms ± 715 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

但是对于 NumPy 数组的成对计算,我们可以进一步利用 NumPy 的逐元素广播以完全消除内部 for 循环:

def my_sum4(A, B):
    length = len(A)  # Size of "for" loops
    assert length == len(B)
    Sum = 0
    A = np.deg2rad(np.asarray(A))
    B = np.deg2rad(np.asarray(B))
    cos_A = np.cos(A)
    sin_A = np.sin(A)
    cos_B = np.cos(B)
    sin_B = np.sin(B)
    
    for i in range(length - 1):
        Sum += np.sum(np.rad2deg(np.arccos(
            sin_A[i] * sin_B[i + 1:] +
            cos_A[i] * cos_B[i + 1:] *
            np.cos(A[i] - B[i + 1:]))))

    return Sum
>>> %timeit my_sum4(df['A'], df['B'])
1.31 ms ± 71.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

还有许多其他方法可以对此进行微优化,使用 Cython、scipy 等,但我不会在这里花更多时间。

这种方法的另一个问题是它专门针对 OP 问题的细节,其中每个坐标由于某种原因具有相同的 RA 和 DEC,并且没有被推广。

使用 SkyCoord

Astropy 初学者经常忽略 SkyCoord class(以及 Astropy 中的许多其他 classes)的一点是,单个 SkyCoord 可以是一个容器坐标数组,而不只是一个坐标。

在 OP 的问题中,他们正在创建数百万个 SkyCoord 对象,每个坐标一个。事实上你可以简单地这样做:

>>> c1 = SkyCoord(df['A']*u.deg, df['A']*u.deg, frame='icrs')
>>> c2 = SkyCoord(df['B']*u.deg, df['B']*u.deg, frame='icrs')

SkyCoord.separation 这样的方法也像 NumPy 数组上的其他函数一样按元素工作:

>>> c1.separation(c2)
<Angle [0.0130013 , 1.18683992, 0.82050812, ...] deg>

因此,对于每个成对分离,您可以使用与我的 my_sum4 解决方案中类似的技术,让您不必自己编写计算:

def my_sum5(c1, c2):
    angle_sum = 0
    for idx in range(len(c1)):
        angle_sum += c1[idx].separation(c2[idx + 1:]).sum()
    return angle_sum
>>> my_sum5(c1, c2)
<Angle 2368.14558945 deg>

诚然,这比上一个纯 NumPy 解决方案

>>> %timeit my_sum5(c1, c2)
166 ms ± 10.2 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

这个开销是 Astropy 的一些高级接口的成本,我同意 MSH 在他们的回答中写道:

You should know some times using ready products is faster since all the tools are available. However in some condition, as yours, using ready product makes you slower in execution time.

也就是说,如果您确实对大型数据集有高性能需求,那么使用手动优化的解决方案可能会更好。

但是,我们仍然可以在 Astropy 中做得更好。如果您查看 the source code for SkyCoord.separation we see that it's little more than a higher-level interface to a function called angular_separation,它使用计算成本更高的 Vincenty 公式计算分离,使用坐标球形表示的 lat/lon 分量。

对于这样的计算,您可以在直接使用此函数时消除大量开销(例如 Astropy 的自动坐标转换),例如:

def my_sum6(c1, c2):
    angle_sum = 0
    lon1 = c1.spherical.lon.to(u.rad).value
    lat1 = c1.spherical.lat.to(u.rad).value
    lon2 = c2.spherical.lon.to(u.rad).value
    lat2 = c2.spherical.lat.to(u.rad).value
    
    for idx in range(len(c1)):
        angle_sum += angular_separation(lon1[idx], lat1[idx], lon2[idx+1:], lat2[idx+1:]).sum()
    return np.rad2deg(angle_sum)

这基本上是在做 SkyCoord.separation 正在做的事情,但它是预先计算两个坐标的 lat/lon 数组并首先将它们转换为弧度,然后对它们调用 angular_separation。它还跳过了断言两个坐标都在同一帧中的开销(在这种情况下它们都是 ICRS,所以我们假设它们是)。这几乎和 my_sum4:

一样好
>>> %timeit my_sum6(c1, c2)
2.26 ms ± 123 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

事实上,在这种情况下,使它比 my_sum4 慢的主要原因只是所使用的 Vincenty 公式的复杂性增加,以及它更通用的事实(不假设 RA == DEC for每个坐标)。