from sklearn import linear_model
import numpy as np
import matplotlib.pyplot as plt


# 元データの作成
x = np.arange(start=0.,stop=1., step=0.02)
y = []
for xi in x:
    # y = 4xi+1+noize
    yi = 4*xi + 1 + np.random.random()
    y.append(yi)


# 元データのプロット
plt.xlim(0, 1)
plt.scatter(x,y)
plt.savefig("test_data.jpg")

# 単回帰のパラメータを求める
reg = linear_model.LinearRegression()
x = np.array(x).reshape(-1,1)
y = np.array(y).reshape(-1,1)
result = reg.fit(x,y)
a = result.coef_[0][0]
b = result.intercept_[0]
print("a:", a)
print("b:", b)

# 単回帰の結果プロット
y_hat = []
for xi in x:
    yi_hat = a*xi + b
    y_hat.append(yi_hat)

plt.plot(x, y_hat, c="r", label="$ax+b$")
plt.legend()
plt.savefig("result.jpg")
plt.show()