如何加速地理空间聚类任务中time_matrix的计算？

问题背景 #

处理地理空间数据聚类任务时，需要结合地理距离（Haversine）与时间距离生成自定义距离矩阵。现有代码在处理1000行数据时，时间差矩阵的计算成为性能瓶颈（耗时约25秒），无法支撑200-500k行的大规模数据集。

样例数据：

latitude    longitude   timestamp
412671  52.506136   6.068709    2017-01-01 00:00:23.518
412672  52.503316   6.071496    2017-01-01 00:01:30.764
412673  52.505122   6.068912    2017-01-01 00:02:30.858
412674  52.501792   6.068605    2017-01-01 00:03:38.194
412675  52.508105   6.075160    2017-01-01 00:06:41.116

现有瓶颈代码：

timestamps = np_data[:,2]
time_matrix = np.abs(np.subtract.outer(timestamps, timestamps)) # 速度极慢

优化方案 #

1. 优先将时间戳转换为数值类型（核心优化） #

datetime对象的逐元素运算效率极低，先将时间戳转换为整数类型的时间值（如秒级或毫秒级），再进行矩阵计算：

import pandas as pd
import numpy as np

# 替换原时间戳处理逻辑
# 将datetime转为秒级整数（也可根据需求转毫秒：//10**6）
timestamps = pd.to_datetime(df['timestamp']).values.astype(np.int64) // 10**9

# 用广播替代subtract.outer，效果相同但效率更高
time_matrix = np.abs(timestamps[:, np.newaxis] - timestamps[np.newaxis, :])
time_matrix /= time_matrix.max()

这一步能将时间计算速度提升数十倍，因为整数运算比datetime对象运算的开销小得多。

2. 使用Numba JIT编译加速（进一步优化） #

如果仍需要更快的计算速度，可以用Numba将循环编译为机器码，避免numpy的额外开销：

from numba import jit

@jit(nopython=True)
def compute_time_matrix(timestamps):
    n = len(timestamps)
    time_matrix = np.zeros((n, n), dtype=np.float64)
    for i in range(n):
        ts_i = timestamps[i]
        for j in range(n):
            time_matrix[i, j] = abs(ts_i - timestamps[j])
    max_val = time_matrix.max()
    return time_matrix / max_val

# 先转换时间戳为数值类型
timestamps = pd.to_datetime(df['timestamp']).values.astype(np.int64) // 10**9
time_matrix = compute_time_matrix(timestamps)

3. 避免计算全距离矩阵（针对超大数据集的关键方案） #

对于200-500k行的数据，全距离矩阵的大小会达到(500000,500000)，需要约1.9TB的内存（float64类型），完全无法存储。此时应放弃预计算全矩阵，改用支持自定义距离的聚类算法，在算法运行时按需计算距离：

以sklearn的DBSCAN为例：

from sklearn.cluster import DBSCAN
from sklearn.metrics.pairwise import haversine_distances

# 预先计算全局归一化参数
lat_lon_rad = np.radians(df[['latitude','longitude']].values.astype(float))
max_geo_dist = haversine_distances(lat_lon_rad).max()
timestamps = pd.to_datetime(df['timestamp']).values.astype(np.int64) // 10**9
max_time_diff = timestamps.max() - timestamps.min()

# 自定义距离函数
def custom_distance(x, y):
    # 计算Haversine地理距离
    geo_dist = haversine_distances(
        np.radians(x[:2].reshape(1,2)), 
        np.radians(y[:2].reshape(1,2))
    )[0,0]
    # 计算时间差
    time_diff = abs(x[2] - y[2])
    # 归一化后加权合并
    return 0.5 * (geo_dist / max_geo_dist) + 0.5 * (time_diff / max_time_diff)

# 组装包含数值时间戳的数据集
data = np.hstack([df[['latitude','longitude']].values, timestamps.reshape(-1,1)])

# 运行DBSCAN聚类
dbscan = DBSCAN(eps=0.1, min_samples=5, metric=custom_distance)
cluster_labels = dbscan.fit_predict(data)

这种方式无需存储全矩阵，仅在算法需要时计算两两样本的距离，大幅降低内存占用和计算时间。

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