加快 Astropy 中的计算速度

Question

我正在尝试使用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)

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

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

Answer 1

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

在你正在创建的代码中

1. 一个单位对象，这将是您的角度。
一个天坐标天体，
2. 它是你的天体的坐标

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

现在我们知道可以使用以下方法计算角度分离：

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

对于相同的数据集，结果为：

星体功能：533.3069727968582

纯数学函数：533.3069727982754

不错。

星体功能需要，2.932075262069702 sec完成

纯数学函数：0.07899618148803711 sec完成

Answer 2

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

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

使用纯数字

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类(以及 Astropy 中的许多其他类)的一点是，单个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在回答中写道：

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

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

但是，我们仍然可以在Astropy中做得更好一点。 如果你看一下SkyCoord.separation的源代码，我们会发现它只不过是一个名为angular_separation的函数的更高级别的接口，该函数使用计算成本稍高的Vincenty公式计算分离，使用坐标球面表示的纬度/隆分量。

对于这样的计算，您可以消除大量开销(如 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正在做的事情，但它是预先计算两个坐标的纬度/纬度数组，并首先将它们转换为弧度，然后在它们上调用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)。