使用polyfit对数据进行拟合，选取1-5次方中拟合最好的并返回loss值、次方数、系数

写一个函数，使用polyfit进行数据拟合，从一次拟合（polyfit(data_x, data_y, 1)）到5次拟合（polyfit(data_x, data_y, 5)）中选取loss值最小的，返回loss值、次方数、系数。

先贴一下整体代码

import numpy as np
import matplotlib.pyplot as plt
from numpy import polyfit


def bestFit(data_x, data_y):
    min_loss = fitData(data_x, data_y, 1, False)
    bestFitTime = 1
    for i in range(4):
        m = fitData(data_x, data_y, i + 2, False)
        if m['loss'] < min_loss['loss']:
            min_loss['loss'] = m['loss']
            bestFitTime = i + 2
            coefficient = m['coefficient']

    fitData(data_x, data_y, bestFitTime, True)
    return {'min_loss': min_loss['loss'], 'bestFitTime': bestFitTime, 'coefficient': coefficient}


def fitData(data_x, data_y, fit_time, draw):	#进行拟合
    coefficient = polyfit(data_x, data_y, fit_time)		#使用polyfit进行fit_time次拟合
    loss = dealData(data_x, data_y, coefficient, draw)
    return {'loss': loss, 'coefficient': coefficient}


def dealData(data_x, data_y, coefficient, draw):
    # plt.plot(data_x , data_y , 'rx')
    xline = np.linspace(min(data_x), max(data_x), len(data_x) * 8)
    func = ""
    for i in range(len(coefficient)):
        func += str(coefficient[i]) + " * pow(x , " + str(len(coefficient) - i - 1) + ") + "

    func += "0"
    y = []
    for x in xline:
        p = eval(func)
        y.append(p)
    # print(func)
    if draw:
        plt.plot(data_x, data_y, 'rx')
        plt.plot(xline, y, 'b-')
        plt.show()

    loss = 0
    for i in range(len(data_x)):
        x = data_x[i]
        p = pow(eval(func) - data_y[i], 2)
        loss += p
    loss = loss / len(data_x)
    return loss


def Test():		#测试
    xdemo = [1, 2, 3, 4, 5]
    ydemo = [1, 2, 3, 4, 5]
    # loss = fitData(xdemo, ydemo, 5, True)
    # print(loss)

    x = np.linspace(-5, 5, 100)
    y = 4 * x + 1.5
    noise_y = y + np.random.randn(y.shape[-1]) * 2.5
    loss = bestFit(xdemo, ydemo)
    print(loss)

Test()

注释

拟合函数：使用polyfit进行fit_time次拟合，同时调用处理数据函数dealData()得到loss值。传入数据分别为x,y值，fit_time采用几次拟合，draw代表是否绘制图像，值为True或False。

def fitData(data_x, data_y, fit_time, draw):	#进行拟合
    coefficient = polyfit(data_x, data_y, fit_time)		#使用polyfit进行fit_time次拟合
    loss = dealData(data_x, data_y, coefficient, draw)
    return {'loss': loss, 'coefficient': coefficient}

处理数据计算loss、绘图

def dealData(data_x, data_y, coefficient, draw):
    # plt.plot(data_x , data_y , 'rx')
    xline = np.linspace(min(data_x), max(data_x), len(data_x) * 8)
    func = ""
    for i in range(len(coefficient)):
        func += str(coefficient[i]) + " * pow(x , " + str(len(coefficient) - i - 1) + ") + "

    func += "0"
    y = []
    for x in xline:
        p = eval(func)
        y.append(p)
    # print(func)
    if draw:
        plt.plot(data_x, data_y, 'rx')
        plt.plot(xline, y, 'b-')
        plt.show()

    loss = 0
    for i in range(len(data_x)):
        x = data_x[i]
        p = pow(eval(func) - data_y[i], 2)
        loss += p
    loss = loss / len(data_x)
    return loss

这里有个eval()函数，直接看这里的解释，我这里写了一个小函数，通过你给我的系数list，我直接返还给你整个f(x),以下面的代码举个例子。

def Test2(coefficient):
    func = ""
    for i in range(len(coefficient)):
        func += str(coefficient[i]) + " * pow(x , " + str(len(coefficient) - i - 1) + ") + "
    func += "0"
    print(func)
    
Test2([5,4,3,2,1])

运行结果：

这里再使用一次以x为变量的循环再配上eval()函数，就可以求得对应的f(x)的list，大致这样：

    for i in range(len(data_x)):
        x = data_x[i]
        p = pow(eval(func) - data_y[i], 2)
        loss += p
    loss = loss / len(data_x)

从1至5进行遍历，选出loss最小的返回

def bestFit(data_x, data_y):
    min_loss = fitData(data_x, data_y, 1, False)
    bestFitTime = 1
    for i in range(4):
        m = fitData(data_x, data_y, i + 2, False)
        if m['loss'] < min_loss['loss']:
            min_loss['loss'] = m['loss']
            bestFitTime = i + 2
            coefficient = m['coefficient']

    fitData(data_x, data_y, bestFitTime, True)
    return {'min_loss': min_loss['loss'], 'bestFitTime': bestFitTime, 'coefficient': coefficient}