如何在Plotnine的facet_wrap/grid中添加回归线方程？

要在Plotnine的分面图中为每个子图添加线性回归方程和R²，核心是按分面分组计算回归参数，再将标注文本映射到对应分面，具体实现如下：

完整实现代码 #

import plotnine as p9
from scipy import stats
from plotnine.data import mtcars as df
import pandas as pd

# 1. 按分面变量（gear）分组计算回归参数与标注文本
def calc_regression_stats(group):
    slope, intercept, r_value, _, _ = stats.linregress(group['wt'], group['mpg'])
    r_squared = r_value ** 2
    # 格式化方程文本，保留合适的小数位数
    eq_text = f'y = {slope:.2f}x + {intercept:.2f}; R² = {r_squared:.2f}'
    # 返回每个分组的标注位置（选每个分面的右上角，可根据数据调整）
    return pd.DataFrame({
        'gear': [group['gear'].iloc[0]],
        'eq_text': [eq_text],
        'x': [group['wt'].max() * 0.7],  # x位置设为权重最大值的70%
        'y': [group['mpg'].max() * 0.9]   # y位置设为油耗最大值的90%
    })

# 生成每个分面对应的标注数据
annot_df = df.groupby('gear').apply(calc_regression_stats).reset_index(drop=True)

# 2. 绘制分面图并添加标注
plot = (
    p9.ggplot(data=df, mapping=p9.aes('wt', 'mpg', color='factor(gear)'))
    + p9.geom_point()
    + p9.facet_wrap('~ gear')
    + p9.geom_smooth(method='lm', formula='y~x', se=False, color='black')
    + p9.xlab('Wt')
    + p9.ylab('MPG')
    # 添加每个分面的回归方程标注
    + p9.geom_text(
        data=annot_df,
        mapping=p9.aes(x='x', y='y', label='eq_text'),
        color='black',
        size=9
    )
)

print(plot)

关键细节说明 #

• 分组计算回归参数：用groupby('gear')按分面变量分组，通过自定义函数calc_regression_stats计算每组的斜率、截距和R²，并格式化方程文本。
• 标注位置调整：x和y的取值根据每组数据的极值动态计算，保证标注在每个分面的合理位置（避免遮挡数据），可根据你的数据集调整比例。
• geom_text映射分面：通过指定data=annot_df，让标注文本自动匹配对应的分面（gear字段对应分面变量），实现每个子图显示专属的回归方程。

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