python-3.x - 将回归线方程和 R 方添加到 PLOTNINE

Question

使用 stat_smooth(method="gls") 很容易在 plotnine 中获得数据的线性最佳拟合。但是，我不知道如何得出最佳拟合线或 R2 值的系数。 R 中的 Ggplot 有一个执行此操作的 stat_regline_equation() 函数，但我在 plotnine 中找不到类似的工具。

目前，我正在使用 statsmodels.formula.api.ols 来获取这些值，但在 plotnine 中必须有更好的方法。

PS:我是所有编码方面的新手。

Answer 1

我最终使用了以下代码；不是 PlotNine 但很容易实现。

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

#calculate best fit line
slope, intercept, r_value, p_value, std_err = stats.linregress(df['wt'],df['mpg'])
df['fit']=df.wt*slope+intercept
#format text 
txt= 'y = {:4.2e} x + {:4.2E};   R^2= {:2.2f}'.format(slope, intercept, r_value*r_value)
#create plot. The 'factor' is a nice trick to force a discrete color scale
plot=(p9.ggplot(data=df, mapping= p9.aes('wt','mpg', color = 'factor(gear)'))
    + p9.geom_point(p9.aes())
    + p9.xlab('Wt')+ p9.ylab(r'MPG')
    + p9.geom_line(p9.aes(x='wt', y='fit'), color='black')
    + p9.annotate('text', x= 3, y = 35, label = txt))
#for some reason, I have to print my plot 
print(plot)