Global fit coupled odes system lmfit - global

Question

I'm trying to get the global fit of multiple set of data considering a system of 4 coupled ODEs.
I have the working code that solves the system of 4 coupled ODEs for a single set of data, and I have the working code that do the global fit with an arbitrary function (but not using odeint).
My problem is that I'm not able to merge the two codes...
Code for coupled ODEs
t =
data=
def gauss(x, amp, sigma, center):
"""Gaussian lineshape."""
return amp * np.exp(-(x-center)**2 / (2.*sigma**2))
def f(xs, t, ps):
"""Lotka-Volterra predator-prey model."""
try:
amp = ps['amp'].value
center = ps['center'].value
sigma = ps['sigma'].value
T1 = ps['T1'].value
Teq = ps['Teq'].value
except Exception:
amp, center, sigma, T1, Teq = ps
s0,s1,s2,s3 = xs
return [- gauss(t,amp,sigma,center) * (s0-s1),\
gauss(t,amp,sigma,center) * (s0-s1) - s1/T1,\
(s1/T1 - s2/Teq),\
(s2/Teq)]
def g(t, x0, ps):
x = odeint(f, x0, t, args=(ps,))
return x
def residual(ps, ts, data):
x0 = ps['s0'].value, ps['s1'].value, ps['s2'].value, ps['s3'].value
b = ps['b'].value
model = (((g(ts, x0, ps)[:,0]-g(ts, x0, ps)[:,1]+g(ts, x0, ps)[:,2]+b*g(ts, x0, ps)[:,3]))**2)/((g(ts, x0, ps)[0,0]))**2
return (model - data).ravel()
# set parameters incluing bounds
params = Parameters()
params.add('s0', value=1, vary=False)
params.add('s1', value=0, vary=False)
params.add('s2', value=0, vary=False)
params.add('s3', value=0, vary=False)
params.add('amp', value=0.02)
params.add('center', value=5)
params.add('sigma', value=0.1)
params.add('T1', value=0.3)
params.add('Teq', value=0.7)
params.add('b', value=-1)
# fit model and find predicted values
result = minimize(residual, params, args=(t, data), method='leastsq')
final = data + result.residual.reshape(data.shape)
Considering the code here: https://lmfit.github.io/lmfit-py/examples/example_fit_multi_datasets.html
I've tried to do by myself the code for global fit in this case
def gauss(x, amp, sigma, center):
"""Gaussian lineshape."""
return amp * np.exp(-(x-center)**2 / (2.*sigma**2))
def f(xs, t, ps):
"""Lotka-Volterra predator-prey model."""
try:
amp = ps['amp'].value
center = ps['center'].value
sigma = ps['sigma'].value
T1 = ps['T1'].value
Teq = ps['Teq'].value
except Exception:
amp, center, sigma, T1, Teq = ps
s0,s1,s2,s3 = xs
return [- gauss(t,amp,sigma,center) * (s0-s1),\
gauss(t,amp,sigma,center) * (s0-s1) - s1/T1,\
(s1/T1 - s2/Teq),\
(s2/Teq)]
def g(t, x0, params):
"""
Solution to the ODE x'(t) = f(t,x,k) with initial condition x(0) = x0
"""
x = odeint(f, x0, t, args=(params,))
return x
def testmodel(params, ts, data):
x0 = params['s0'].value, params['s1'].value, params['s2'].value, params['s3'].value
b = params['b'].value
model = (((g(ts, x0, params)[:,0]-g(ts, x0, params)[:,1]+g(ts, x0, params)[:,2]+b*g(ts, x0, params)[:,3]))**2)/((g(ts, x0, params)[0,0]))**2
return model
def testmodel_dataset(params, i, x):
"""Calculate Gaussian lineshape from parameters for data set."""
x0 = params[f's0_{i+1}'], params[f's1_{i+1}'], params[f's2_{i+1}'], params[f's3_{i+1}']
amp = params[f'amp_{i+1}']
center = params[f'center_{i+1}']
sigma = params[f'sigma_{i+1}']
T1 = params[f'T1_{i+1}']
Teq = params[f'Teq_{i+1}']
b = params[f'b_{i+1}']
return testmodel(params, x, data)
def objective(params, x, data):
"""Calculate total residual for fits of Gaussians to several data sets."""
ndata, _ = data.shape
resid = 0.0*data[:]
# make residual per data set
for i in range(ndata):
resid[i, :] = data[i, :] - testmodel_dataset(params, i, x)
# now flatten this to a 1D array, as minimize() needs
return resid.flatten()
fit_params = Parameters()
for iy, y in enumerate(data):
fit_params.add(f's0_{iy+1}', value=1)
fit_params.add(f's1_{iy+1}', value=0)
fit_params.add(f's2_{iy+1}', value=0)
fit_params.add(f's3_{iy+1}', value=0)
fit_params.add(f'amp_{iy+1}', value=0.5)
fit_params.add(f'center_{iy+1}', value=0.5)
fit_params.add(f'sigma_{iy+1}', value=0.5)
fit_params.add(f'T1_{iy+1}', value=0.5)
fit_params.add(f'Teq_{iy+1}', value=0.4)
fit_params.add(f'b_{iy+1}', value=0.3)
for iy in (2, 3, 4, 5, 6):
fit_params[f'sigma_{iy}'].expr = 'sigma_1'
out = minimize(objective, fit_params, args=(x, data))
report_fit(out.params)
Result -> KeyError: 's0'
There is a problem with x0 and s0,s1,s2,s3, population of the four states.
I'm sorry if the question may be very naive...
Thank you for your help.