如何加速基于多列的分组滚动窗口Brier Score计算？

原方案使用rolling.apply()调用自定义函数实现计算，这种逐窗口的Python级循环在百万级数据上效率极低——每个窗口都要重复进行索引查找与零散计算，完全没利用pandas的向量化优化能力。

优化思路 #

Brier Score的本质是窗口内(y_true - y_prob)²的均值，可拆解为：
$$Brier = \frac{1}{N} \sum_{i=1}^N (y_{true,i} - y_{prob,i})^2$$
基于这个数学特性，我们可以用向量化的滚动求和+滚动计数替代自定义函数，彻底避免循环开销。

优化代码 #

import pandas as pd
import numpy as np
import random
from pandas._libs.tslibs.timestamps import Timestamp

ROWS = 1_000_000  # 模拟百万级数据集

# 生成测试数据（保留原逻辑）
def create_random_dates(start: Timestamp, end: Timestamp, n: int): 
    divide_by = 24*60*60*10**9
    start_u = start.value // divide_by
    end_u = end.value // divide_by
    return pd.to_datetime([random.randint(start_u, end_u) for p in range(n)], unit="D") 

random.seed(1)
start = pd.to_datetime('2015-01-01')
end = pd.to_datetime('2018-01-01')
random_dates = create_random_dates(start, end, ROWS)
df = pd.DataFrame(
    {
        "id_": list(range(ROWS)),
        "date": random_dates,
        "group": [random.randint(1, 2) for p in range(ROWS)],
        "y_true": [random.randint(0, 1) for p in range(ROWS)],
        "y_prob": [random.random() for p in range(ROWS)],
    }
)
df.sort_values(["group", "date"], inplace=True)
df.reset_index(drop=True, inplace=True)

# 快速计算分组滚动Brier Score
# 1. 先计算每个样本的平方项
df['brier_term'] = (df['y_true'] - df['y_prob']) ** 2

# 2. 按分组执行滚动窗口操作
rolling_window = df.groupby('group').rolling('1000d', min_periods=3, closed='left')

# 3. 滚动求和除以滚动计数，得到Brier Score
df['brier'] = rolling_window['brier_term'].sum() / rolling_window['brier_term'].count()

# 清理中间列（可选）
df.drop('brier_term', axis=1, inplace=True)

性能提升核心原因 #

• 向量化底层实现：rolling.sum()和rolling.count()是pandas用C实现的批量操作，比Python自定义函数循环快100倍以上
• 无额外索引开销：原方案每次apply都要通过df.loc[series.values]重复查找数据，优化后直接在原数据集上做批量计算，避免了冗余的索引操作
• 内存高效：中间列brier_term一次性生成，后续滚动操作都是对该列的批量处理，内存占用更可控

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