Fit a Polynomial Model

polyplot.py

import matplotlib.pyplot as plt
import numpy as np

x_data = np.arange(0, 10)
y_data = [n + np.random.normal() for n in x_data]

def poly_model(x_data, coeffs):
    '''Fit data to a linear polynomial model
    
    Args:
        x_data: The input data
        coeffs: the polynomial coefficients as a numpy array
    Return:
        The model response data.
    '''
    M = [x_data**i for i, b in enumerate(coeffs)]
    X = np.array(M)
    b = np.array(coeffs).reshape((-1, 1)) # get column vector
    return X.T@b

def polyplot(x_data, y_data, degree, ax):
    '''Plot a polynomial to input data
    
    Args:
        x_data: the input data
        y_data: the reponse data
        degree: polynomial degree
        ax: a motplotlib axis for drawing.
    Return:
        Tuple of the axis with the plotted data including a default title and a tuple of fit coefficients
    '''
    coeffs = np.polyfit(x_data, y_data, degree)
    model = poly_model(x_data, np.flip(coeffs))
    ax.scatter(x_data, y_data)
    ax.plot(x_data, model)
    ax.set_title(f"Polynomial Fit: Degree {len(coeffs) - 1}")
    return ax, coeffs

fig, ax = plt.subplots()
ax, coeffs = polyplot(x_data, y_data, 1, ax)
fig.tight_layout()
print(coeffs)
plt.show()