python应用 曲线拟合04

python应用 曲线拟合04 → 多项式拟合

---

• 主要是使用 numpy 库中的 polyfit() 函数，见第 66 行， z = np.polyfit(x_proton, y, 3) ，其中待拟合曲线的横纵坐标分别为第 1, 2 个参数，第 3 个参数 3 表示使用 3 次多项式拟合。得到的拟合结果写入返回列表 z，按照幂指数从高到低排序，比如这里 z[0] 代表 x³ 前面的系数，z[1] 代表 x² 前面的系数，以此类推。

import xlrd
import xlwt
import numpy as np
import matplotlib.pyplot as plt


def pol2(x, a, b, c):
    x = np.array(x)
    return a + b*x + c*pow(x, 2)


def pol3(x, a, b, c, d):
    x = np.array(x)
    return a + b*x + c*pow(x, 2) + d*pow(x, 3)


def pol4(x, a, b, c, d, e):
    x = np.array(x)
    return a + b*x + c*pow(x, 2) + d*pow(x, 3) + e*pow(x, 4)


def pol5(x, a, b, c, d, e, f):
    x = np.array(x)
    return a + b*x + c*pow(x, 2) + d*pow(x, 3) + e*pow(x, 4) + f*pow(x, 5)


def pol6(x, a, b, c, d, e, f, g):
    x = np.array(x)
    return a + b*x + c*pow(x, 2) + d*pow(x, 3) + e*pow(x, 4) + f*pow(x, 5) +
        g*pow(x, 6)


def main():
    data_file = xlrd.open_workbook('./data_78As.xlsx')
    sheet_proton = data_file.sheet_by_name('proton')
    sheet_neutron = data_file.sheet_by_name('neutron')

    write_excl = xlwt.Workbook(encoding='utf-8')
    write_sheet_proton = write_excl.add_sheet('proton')
    write_sheet_neutron = write_excl.add_sheet('neutron')

    x_proton = sheet_proton.col_values(0)
    x_proton_add = x_proton[:]
    x_proton_add.insert(0, 0.23)
    x_proton_add.insert(0, 0.22)
    x_proton_add.insert(0, 0.21)
    x_proton_add.insert(0, 0.20)
    x_proton_add.append(0.36)
    x_proton_add.append(0.37)
    x_proton_add.append(0.38)
    x_proton_add.append(0.39)
    x_proton_add.append(0.40)

    j = 0
    for temp in x_proton_add:
        write_sheet_proton.write(j, 0, temp)
        j = j + 1

    for i in range(18):
        y = sheet_proton.col_values(i+1)
        plt.scatter(x_proton, y)
        print(y)
        #  popt, pcov = curve_fit(pol5, x_proton, y, [1, 1, 1, 1, 1, 1])
        #  print(popt)

        z = np.polyfit(x_proton, y, 3)

        y_fit = pol3(x_proton_add, z[3], z[2], z[1], z[0])
        plt.plot(x_proton_add, y_fit, color='green', linewidth=1.0)
        plt.show()
        j = 0
        for temp in y_fit:
            print(x_proton_add[j])
            if x_proton_add[j] <= 0.23:
                write_sheet_proton.write(j, i+1, temp)
                j = j + 1
            elif x_proton_add[j] > 0.23 and x_proton_add[j] < 0.35:
                write_sheet_proton.write(j, i+1, y[j-4])
                j = j + 1
            else:
                write_sheet_proton.write(j, i+1, temp)
                j = j + 1

    plt.clf()
    x_neutron = sheet_neutron.col_values(0)
    x_neutron_add = x_neutron[:]
    x_neutron_add.insert(0, 0.23)
    x_neutron_add.insert(0, 0.22)
    x_neutron_add.insert(0, 0.21)
    x_neutron_add.insert(0, 0.20)
    x_neutron_add.append(0.36)
    x_neutron_add.append(0.37)
    x_neutron_add.append(0.38)
    x_neutron_add.append(0.39)
    x_neutron_add.append(0.40)

    j = 0
    for temp in x_neutron_add:
        write_sheet_neutron.write(j, 0, temp)
        j = j + 1

    for i in range(18):
        y = sheet_neutron.col_values(i+1)
        plt.scatter(x_neutron, y)
        print(y)
        #  popt, pcov = curve_fit(pol5, x_proton, y, [1, 1, 1, 1, 1, 1])
        #  print(popt)

        z = np.polyfit(x_neutron, y, 3)

        y_fit = pol3(x_neutron_add, z[3], z[2], z[1], z[0])
        plt.plot(x_neutron_add, y_fit, color='green', linewidth=1.0)
        plt.show()
        j = 0
        for temp in y_fit:
            if x_neutron_add[j] <= 0.23:
                write_sheet_neutron.write(j, i+1, temp)
                j = j + 1
            elif x_neutron_add[j] > 0.23 and x_neutron_add[j] < 0.35:
                write_sheet_neutron.write(j, i+1, y[j-4])
                j = j + 1
            else:
                write_sheet_neutron.write(j, i+1, temp)
                j = j + 1

    write_excl.save("result_78As.xls")


if __name__ == '__main__':
    main()

得到结果：