Python中集群级别Bootstrap统计量计算的提速方法求助

原代码性能瓶颈分析 #

• 数据拼接效率极低：循环中逐个执行pd.concat，每次拼接都会生成新DataFrame，带来大量内存开销和时间消耗。
• 重复筛选集群数据：每次循环都重新执行df[df[cluster] == p]，没有复用已有的集群子集。
• 单线程串行执行：上万次Bootstrap迭代依次运行，完全没利用多核CPU资源。

优化步骤与完整代码 #

核心优化点：预存集群子集+批量拼接+并行计算

import numpy as np
import pandas as pd
import statsmodels.api as sm
from joblib import Parallel, delayed

def bootstrap_rsquared_single_iteration(cluster_groups, model_formula, time, cov_struct, dep_data, fam, cluster, seed):
    """单个Bootstrap迭代的计算逻辑，供并行调用"""
    ids = list(cluster_groups.keys())
    # 有放回重采样集群ID
    resampled_ids = np.random.choice(ids, size=len(ids), replace=True, random_state=seed)
    # 批量拼接选中的集群数据，替代循环逐个拼接
    dfs = pd.concat([cluster_groups[cid] for cid in resampled_ids], axis=0)
    # 拟合模型并计算伪R平方
    model = sm.GEE.from_formula(
        formula=model_formula,
        groups=cluster,
        data=dfs,
        time=time,
        cov_struct=cov_struct,
        dep_data=dep_data,
        family=fam
    )
    return model.fit().pseudo_rsquared()

def bootstrap_rsquared(df, model_formula, time, cov_struct, dep_data, fam,
                       cluster='ID', n_resamples=10000, n_jobs=-1):
    '''提供分组Bootstrap的GEE伪R平方值'''
    # 预存每个集群的数据，避免重复筛选
    cluster_groups = {cid: group for cid, group in df.groupby(cluster)}
    
    # 并行执行所有Bootstrap迭代，n_jobs=-1表示使用所有可用CPU核心
    rsqs = Parallel(n_jobs=n_jobs)(
        delayed(bootstrap_rsquared_single_iteration)(
            cluster_groups, model_formula, time, cov_struct, dep_data, fam, cluster, seed=i
        ) for i in range(n_resamples)
    )
    
    return rsqs

额外优化建议 #

• 先小批量测试：先用n_resamples=100验证代码正确性，确认没问题再跑全量10000次。
• 内存优化：如果数据集极大，可以考虑将集群子集存储为numpy数组（需注意Statsmodels兼容性），进一步降低内存占用。
• 随机状态复用：如果不需要严格按迭代数作为种子，可以提前生成一组随机种子数组，避免重复创建随机数生成器。

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