python 回归线性（或多项式）

Posted 2021-05-09

def polyfit(namex,x,namey, y, degree):
    import numpy as np
    import matplotlib.pyplot as plt

    # initialize
    results = {}

    # calculate polynomial regression
    coeffs = np.polyfit(x, y, degree)

    # build model (function)
    p = np.poly1d(coeffs)

    # fit values, and mean
    yhat = p(x)                      # or [p(z) for z in x]
    ybar = np.sum(y)/len(y)          # or sum(y)/len(y)
    ssreg = np.sum((yhat-ybar)**2)   # or sum([ (yihat - ybar)**2 for yihat in yhat])
    sstot = np.sum((y - ybar)**2)    # or sum([ (yi - ybar)**2 for yi in y])
 
    # Set regression info 
    results['polynomial'] = coeffs.tolist() # Polynomial Coefficients
    results['R2'] = ssreg / sstot # R**2
    results['correlation'] = np.corrcoef(x,y)[0,1]  # correlation index

    ## PLOT

    # build object
    fig, ax = plt.subplots()
    # plot scatter
    ax.scatter(x,y, color="blue")
    # set title    
    plt.title("REGRESSION grade %s- R**2: %.3f Correlation: %.3f"%(degree,results['R2'],results['correlation']))
    # axes labels
    plt.xlabel(namex)
    plt.ylabel(namey)
    # plot grid axis
    ax.xaxis.grid(True)
    ax.yaxis.grid(True)
    # plot fitted line
    ax.plot(x,yhat, label='fit', color="red")
    # display
    plt.show()

    # return info results
    return results