加速 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. SkyCoord 對象，它是您天體的坐標

然后您只需使用separation計算它們之間的距離。 這些對象比你使用的更強大，這就是為什么它們更慢。

現在我們知道可以使用以下方法計算角距：

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

請參閱： https : //en.wikipedia.org/wiki/Angular_distance

現在您可以實施它。 只是不要忘記您的角度以degrees radians ，三角函數接受以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

不錯。

Astropy Function 耗時， 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] 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_B = np.cos(A[:, np.newaxis] - B)速度（我們也可以為cos(A[i] - B[j]) cos_A_B = np.cos(A[:, np.newaxis] - B)構建矩陣cos_A_B = np.cos(A[:, np.newaxis] - B) cos(A[i] - B[j])因素，但如果 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

對於SkyCoord類（以及SkyCoord許多其他類），Astropy 初學者經常錯過的SkyCoord是，單個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 數組上的其他函數一樣SkyCoord.separation元素工作：

>>> 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慢的主要my_sum4只是所使用的 Vincenty 公式的復雜性增加，以及它更通用的事實（不假設每個坐標的 RA == DEC）。