最小二乘法Savitzky-Golay平滑滤波

参考：https://blog.csdn.net/sinat_21258931/article/details/79298478

# -*- coding: utf-8 -*-  
import numpy as np

from numpy import *
import matplotlib.pyplot as plt
from random import *
from scipy.signal import savgol_filter



def loadData():
    x = arange(-1,1,0.02)
    y = ((x*x-1)**3+1)*(cos(x*2)+0.6*sin(x*1.3))
    xr = []
    yr = []
    i = 00
    for xx in x:
        yy = y[i]
        d=float(randint(90,110))/100
        dd=float(randint(80,120))/100
        i+=1  
        xr.append(xx*d)  
        yr.append(yy*dd)  
    return x,y,xr,yr



"""
* 创建系数矩阵X
* size - 2×size+1 = window_size
* rank - 拟合多项式阶次
* x - 创建的系数矩阵
"""
def create_x(size, rank):
    x = []
    for i in range(2 * size + 1):
        m = i - size
        row = [m**j for j in range(rank)]
        x.append(row) 
    x = np.mat(x)
    print("x",x)
    return x


"""
 * Savitzky-Golay平滑滤波函数
 * data - list格式的1×n纬数据
 * window_size - 拟合的窗口大小
 * rank - 拟合多项式阶次
 * ndata - 修正后的值
"""
def savgol(data, window_size, rank):
    m = int((window_size - 1) / 2)
    odata = data[:]
    # 处理边缘数据，首尾增加m个首尾项
    for i in range(m):
        odata.insert(0,odata[0])
        odata.insert(len(odata),odata[len(odata)-1])
    # 创建X矩阵
    x = create_x(m, rank)
    # 计算加权系数矩阵B
    b = (x * (x.T * x).I) * x.T
    print("b",b)
    a0 = b[m]
    print("a0",a0)
    a0 = a0.T
    # 计算平滑修正后的值
    ndata = []
    y2 = []
    for i in range(len(data)):
        y = [odata[i + j] for j in range(window_size)]
        y2.append(y)
    ndata = np.mat(y2) * a0
    return ndata



def figPlot(x1,y1,x2,y2):
    plt.plot(x1,y1,color='g',linestyle='-',marker='')
    plt.plot(x2,y2,color = 'm',linestyle='',marker='.')
    plt.show()

def Main():
    x,y,xr,yr = loadData()
    print(type(x),type(xr))
    myY = savgol(yr,21,3)
    figPlot(x,myY,xr,yr)
    
Main()

效果甚好：