from sklearn import linear_model
from matplotlib import pyplot as plt
import numpy as np
raw_xtrain = [8, 3, 9, 7, 16, 5, 13, 10, 4, 6]
ytrain = [30, 21, 33, 27, 42, 24, 36, 33, 22, 25]
raw_xtest = [4.5, 6, 14]
ytest = [24, 25, 38]
xtrain = []
xtest = []
for i in raw_xtrain:
xtrain.append([i])
for i in raw_xtest:
xtest.append([i])
clf = linear_model.LinearRegression()
clf.fit(xtrain, ytrain)
y_pred = clf.predict(xtest)
k = clf.coef_
b = clf.intercept_
x0 = np.arange(0, 20, 0.2)
y0 = k* x0 + b
# y = kx + b
mt = k * np.array(raw_xtrain) + b - np.array(ytrain)
mt = np.multiply(mt, mt)
print('train MSE={:.2f}'.format(1/10 * sum(mt)))
mt =k * np.array(raw_xtest) + b - np.array(ytest)
mt = np.multiply(mt, mt)
print('test MSE={:.2f}'.format(1/3 * sum(mt)))
print('figure:')
plt.title('LinearRegression')
plt.xlabel('X')
plt.ylabel('Y')
plt.grid()
plt.plot(raw_xtrain, ytrain, 'bo', label='train data')
plt.plot(raw_xtest, ytest, 'ro', label='test data')
plt.plot(raw_xtest, y_pred, 'go', label='pred')
plt.plot(x0, y0, 'black',label='pred line')
plt.legend()
plt.show()
# train MSE=0.92
# test MSE=0.67