Python使用Prophet时循环映射字典触发KeyError的解决方法

问题背景 #

使用Facebook Prophet库预测运动员肌肉围度随时间的变化，数据包含各肌肉（如Neck、Arm）的测量值及时间戳。先通过相关性分析筛选回归变量，编写训练预测函数后，遍历肌肉调用函数时出现KeyError: 'Neck'，提示目标肌肉不在regressors字典中。

错误根源 #

1. 字典赋值笔误：相关性计算代码中定义了regressors={}，但赋值时误用了regressor[target]（少了末尾的s），导致实际生成的有效字典是regressor而非regressors，后续调用regressors自然找不到对应键。
2. 循环逻辑混乱：遍历预测的代码存在多处错误：
   • necessary_columns定义在循环外，且提前调用regressors[Muscle]时Muscle变量未初始化
   • 循环内重复赋值group = group[necessary_columns]会导致数据丢失
   • 调用train_and_forecast时参数顺序错误，缺少muscle参数

修复方案及完整代码 #

1. 修正相关性分析的字典赋值 #

将赋值语句的regressor改为regressors，确保回归变量字典正确填充：

df_selected = df[['Taille', 'Neck', 'Arm', 'Shoulder', 'Chest', 'Waist', 'Hips','Leg']]

# 计算相关矩阵
correlation_matrix = df_selected.corr()

regressors = {}
# 遍历每个目标变量
for target in df_selected.columns:
    # 筛选强相关变量（绝对值>0.4）
    strong_correlations = correlation_matrix[target].abs() > 0.40
    strong_correlators = strong_correlations[strong_correlations].index.tolist()
    
    # 移除目标变量自身
    if target in strong_correlators:
        strong_correlators.remove(target)
    
    # 仅当存在强相关变量时存入字典
    if strong_correlators:
        regressors[target] = strong_correlators  # 修正为regressors

# 输出结果
for target, correlators in regressors.items():
    print(f"{target}: {correlators}")

2. 修正遍历预测的循环逻辑 #

将列定义移到循环内，使用字典get方法避免KeyError，同时修正参数传递和数据处理：

for_loop_forecast = pd.DataFrame()

# 遍历每个肌肉部位
for muscle in Muscle_list:
    # 动态获取当前肌肉的回归变量，无匹配则返回空列表
    relevant_regressors = regressors.get(muscle, [])
    # 构建需要的列
    necessary_columns = ['ds', 'Muscle', 'y'] + relevant_regressors
    # 获取当前肌肉的分组数据（按需调整，示例为从原数据筛选对应肌肉）
    group = df[df['Muscle'] == muscle][necessary_columns].copy()
    # 调用训练函数，传递正确参数
    forecast = train_and_forecast(group, muscle, regressors)
    # 拼接预测结果
    for_loop_forecast = pd.concat([for_loop_forecast, forecast], ignore_index=True)

# 查看结果
for_loop_forecast.head()

3. 保留训练函数的鲁棒性 #

原训练函数已通过regressors.get(muscle, [])处理无回归变量的情况，保持不变即可：

def train_and_forecast(group, muscle, regressors):
    m2 = Prophet()
    
    relevant_regressors = regressors.get(muscle, [])
    necessary_columns = ['ds', 'y'] + relevant_regressors
    filtered_group = group[necessary_columns].copy()
    filtered_group['Muscle'] = muscle

    # 添加回归变量
    for regressor in relevant_regressors:
        m2.add_regressor(regressor)
    
    m2.fit(filtered_group)
    
    # 生成未来时间序列
    future = m2.make_future_dataframe(periods=45)
    # 注意：若回归变量为时间序列，需提前填充future中对应列的未来值
    forecast = m2.predict(future)[['ds', 'yhat', 'yhat_lower', 'yhat_upper']]
    forecast['Muscle'] = muscle
    
    return forecast[['ds', 'Muscle', 'yhat', 'yhat_upper', 'yhat_lower']]

额外提示 #

• 若回归变量本身也是需要预测的时间序列，生成future后必须填充这些变量的未来值，否则Prophet会报错。
• 检查Muscle_list中的肌肉名称与regressors的键是否完全匹配（大小写、拼写一致）。

内容的提问来源于stack exchange，提问作者Julien C.