Is there a cleaner way to achieve curve fitting with multiple models?

Question

In my project I have multiple family of functions predefined to fit a curve. Let's look at the simplest:

def polyfit3(x, b0, b1, b2, b3):
    return b0+b1*x+b2*x**2+b3*x**3

def polyfit2(x, b0, b1, b2):
    return b0+b1*x+b2*x**2

def polyfit1(x, b0, b1):
    return b0+b1*x

Note: I know that in this particular case np.polyfit would be a better choice

The (much more simplyfied) function, which makes the fitting looks like this:

from scipy.optimize import curve_fit
try:
    from lmfit import Model
    _has_lmfit = True
except ImportError:
    _has_lmfit = False

def f(x, y, order=3):
    if _has_lmfit:
        if order == 3:
            fitModel = Model(polyfit3)
            params = fitModel.make_params(b0=0, b1=1, b2=1, b3=1)
            result = fitModel.fit(y, x=x, params=params)
        elif order == 2:
            fitModel = Model(polyfit2)
            params = fitModel.make_params(b0=0, b1=1, b2=1)
            result = fitModel.fit(y, x=x, params=params)
        elif order == 1:
            fitModel = Model(polyfit1)
            params = fitModel.make_params(b0=0, b1=1)
            result = fitModel.fit(y, x=x, params=params)
        else:
            raise ValueError('Order is out of range, please select from [1, 3].')
    else:
        if order == 3:
            popt, pcov = curve_fit(polyfit3, x, y)
            _function = polyfit3
        elif order == 2:
            popt, pcov = curve_fit(polyfit2, x, y)
            _function = polyfit2
        elif order == 1:
            popt, pcov = curve_fit(polyfit1, x, y)
            _function = polyfit1
        else:
            raise ValueError('Order is out of range, please select from [1, 3].')
    # more code there.. mostly working with the optimized parameters, plotting, etc.

My problem is this gets really ugly quickly and I repeat myself over and over again. Is there a way to build this better?

EDIT:

I've tried this:

def poly_fit(x, *args):
    return sum(b*x**i for i, b in enumerate(args))

...

fitModel = Model(poly_fit)
fitModel.make_params(**{f'b{i}': 1 for i in range(order+1)})

but unfortunately lmfit throws an error:

ValueError: varargs '*args' is not supported

Answer 1

I think that lmfit.models.PolynomialModel() does exactly what you are looking for. That model takes the polynomial degree n as an argument and uses coefficients named c0, c1, ..., cn (handling up to n=7):

from lmfit.models import PolynomialModel

def f(x, y, degree=3):
    fitModel = PolynomialModel(degree=degree)
    params = fitModel.make_params(c0=0, c1=1, c2=1, c3=0, 
                                  c4=0, c5=0, c6=0, c7=0)
    # or if you prefer to do it the hard way:
    params = fitModel.make_params(**{'c%d'%i:0 for i in range(degree+1)})

    return fitModel.fit(y, x=x, params=params)

Note that it is OK to over-specify coefficients here. That is, if degree=3, that call to fitModel.make_params(c0=0, ..., c7=0) will actually not make parameters for c4, c5, c6, or c7.

PolynomialModel will raise a TypeError if degree > 7, so I left your explicit test for that out.

I hope that gets you started, but it does seem like you might have wanted to include other model functions too. In a case like that, what I have done is to make a dictionary of class names:

from lmfit.models import LinearModel, PolynomialModel, GaussianModel, ....

KnownModels = {'linear': LinearModel, 'polynomial': PolynomialModel, 
              'gaussian': GaussianModel, ...}

and then use that to construct the model:

modelchoice = 'linear' # probably really came from user selection in a GUI

if modelchoice in KnownModels:
    model = KnownModels[modelchoice]()
else:
    raise ValueError("unknown model '%s'" % modelchoice)

params = model.make_params(....) # <- might know and store what the parameter names are
.....