How to write a wrapper to fix arbitrary parameters of the jacobian of a function

Question

This is an extension of a previous stack-exchange question I posted. link

Context:

My goal is to fit data to a function f(t, *p) using the scipy.optimize.curve_fit function. I happen to know some parameters pfix = {p_j, ..., p_k} and want to fit f(t, *p) to my data with the parameters in pfix... fixed.

Above I have a link asking how I could write a wrapper for fixing parameters in a function. Now I want to do the same thing, but for the jacobian of f(t, *p), fixing parameters pfix. I am not sure how to do this.

Wrapper for func(x, *p)

Below is the wrapper for my function:

def fix_params(f, fix_pars):
    # fix_pars = ((1, A), (2, B))
    def new_func(x, *pars):
        new_pars = [None]*(len(pars) + len(fix_pars))
        for j, fp in fix_pars:
            new_pars[j] = fp
        for par in pars:
            for j, npar in enumerate(new_pars):
                if npar is None:
                    new_pars[j] = par
                    break
        return f(x, *new_pars)
    return new_func

Issue

Naively I would just use this wrapper for my Jacobian function. However, here is the problem.

Let's say I have N parameters and M values for x. Then my jacobian function returns a (M, N) numpy array. Now this is fine if I don't fix any parameters. However, even when I fix just one parameter, my wrapped jacobian function still returns a (M, N) numpy array. This causes curve_fit to complain, as the number of parameters I use is now less than the parameter dimension of my jacobian. I am not sure how to get around this.

Any suggestions?

Answer 1

This should work (using as in my comment scipy.optimize.leatsq)

import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import leastsq

def arb_func( x, a, b, c, d ):
    return a * np.exp( b * np.sin( c * x + d ) )

def arb_fixed( fixed ):
    def g( x , *pars ):
        locPars = list( pars )
        for f in sorted( list(fixed), key=lambda x: x[0] ):
            locPars.insert( f[0], f[1] )
        return arb_func( x, *locPars)
    return g

def func_da( x, a, b, c, d ):
    return np.exp( b * np.sin( c * x + d ) )

def func_db( x, a, b, c, d ):
    return arb_func( x, a, b, c, d ) * np.sin( c * x + d )

def func_dc( x, a, b, c, d ):
    return arb_func( x, a, b, c, d ) * b * np.cos( c * x + d ) * x

def func_dd( x, a, b, c, d ):
    return arb_func( x, a, b, c, d ) * b * np.cos( c * x + d )

dList = [ func_da, func_db, func_dc, func_dd ]

def jac( pars, x, y ):
    return [ func( x, *pars ) for func in dList ]

def jac_fixed( Fixed = None ):
    def h( pars, x, y, z ):
        funcList = dList[::]
        locFixed = sorted( list(Fixed), key=lambda x: x[0] )
        locPars = list( pars )
        for f in locFixed:
            locPars.insert( f[0], f[1] )
        locFixed.reverse()
        for f in locFixed:
            del funcList[ f[0] ]
        out = [ func( x, *locPars ) for func in funcList ]
        return out
    return h

a0 = +3
b0 = -0.6
c0 = +1.44
d0 = +0.4

xList = np.linspace( -2, 4, 100 )
y0List = np.fromiter( ( arb_func( x, a0, b0, c0, d0 ) for x in xList ), np.float )
yNoiseList = np.fromiter( ( y + .2 * np.random.normal() for y in y0List[::4] ), np.float )
xNoiseList = xList[::4]

def residuals_fixed( params, xList, yList, Fixed=None ):
    if Fixed is None:
        fixedSorted = []
    else:
        fixedSorted = sorted( list(Fixed), key=lambda x: x[0] )
    locParams = list( params )
    for f in fixedSorted:
            locParams.insert( f[0], f[1] )
    diff = [ arb_func( x, *locParams ) - y for x, y in zip( xList, yList )]
    return diff

def leastsq_wrapper( xList, yList, p0, **kwargs ):
    fixed = kwargs.pop( 'Fixed', None )
    if fixed is None:
        locFixed = []
    else:
        locFixed = fixed
        s = np.array( locFixed ).shape
        if len(s) !=2 or s[-1] !=2:
            raise ValueError( 'fixed value list has wrong shape. Must be n by 2, but is {}'.format(s) )
    if len( p0 ) + len( locFixed ) != 4:
            raise TypeError( 'Total number of arguments (variable + fixed) is not 4' )
    fixedSorted = sorted( list( locFixed ), key=lambda x: x[0] )
    if not all( [ ( type( item[0] ) is int )  and ( item[0] > -1 ) and ( item[0] < 4 ) for item in fixedSorted ] ):
        raise ValueError( 'list indices i for fixed values are not int with 0 <= i < 4' )
    my_jac = jac_fixed( Fixed=fixedSorted )
    baseDict = { 'args':( xList, yList, fixed ), 'Dfun':my_jac, 'col_deriv':1}
    baseDict.update(kwargs) ## allows to send e.g.  full_output=True
    out = leastsq( residuals_fixed, p0, **baseDict )
    return out

myFitStd, err = leastsq( residuals_fixed, [ a0, b0 ,c0 , d0 ], args=( xNoiseList, yNoiseList ) )
print myFitStd
myFit0, err = leastsq_wrapper( xNoiseList, yNoiseList,  [ a0, b0 ,c0 , d0 ] )
print myFit0
myFixed1 = [[0,3.3]]
myFit1, err = leastsq_wrapper( xNoiseList, yNoiseList,  [ b0 ,c0 , d0 ], Fixed=myFixed1 )
arb1 = arb_fixed( myFixed1 )
print myFit1

myFixed2 = [ [ 3, .8], [2, 1.2 ] ]
myFit2, err = leastsq_wrapper( xNoiseList, yNoiseList,  [ a0, b0 ], Fixed=myFixed2 )
arb2 = arb_fixed( myFixed2 )
print myFit2

fig = plt.figure()
ax = fig.add_subplot( 1, 1, 1 )

ax.plot( xList, y0List )
ax.plot( xNoiseList, yNoiseList, marker='o', linestyle='' )
ax.plot( xList, np.fromiter( ( arb_func( x, *myFitStd ) for x in xList ), np.float), linestyle='--' )
ax.plot( xList, np.fromiter( ( arb_func( x, *myFit0 ) for x in xList ), np.float), linestyle=':' )
ax.plot( xList, np.fromiter( ( arb1( x, *myFit1 ) for x in xList ), np.float) )
ax.plot( xList, np.fromiter( ( arb2( x, *myFit2 ) for x in xList ), np.float) )

plt.show()

Giving the output

[ 3.03802692 -0.57275564  1.43380277  0.38557492]
[ 3.03802692 -0.57275564  1.43380277  0.38557493]
[-0.49087778  1.422561    0.40503389]
[ 3.31028289 -0.46678563]