如何提升Astropy中大量点对天球距离求和的计算速度

核心性能瓶颈 #

你的现有代码存在两个致命的性能问题：

• 时间复杂度为O(n²)，嵌套循环遍历所有两两组合，当n达到百万级时，运算量达到10^12次，常规算力根本无法完成
• 大量重复开销：循环内反复初始化SkyCoord对象、单次计算分离角，没有利用numpy向量化运算的优势，额外增加了百倍以上的耗时

优化步骤 #

1. 替换循环为向量化运算 #

Astropy的SkyCoord原生支持数组输入，不需要逐点初始化，同时separation方法支持广播运算，可以一次性计算所有两两组合的距离。如果不需要额外的坐标转换功能，直接用numpy实现哈弗辛公式计算天球角距离，速度会比调用Astropy接口快3~10倍：

import numpy as np
import pandas as pd

# 向量化哈弗辛公式计算天球角距离，输入单位为角度，输出单位为角度
def haversine(ra1: np.ndarray, dec1: np.ndarray, ra2: np.ndarray, dec2: np.ndarray) -> np.ndarray:
    ra1, dec1, ra2, dec2 = np.deg2rad([ra1, dec1, ra2, dec2])
    dlon = ra2 - ra1
    dlat = dec2 - dec1
    a = np.sin(dlat / 2) ** 2 + np.cos(dec1) * np.cos(dec2) * np.sin(dlon / 2) ** 2
    c = 2 * np.arcsin(np.sqrt(a))
    return np.rad2deg(c)

2. 分块计算适配百万级数据 #

百万级数据的完整距离矩阵大小为(1e6, 1e6)，需要8PB内存完全无法存储，因此采用分块计算的方案，每次只计算部分数据的距离再累加，控制内存占用：

def calc_total_sum(df: pd.DataFrame, chunk_size: int = 1000) -> float:
    ra_a = df['A'].values
    dec_a = df['A'].values
    ra_b = df['B'].values
    dec_b = df['B'].values
    n = len(df)
    total = 0.0

    for i_start in range(0, n, chunk_size):
        i_end = min(i_start + chunk_size, n)
        # 取当前块的A列坐标
        chunk_ra_a = ra_a[i_start:i_end, None]
        chunk_dec_a = dec_a[i_start:i_end, None]
        for j_start in range(i_start + 1, n, chunk_size):
            j_end = min(j_start + chunk_size, n)
            # 取当前块的B列坐标
            chunk_ra_b = ra_b[None, j_start:j_end]
            chunk_dec_b = dec_b[None, j_start:j_end]
            # 计算当前块对的距离和，累加
            total += haversine(chunk_ra_a, chunk_dec_a, chunk_ra_b, chunk_dec_b).sum()
    return total

# 调用示例
# total = calc_total_sum(df, chunk_size=2000)

可以根据自己的显存/内存大小调整chunk_size参数，chunk越大速度越快，内存占用也越高。

3. 可选极致优化方案 #

• 用numba的@njit装饰器编译哈弗辛函数，运算速度可以再提升2~5倍
• 如果你需要保留Astropy的坐标转换功能，先一次性将所有坐标转为数组格式的SkyCoord对象，再调用separation方法分块计算，不要在循环内初始化SkyCoord
• 确认业务需求是否真的需要计算所有两两组合的距离和，如有可能尽量降维、采样或者用近似统计量替代，从根本上降低运算量

内容的提问来源于stack exchange，提问作者Cruz