如何修复回归拟合绘图问题并为训练数据集添加预测列？

嘿，我来帮你搞定这两个问题！咱们一步步来：

一、修复绘图的散点颜色和图例混乱问题 #

你当前的问题出在每次调用plt.plot()时都重复传入了X, Y, 'yo'，这相当于每次都重新绘制了原始数据的黄色散点，不仅覆盖了之前的黑色散点，还导致图例里混入了重复的原始数据条目。

解决方法：只保留一开始的plt.scatter()绘制黑色原始散点，后续的plt.plot()只专注于绘制各阶拟合的回归线即可：

修改后的绘图代码片段：

plt.scatter(X,Y, color = 'black', label = 'Original Data')
# 仅绘制各阶拟合线，不再重复绘制原始数据
plt.plot(X, linear_fn(X), color='red', linestyle='--', label='Linear')
plt.plot(X, quadratic_fn(X), color='green', linestyle='--', label='Quadratic')
plt.plot(X, order7_fn(X), color='purple', linestyle='--', label='Order 7')

plt.title('Training Data - Fits')
plt.legend()
plt.show()

这样修改后，黑色散点只会显示一次，每条回归线用你指定的颜色和虚线样式，图例也会清晰对应各阶模型。

二、高效添加预测列到训练DataFrame #

np.poly1d返回的拟合函数是向量化的，直接作用于整个Series就能生成预测值，这是最高效的方式（比循环、apply快得多，因为底层是numpy的C级运算）。

直接用以下代码添加预测列：

# 向量化生成各阶模型的预测值，直接赋值为新列
train['linear_pred'] = linear_fn(train['X'])
train['quadratic_pred'] = quadratic_fn(train['X'])
train['order7_pred'] = order7_fn(train['X'])

运行后，你的train DataFrame就会新增三列，分别对应1阶、2阶、7阶模型的预测值。

完整修改后的代码 #

%matplotlib inline
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

X = np.linspace(0, 18, num=10)
Y = (0.05*X**2)+4*np.sin(X)+1.15*X+2*np.random.randn()
combi = {'X':X, 'Y':Y}
train = pd.DataFrame(combi)
print(train)

# 拟合各阶模型
linear = np.polyfit(X, Y, 1)
linear_fn = np.poly1d(linear)
quadratic = np.polyfit(X, Y, 2)
quadratic_fn = np.poly1d(quadratic)
order7 = np.polyfit(X, Y, 7)
order7_fn = np.poly1d(order7)

# 修复后的绘图逻辑
plt.scatter(X,Y, color = 'black', label = 'Original Data')
plt.plot(X, linear_fn(X), color='red', linestyle='--', label='Linear')
plt.plot(X, quadratic_fn(X), color='green', linestyle='--', label='Quadratic')
plt.plot(X, order7_fn(X), color='purple', linestyle='--', label='Order 7')

plt.title('Training Data - Fits')
plt.legend()
plt.show()

# 高效添加预测列
train['linear_pred'] = linear_fn(train['X'])
train['quadratic_pred'] = quadratic_fn(train['X'])
train['order7_pred'] = order7_fn(train['X'])
print("\n添加预测列后的数据集：")
print(train)

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