高效计算时间序列相关性的性能优化方案咨询

数据格式 #

我的销售数据格式如下：

[Date, Product, Sales]
01-01, A, 3
01-02, A, 0
01-03, A, 2
...
01-01, Z, 10
01-02, Z, 11
01-03, Z, 13

需求与当前实现问题 #

我需要找出销售模式高度相似或相异的产品，当前采用计算每对产品销售额皮尔逊相关系数的方案，代码如下：

df = pd.DataFrame(..., columns=['Date', 'Product', 'Sales'])
products = df.Product.unique()
correlations = []

for i in range(len(products)):
    a = df[df.Product == products[i]]

    for j in range(i+1, len(products)):
        b = df[df.Product == products[j]]
        
        # 内连接确保日期一致（部分产品存在数据缺失）
        merged = pd.merge(a, b, on='Date', how='inner')
        
        # 计算皮尔逊相关系数
        corr = merged['Sales_x'].corr(merged['Sales_y'])
        correlations.append((products[i], products[j], corr))

目前有大约10k个产品，内连接后日期重叠最多约2k个。上述代码时间复杂度为O(产品数²×日期数)，当前运行时长约30小时，添加线程后性能无明显提升。希望优化代码以高效计算产品间销售相关性，后续计划为每个产品选取相关性Top和Bottom X的产品用于决策（比如分析产品蚕食效应）。

补充示例 #

假设数据包含3个产品A、B、C：

[Date, Product, Sales]
01-01, A, 3
01-02, A, 0
01-03, A, 2

01-01, B, 2
01-02, B, 0
01-03, B, 3

01-01, C, 0
01-02, C, 4
01-03, C, 0 

A与B销售模式高度相似，皮尔逊相关系数接近1；A与C销售模式完全相反，相关系数接近-1，这类关系是我需要分析的核心。

优化方案 #

1. 数据重塑为宽表，利用矩阵运算加速 #

Python循环的效率极低，直接用pandas的向量化/矩阵运算功能，底层是C实现的优化代码，能把运行时间压缩到分钟级：

import pandas as pd

# 将长表转为宽表：行=日期，列=产品，值=销售额
wide_df = df.pivot(index='Date', columns='Product', values='Sales')

# 自动计算所有产品对的皮尔逊相关系数矩阵
# 注：计算时会自动忽略任意一方数据缺失的日期，等价于手动内连接的逻辑
corr_matrix = wide_df.corr(method='pearson')

这个方案直接跳过了Python层面的双重循环，把计算逻辑交给底层优化的矩阵运算，性能提升几个数量级。

2. 内存友好的替代方案（scipy） #

如果宽表内存占用过高，可使用scipy的pdist计算两两相关系数，同样是底层优化的实现：

import pandas as pd
from scipy.spatial.distance import pdist, squareform

wide_df = df.pivot(index='Date', columns='Product', values='Sales')
# 转置为产品×日期的矩阵，适配pdist的输入要求
product_sales = wide_df.T.values

# pdist的correlation metric返回的是(1-皮尔逊相关系数)，转换回相关系数
corr_dist = pdist(product_sales, metric='correlation')
corr_matrix = 1 - squareform(corr_dist)
# 转成DataFrame方便后续处理
corr_matrix = pd.DataFrame(corr_matrix, index=wide_df.columns, columns=wide_df.columns)

3. 高效提取Top/Bottom X相关产品 #

不需要保存完整的10k×10k相关矩阵，直接遍历每个产品提取目标结果，节省内存：

top_x = 5  # 自定义需要的Top相似产品数量
bottom_x = 5  # 自定义需要的Bottom相异产品数量
product_correlations = {}

for product in corr_matrix.columns:
    # 排除产品自身的相关系数（值为1）
    corr_series = corr_matrix[product].drop(product)
    # 获取Top X相似产品
    top_similar = corr_series.nlargest(top_x).to_dict()
    # 获取Bottom X相异产品
    top_dissimilar = corr_series.nsmallest(bottom_x).to_dict()
    product_correlations[product] = {
        'top_similar': top_similar,
        'top_dissimilar': top_dissimilar
    }

4. 并行计算的正确打开方式 #

之前加线程无提升是因为Python GIL限制，CPU密集型任务要用多进程：

import pandas as pd
from joblib import Parallel, delayed
import numpy as np

wide_df = df.pivot(index='Date', columns='Product', values='Sales')

# 将产品列表分块，避免单进程内存压力过大
product_chunks = np.array_split(wide_df.columns, 8)  # 分8块，可根据CPU核心数调整

def compute_corr_chunk(chunk):
    # 计算当前块与全量产品的相关系数
    return wide_df[chunk].corrwith(wide_df)

# 多进程并行计算
chunk_results = Parallel(n_jobs=-1)(delayed(compute_corr_chunk)(chunk) for chunk in product_chunks)
# 合并结果为完整的相关矩阵
corr_matrix = pd.concat(chunk_results, axis=1)

注：如果用pandas/scipy的原生corr已经足够快，这个方案可能提升有限，但在超大规模数据下能进一步缩短时间。

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