三自变量线性回归模型3D拟合曲面绘制IndexError报错求解

错误原因 #

LinearRegression的predict()方法输出的y_pred是一维数组，仅包含每个样本对应的预测结果，不存在第二个维度，你使用y_pred[0:,1]对第二维进行索引自然会触发维度不匹配报错。同时plot_trisurf()需要传入三个一维数组分别对应X、Y、Z三个轴的数值，你之前的写法仅传入两个参数也不符合调用要求。

修正后的3D绘图代码 #

这里选取测试集中的人均GDP（x_test第一列）、识字率（x_test第二列）作为3D图的两个横轴，预测的婴儿死亡率作为Z轴绘制拟合曲面，同时叠加真实测试样本点做对照：

from mpl_toolkits.mplot3d import Axes3D

plt3d = plt.figure(figsize=(10,7)).gca(projection='3d')
# 设定三个轴的数值：x轴为人均GDP，y轴为识字率，z轴为预测的婴儿死亡率
x_axis = x_test[:, 0]
y_axis = x_test[:, 1]
z_axis_pred = y_pred
# 绘制拟合曲面
plt3d.plot_trisurf(x_axis, y_axis, z_axis_pred, alpha=0.5, cmap='viridis', antialiased=True)
# 叠加真实测试样本点
plt3d.scatter(x_axis, y_axis, y_test, c='red', s=30, label='真实样本')
# 调整视角
plt3d.view_init(azim=135, elev=20)
# 坐标轴标注
plt3d.set_xlabel('人均GDP（美元）')
plt3d.set_ylabel('识字率（%）')
plt3d.set_zlabel('婴儿死亡率')
plt3d.legend()
plt.show()

可选优化 #

如果想要更平滑的拟合曲面，可以对自变量范围做网格采样后再预测生成Z轴数值，代码如下：

# 生成网格采样点
x_min, x_max = x_axis.min(), x_axis.max()
y_min, y_max = y_axis.min(), y_axis.max()
xx, yy = np.meshgrid(np.linspace(x_min, x_max, 100), np.linspace(y_min, y_max, 100))
# 固定第三个自变量为样本均值（也可以选其他固定值）
third_feature_mean = np.full_like(xx.ravel(), x_test[:,2].mean())
# 生成输入特征矩阵预测Z值
grid_x = np.c_[xx.ravel(), yy.ravel(), third_feature_mean]
zz = regressor.predict(grid_x).reshape(xx.shape)

plt3d = plt.figure(figsize=(10,7)).gca(projection='3d')
plt3d.plot_surface(xx, yy, zz, alpha=0.5, cmap='viridis', antialiased=True)
plt3d.scatter(x_axis, y_axis, y_test, c='red', s=30, label='真实样本')
plt3d.view_init(azim=135, elev=20)
plt3d.set_xlabel('人均GDP（美元）')
plt3d.set_ylabel('识字率（%）')
plt3d.set_zlabel('婴儿死亡率')
plt3d.legend()
plt.show()

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