Search code examples
pythonnumpyscipycurve-fittingpolynomials

How to fit a polynomial with some of the coefficients constrained?


Using NumPy's polyfit (or something similar) is there an easy way to get a solution where one or more of the coefficients are constrained to a specific value?

For example, we could find the ordinary polynomial fitting using:

x = np.array([0.0, 1.0, 2.0, 3.0,  4.0,  5.0])
y = np.array([0.0, 0.8, 0.9, 0.1, -0.8, -1.0])
z = np.polyfit(x, y, 3)

yielding

array([ 0.08703704, -0.81349206,  1.69312169, -0.03968254])

But what if I wanted the best fit polynomial where the third coefficient (in the above case z[2]) was required to be 1? Or will I need to write the fitting from scratch?


Solution

  • In this case, I would use curve_fit or lmfit; I quickly show it for the first one.

    import numpy as np
    import matplotlib.pyplot as plt
    from scipy.optimize import curve_fit
    
    def func(x, a, b, c, d):
      return a + b * x + c * x ** 2 + d * x ** 3
    
    x = np.array([0.0, 1.0, 2.0, 3.0,  4.0,  5.0])
    y = np.array([0.0, 0.8, 0.9, 0.1, -0.8, -1.0])
    
    print(np.polyfit(x, y, 3))
    
    popt, _ = curve_fit(func, x, y)
    print(popt)
    
    popt_cons, _ = curve_fit(func, x, y, bounds=([-np.inf, 2, -np.inf, -np.inf], [np.inf, 2.001, np.inf, np.inf]))
    print(popt_cons)
    
    xnew = np.linspace(x[0], x[-1], 1000)
    
    plt.plot(x, y, 'bo')
    plt.plot(xnew, func(xnew, *popt), 'k-')
    plt.plot(xnew, func(xnew, *popt_cons), 'r-')
    plt.show()
    

    This will print:

    [ 0.08703704 -0.81349206  1.69312169 -0.03968254]
    [-0.03968254  1.69312169 -0.81349206  0.08703704]
    [-0.14331349  2.         -0.95913556  0.10494372]
    

    So in the unconstrained case, polyfit and curve_fit give identical results (just the order is different), in the constrained case, the fixed parameter is 2, as desired.

    The plot looks then as follows:

    enter image description here

    In lmfit you can also choose whether a parameter should be fitted or not, so you can then also just set it to a desired value (check this answer).