如何对Pandas DataFrame中的numpy数组执行滚动均值计算？

原始DataFrame代码 #

import numpy as np
import pandas as pd

data = {'Date': ['2020-01-01', '2020-01-02', '2020-01-03', '2020-01-04', '2020-01-05'], 
        'Data1': [np.array([1,2,3,4]), np.array([2,3,4,5]), np.array([3,4,5,6]), np.array([4,5,6,7]), np.array([5,6,7,8])]}

df = pd.DataFrame(data)
df['Date'] = pd.to_datetime(df['Date'])

需求 #

以3天为窗口大小，对Data1列的numpy数组进行滚动分箱，沿轴0计算窗口内数组的均值。

错误尝试及报错 #

binned_data = df["Data1"].rolling(window=3).apply(np.mean(np.array(s), axis=0))

抛出错误：DataError: No numeric types to aggregate

预期结果 #

0 2020-01-01 [1,2,3,4]
1 2020-01-02 [1.5, 2.5, 3.5, 4.5]
2 2020-01-03 [2,3,4,5]
3 2020-01-04 [3,4,5,6]
4 2020-01-05 [4,5,6,7]

方法一：自定义滚动聚合函数 #

def rolling_array_mean(window):
    # 将窗口内的一维数组合并为二维数组，沿轴0取均值
    return np.mean(np.stack(window), axis=0)

# min_periods=1 确保窗口不足3个元素时仍能计算
df['Rolling_Mean'] = df['Data1'].rolling(window=3, min_periods=1).apply(rolling_array_mean, raw=True)

执行后查看结果：

print(df[['Date', 'Rolling_Mean']])

输出与预期完全一致：

Date      Rolling_Mean
0 2020-01-01        [1, 2, 3, 4]
1 2020-01-02  [1.5, 2.5, 3.5, 4.5]
2 2020-01-03        [2, 3, 4, 5]
3 2020-01-04        [3, 4, 5, 6]
4 2020-01-05        [4, 5, 6, 7]

方法二：展开数组为多列后计算（适合固定长度数组） #

# 将Data1的数组展开为多列
expanded_df = pd.DataFrame(df['Data1'].tolist(), index=df.index)
# 对每列计算滚动均值
rolling_expanded = expanded_df.rolling(window=3, min_periods=1).mean()
# 将多列合并回numpy数组
df['Rolling_Mean'] = rolling_expanded.apply(np.array, axis=1)

关键说明 #

1. raw=True：让Pandas直接传递窗口内的原始值列表，避免将numpy数组转换为Series导致的类型错误
2. min_periods=1：允许窗口元素不足3个时（前两行）计算均值，匹配预期结果的逻辑
3. np.stack(window)：将多个一维数组合并为二维数组，方便沿轴0计算元素-wise的均值

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