pythonNA_lagrange.py

# -*- coding: utf8 -*-

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import scipy.interpolate

xSize = 15
fig = plt.figure()

def onclick(event):
    print('button=%d, x=%d, y=%d, xdata=%f, ydata=%f' %
          (event.button, event.x, event.y, event.xdata, event.ydata))

def onkey(event):
    global xSize
    if(event.key == "up"):
        xSize  = xSize +1
    elif (event.key == "down"):
        xSize = xSize - 1
    testInterpolateFit()


fig.canvas.mpl_connect('button_press_event', onclick)
fig.canvas.mpl_connect('key_press_event', onkey)

def func_cubic(x, a, b, c, d):
    return a * x ** 3 + b * x ** 2 + c*x + d

def genData():
    xdata = np.linspace(-3, 4, xSize)
    y = func_cubic(xdata, 1, 1, 1, 1)
    y_noise = 3 * np.random.normal(size=xdata.size)
    ydata = y + y_noise
    return xdata, ydata

def testInterpolateFit():
    global xSize
    plt.clf()
    xdata, ydata = genData()
    plt.plot(xdata, ydata, 'b.', label='data,  size=' + str(xSize) )

    # 按照三次曲线拟合, popt 是方程 f 的系数， pcov是预估的协方差
    popt_cubic, pcov_cubic = curve_fit(func_cubic, xdata, ydata)
    plt.plot(xdata, func_cubic(xdata, *popt_cubic), 'g-', label='cubic fit')

    lagarangePoly = scipy.interpolate.lagrange(xdata, ydata)
    print(lagarangePoly)
    plt.plot(xdata, lagarangePoly(xdata), 'r--.', label='p(x)')

    plt.xlabel('x'); plt.ylabel('y'); plt.legend()
    plt.show()

testInterpolateFit()