sklearn方法求最小

练习题：
'''
给定训练集为     x=1, y=6.8
          x=2, y=9.8
          x=3, y=13.2
          x=4, y=16.2
测试集       x=5, y=? '''

import numpy as np
from matplotlib import pyplot as plt
x_data = [1,2,3,4]
y_data = [6.8,9.8,13.2,16.2]

loss_list = list()
def forward(a,x,b):
    return a*x+b

def lossFunction(a,x,y,b):
    y_pred = forward(a,x,b)
    loss = (y_pred - y)**2
    return loss

a_list = list()
b_list = list()
if __name__ == '__main__':
    for a in np.arange(0,6,0.1):
        for b in np.arange(0,6,0.1):
            sum_loss = 0
            for i in range(4):
                sum_loss += lossFunction(a, x_data[i], y_data[i],b)
            loss_list.append(sum_loss/4)
            a_list.append(a)
            b_list.append(b)

    plt.plot(a_list,loss_list)
    plt.xlabel('a')
    plt.ylabel('loss')
    print(min(loss_list))
    loss_min_index = loss_list.index(min(loss_list))
    print(loss_min_index)
    a_wanted = a_list[loss_min_index]
    b_wanted = b_list[loss_min_index]
    print(f'a_wanted = {a_wanted}, b_wanted ={b_wanted}')
    # plt.show()
    # a_wanted = a_list[loss_list.index(min(loss_list))]
    # print(forward(a_wanted, 4))
    print(forward(a_wanted, 5, b_wanted))