Statsmodels OLS 一元线性回归

OLS实现一元线性回归

工作年限和薪资情况数据集：

代码实现：

import pandas as pd
import matplotlib.pyplot as plt
import statsmodels.api as sm

'''
工作年限和薪资情况的散点图
'''
# 读入数据
income = pd.read_csv(r'Salary_Data.csv')
# 绘制散点图
plt.scatter(x=income.YearsExperience,  # 指定散点图的x轴数据
            y=income.Salary,  # 指定散点图的y轴数据
            color='steelblue',  # 指定散点图中点的颜色
            label='data'  # 标签
            )
'''
最小二乘法求出一元线性关系模型的a、b参数；输出拟合得到的各个x对应的y值；输出汇总表
'''
model = sm.formula.ols(formula='Salary~YearsExperience', data=income).fit()
print('model params: ', model.params, '\n')
print('model fittedvalues: ', model.fittedvalues, '\n')
print('model summary: ', model.summary(), '\n')
'''
绘制拟合直线
'''
# 'r-'表示红色的 - 线条
plt.plot(income.YearsExperience, model.fittedvalues, 'r-', label='OLS')
plt.legend(loc='best')
# 设置字体Microsoft YaHei或者SimHei
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
plt.style.use('ggplot')
# 添加x轴和y轴标签
plt.xlabel('YearsExperience')
plt.ylabel('Salary')
# 添加标题
plt.title('工作年限和薪资情况的关系')
# 显示图形
plt.show()

plt 绘制散点图和拟合直线：