最小二乗法の三次元版

最小二乗法の三次元版っていうのを scikit-learn でやってみました。

# データセットの取得
from sklearn.datasets import load_iris

dataset = load_iris()
X = dataset.data[50:100, [0, 1]]
Y = dataset.data[50:100, 2]

# 最小二乗法による線形回帰
from sklearn.linear_model import LinearRegression

model = LinearRegression()
model.fit(X, Y)

LinearRegression()

# 三次元プロット
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure()
ax = Axes3D(fig)
ax.scatter(X[:, 0], X[:, 1], Y)

plot = np.array([
                 [X.min(axis=0)[0], X.min(axis=0)[1]], 
                 [X.max(axis=0)[0], X.max(axis=0)[1]]
                 ])

ax.plot(plot[:, 0], plot[:, 1], model.predict(plot))

ax.set_xlabel('X0')
ax.set_ylabel('X1')
ax.set_zlabel('Y')

plt.show()

こんな感じで合ってますか。