使用matplotlib绘制线性回归预测曲线出现异常折线该如何修改？

问题根因 #

你出现绘制结果错乱的问题，核心是两个错误：

• 绘图时x轴参数错误：plt.plot(X_test, y_pred, "g-")中传入的第一个参数是经过多项式特征转换后的X_test矩阵，该矩阵包含多列衍生特征，并非原始的一维x坐标，直接用来当x轴会出现大量错乱折线
• 测试集点顺序混乱：train_test_split会随机打乱数据集再拆分，测试集对应的原始x值是无序的离散点，即使使用正确的x坐标连线，也会画出来回跳转的折线，无法得到平滑的预测曲线

修改方案 #

你需要额外生成一组0~1区间内的等间距连续点，用训练好的模型预测这组点的输出后再绘图，完整修改后的代码如下：

import matplotlib.pyplot as plt
import numpy as np
from sklearn import linear_model
from sklearn.preprocessing import PolynomialFeatures
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
from math import sqrt

def f(x):
    return np.sin(np.pi * 2 * x) + np.random.normal(scale=0.1, size=len(x))

def main():
    NUM_SAMPLES = 15
    x = np.random.uniform(0, 1, NUM_SAMPLES)
    y = f(x)

    plt.plot(x, y, "bo")

    poly = PolynomialFeatures(degree=3)
    X = poly.fit_transform(np.array(x).reshape(-1,1))

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33)

    regr = linear_model.LinearRegression()
    regr.fit(X_train, y_train)
    y_pred = regr.predict(X_test)

    print(regr.coef_)
    # 新增：生成连续点绘制平滑曲线
    x_plot = np.linspace(0, 1, 100).reshape(-1, 1)
    X_plot = poly.transform(x_plot)
    y_plot_pred = regr.predict(X_plot)
    plt.plot(x_plot, y_plot_pred, "g-")
    
    plt.show()

    rmse = sqrt(mean_squared_error(y_test, y_pred))
    print("RMSE = %.2f" % rmse)

if __name__ == "__main__":
    main()

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