Python中计算5000维销售相关矩阵的高效方法（解决pandas df.corr()内存问题）

问题根源 #

pandas原生df.corr()的高内存开销来自三方面：默认采用float64全精度存储与计算、计算过程中生成大量临时Series对象、未利用相关矩阵天然对称的特性做计算与存储的冗余裁剪。5000个产品的规模无需动用高配置内存机器，通过以下方案即可完成计算：

具体优化手段 #

• 全链路降低数值精度
  把输入销售数据、中间计算变量、最终存储的相关矩阵全部从默认的float64转为float32，单份矩阵的内存占用直接降低50%。计算皮尔逊/斯皮尔曼相关性时，float32的精度完全匹配销售数据的分析要求，不会产生可感知的误差。实现时可在数据加载阶段就通过.astype('float32')指定列类型，避免后续pandas自动提升精度造成的内存浪费。
• 替换为numpy原生矩阵计算
  放弃pandas逐列循环计算的逻辑，直接用numpy.corrcoef()完成全量矩阵计算。该方法会自动基于矩阵对称特性减少重复运算，全程无多余的临时DataFrame/Series对象生成，同等数据量下内存开销仅为pandas原生corr()的1/3~1/2。对应实现代码为：

  import numpy as np
  # rowvar=False表示每一列是一个独立变量，即单个产品的销售时序
  corr_matrix = np.corrcoef(df.values.astype('float32'), rowvar=False)
  

• 利用对称矩阵特性减半存储开销
  相关矩阵天然满足corr(i,j) = corr(j,i)，计算时仅需遍历所有i<j的位置计算对应值，无需计算重复的对称位置元素；存储时也仅需保存上三角部分的数值，需要全矩阵展示时再临时拼接对称部分即可，可再减少一半的内存占用。
• 分块加载计算，避免全量数据驻留
  如果原始销售数据的时间维度较长（如逐小时/逐门店拆分的明细数据），不需要一次性把5000个产品的全量数据加载到内存。可将产品按每200~500个为一组拆分，每次仅加载当前组的销售数据，计算该组与所有其他产品的相关性，算完就释放原始数据的内存，仅把最终的相关值写入预分配好的矩阵数组中，全程内存占用可控制在常规办公设备可承载的范围内。
• 按需裁剪矩阵存储
  如果最终分析仅需要保留强相关性的产品对（如相关系数绝对值大于0.5），计算过程中遇到低于阈值的相关值直接丢弃，最终用稀疏矩阵存储有效相关值，整体内存占用还能再降低80%以上。

注意：如果之前使用df.corr()时传入了自定义的逐行/逐列函数作为计算方法，会进一步放大内存开销，这类场景必须替换为向量化实现才能解决问题。

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