Pandas滚动窗口多均值计算报错及实现方案问询

错误原因分析 #

你遇到的IndexError: too many indices for array是因为pandas的rolling.apply()默认会逐列处理窗口数据：当你对DataFrame调用该方法时，它会把每一列的51个窗口元素作为一维numpy数组传给test函数，而不是你预期的二维数组（51行×3列）。这就导致你在函数里使用data[0:31,:]这种二维索引时出错——一维数组根本没有第二维。

修正方案 #

我们需要调整两个关键点：

1. 让rolling.apply()把整个窗口的二维数据（而非单列一维数据）传给自定义函数；
2. 修改自定义函数，支持多列并行计算，并返回每列的结果以便后续处理。

修正后的完整代码 #

import pandas as pd
import numpy as np

# 生成示例数据
df = pd.DataFrame(np.random.randint(0,800,size=(1000, 3)), columns=list('ABC'))

# 重写自定义计算函数：接收整个窗口的DataFrame，返回每列的计算结果
def test(data):
    # 将DataFrame转为二维numpy数组（51行×3列）
    data_arr = data.to_numpy()
    # 初始化存储每列均值序列的数组（51行×3列）
    meanMov = np.zeros((51, 3))
    
    # 第一阶段：前16个位置用前31行的均值
    mean_first = np.mean(data_arr[0:31, :], axis=0)
    meanMov[0:16, :] = mean_first
    
    # 第三阶段：后16个位置用20-49行的均值（注意data[20:50]是20到49，共30行）
    mean_last = np.mean(data_arr[20:50, :], axis=0)
    meanMov[35:51, :] = mean_last
    
    # 第二阶段：中间16-34位置用滑动的31窗口均值（i-15到i+15，共31个元素）
    for i in range(16, 35):
        meanMov[i, :] = np.mean(data_arr[(i-15):(i+15+1), :], axis=0)
    
    # 返回每列的整体均值（3个元素，对应A/B/C列）
    return meanMov.mean(axis=0)

# 调用rolling.apply：设置raw=False，让函数接收整个窗口的DataFrame；result_type='expand'让结果展开为多列
r = df.rolling(51)
entr = r.apply(test, raw=False, result_type='expand').dropna(axis=0, how='all')

# 查看结果：entr会有3列，对应原DataFrame的A/B/C列的计算结果
print(entr.head())

关键改动说明 #

1. raw=False：告诉pandas不要把数据转成一维数组，而是把整个窗口的DataFrame传给test函数，这样我们就能拿到完整的二维数据。
2. result_type='expand'：让apply返回的多列结果直接展开为DataFrame的列，而不是压缩成一维Series。
3. 函数逻辑调整：
   • 先把传入的DataFrame转为numpy数组方便索引操作；
   • 所有均值计算都保留列维度，最后返回每列的整体均值（而非所有列的均值）；
   • 修正了索引范围：data[20:50]对应的是第20到第49行（左闭右开），和你原来的逻辑一致。

后续处理 #

entr是一个和原DataFrame列数一致的DataFrame，每一列对应原列的计算结果，你可以直接把它和原数据合并，或者单独保存为文件：

# 合并原数据和计算结果
combined_df = pd.concat([df, entr.add_suffix('_rolling_mean')], axis=1).dropna()

# 保存到CSV文件
combined_df.to_csv('rolling_calculations.csv', index=False)

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