CODEWITHSUNDEEP
pythonscipycurve-fittinglmfit
Péter Leéh
Péter Leéh

Reputation: 2119

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

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

Upvotes: 0

Views: 950

Answers (2)

M Newville
M Newville

Reputation: 7862

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
.....

Upvotes: 2

Grigory Feldman
Grigory Feldman

Reputation: 413

I rewrite your code by creating global config for your polyfit functions. This is more pythonic version of if. 

polyfits = {
    1: {
        'f': polyfit1,
        'params': ['b0', 'b1'],
        'vals'  : [  0,    1], 
    },
    2: {
        'f': polyfit2,
        'params': ['b0', 'b1', 'b2'],
        'vals'  : [   0,    1,   1,], 
    },
    3: {
        'f': polyfit3,
        'params': ['b0', 'b1', 'b2', 'b3'],
        'vals'  : [   0,    1,    1,    1], 
    },

}

def f(x, y, order=3):
    if order not in polyfits.keys():
        raise ValueError('Order is out of range, please select from {}.'.format(','.join(map(str, polyfits.keys()))))
    _function = polyfits[order]['f']
    if _has_lmfit:
        fitModel = Model(_function)
        params = dict(zip(polyfits[order]['params'], polyfits[order]['vals']))
        params = fitModel.make_params(**params)
        result = fitModel.fit(y, x=x, params=params)
    else:
        popt, pcov = curve_fit(_function, x, y)

I believe you post very simplified version of your code (because your current version could be minimized more effectively than my code above).

Upvotes: 1

Related Questions