如何优化Python双重循环，加速2万只股票与指数的滚动相关度计算

问题背景 #

我有一个包含20000份股票数据文件的文件夹，需要计算每只股票之间的相关度；另有10个指数（示例列表为dex = ['AAA','BBB','CCC','DDD']），需为每个指数计算其与20000只股票的滚动相关度。当前代码运行完成需约3天，现寻求提速方案。

数据文件样例 #

现有代码 #

dex = ['AAA','BBB','CCC','DDD']
roll = 21

for ndx in dex:
    dx = pd.read_csv(base + '/' + f'{ndx}.csv', usecols= ['TimeStamp','Close'])
    dx.columns = ['Date', ndx]
    dx = dx.set_index('Date')
    
    for files in os.listdir(base):

        df = pd.read_csv(base + '/' + files, usecols= ['TimeStamp','Close'])
        tckr = files[:-4]
    
        if tckr != ndx:
            df.columns = ['Date', tckr]
            df = df.set_index('Date')
            sd.append(df)
            print(files, df.shape, ndx)
    
            loo = pd.concat([dx,df], axis =1)
            poo = loo.pct_change()
            poo.dropna(axis = 0, how = 'any', inplace = True)
    
            poo[f'{tckr}_{ndx}_CR'] = round(poo[ndx].rolling(roll).corr(poo[tckr]), 2)

            poo.pop(tckr)
            poo.pop(ndx)

            os.makedirs(output + '/' + ndx + '/', exist_ok = True)
            poo.reset_index().to_csv(f'{output}' + '/' + ndx  + '/' + f'{tckr}.csv', index = False)

提速优化方案 #

1. 批量预处理数据，减少IO操作 #

• 一次性读取所有股票和指数文件，统一完成日期索引设置、涨跌幅计算等预处理步骤，避免嵌套循环中反复读写文件。IO是当前最大性能瓶颈，减少文件操作次数能直接节省大量时间。
• 若内存不足，可将股票数据分批次读取处理，避免加载全部数据导致内存溢出。

2. 向量化计算替代嵌套循环 #

• 现有代码的嵌套循环（指数→股票）效率极低，利用Pandas向量化特性，针对单个指数，一次性计算其与所有股票的滚动相关度。
• 示例：将所有股票的涨跌幅合并为一个DataFrame，对指数的涨跌幅序列调用rolling(roll).corr(other=stock_returns_df)，直接得到该指数与所有股票的滚动相关结果矩阵。

3. 优化内存占用 #

• 数据类型压缩：将Close列从float64转为float32（若精度满足需求），日期列转为datetime类型后设为索引，减少内存消耗，提升计算速度。
• 清理冗余操作：删除代码中未使用的sd.append(df)语句，避免不必要的内存占用。

4. 并行计算拆分任务 #

• 利用multiprocessing或concurrent.futures实现多进程并行处理，每个进程负责一部分股票与指数的相关度计算（各股票计算逻辑独立，无依赖）。
• 示例：对每个指数，将股票列表拆分为多个子列表，分配给不同进程并行计算，同时输出结果文件。

5. 优化滚动相关度计算逻辑 #

• Pandas内置的rolling.corr存在冗余计算，可手动实现滚动窗口相关度的矩阵运算，或用numba对自定义计算函数做JIT编译，加速数值计算过程。

6. 精简冗余操作 #

• 将os.makedirs(output + '/' + ndx + '/', exist_ok = True)移到指数循环的开头，每个指数仅创建一次目录，无需每次循环重复检查。
• 移除不必要的print语句，减少IO耗时。

7. 替换数据存储格式 #

• 将原始CSV文件批量转换为Parquet或Feather格式，这类格式读写速度远快于CSV，且支持列存储，能进一步减少内存占用和读取时间。后续计算直接读取转换后的文件即可。

内容的提问来源于stack exchange，提问作者Divyansh Kumar Singh