基于lmfit实现多数据集4组耦合ODE全局拟合报错求助

尝试用lmfit对多数据集执行4组耦合常微分方程的全局拟合，单数据集代码可正常运行，但改写后的全局拟合代码出现KeyError: 's0'，问题出在初始状态量s0-s3的参数配置与函数调用不匹配。

单数据集耦合ODE求解代码 #

import numpy as np
from lmfit import minimize, Parameters
from scipy.integrate import odeint

# 示例时间序列与数据
t = np.linspace(0, 10, 100)
data = np.random.rand(100)

def gauss(x, amp, sigma, center):
    """Gaussian lineshape."""
    return amp * np.exp(-(x-center)**2 / (2.*sigma**2))

def f(xs, t, ps):
    """耦合ODE模型"""
    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()

# 设置参数
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)

# 拟合
result = minimize(residual, params, args=(t, data), method='leastsq')
final = data + result.residual.reshape(data.shape)

原改写的全局拟合代码（存在KeyError问题） #

import numpy as np
from lmfit import minimize, Parameters, report_fit
from scipy.integrate import odeint

# 模拟多数据集：6个数据集，每个100个时间点
data = np.random.rand(6, 100)
x = np.linspace(0, 10, 100)

def gauss(x, amp, sigma, center):
    """Gaussian lineshape."""
    return amp * np.exp(-(x-center)**2 / (2.*sigma**2))

def f(xs, t, ps):
    """耦合ODE模型"""
    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):
    """ODE求解"""
    x = odeint(f, x0, t, args=(params,))
    return x

def testmodel(params, ts, data):
    # 尝试读取无后缀的s0，全局参数仅存在带后缀的s0_1、s0_2等，触发KeyError
    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):
    x0 = (params[f's0_{i+1}'].value, params[f's1_{i+1}'].value, params[f's2_{i+1}'].value, params[f's3_{i+1}'].value)
    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}'].value
    # 调用testmodel时未传递当前数据集参数，仍会查找无后缀参数
    return testmodel(params, x, data)

def objective(params, x, data):
    ndata, _ = data.shape
    resid = 0.0*data[:]
    for i in range(ndata):
        resid[i, :] = data[i, :] - testmodel_dataset(params, i, x)
    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)

# 耦合sigma参数
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)

问题分析与修复 #

出现KeyError: 's0'的核心原因：

• testmodel函数仍尝试读取无后缀的s0、amp等参数，但全局拟合的参数都带有数据集后缀（如s0_1、amp_2），无法匹配。
• testmodel_dataset没有正确传递当前数据集的参数给ODE求解流程，而是复用了单数据集的testmodel逻辑。

修复后的全局拟合代码：

import numpy as np
from lmfit import minimize, Parameters, report_fit
from scipy.integrate import odeint

# 模拟多数据集：6个数据集，每个100个时间点
data = np.random.rand(6, 100)
t = np.linspace(0, 10, 100)

def gauss(x, amp, sigma, center):
    """Gaussian lineshape."""
    return amp * np.exp(-(x-center)**2 / (2.*sigma**2))

def f(xs, t, ps):
    """耦合ODE模型，ps为当前数据集的参数字典"""
    amp = ps['amp']
    center = ps['center']
    sigma = ps['sigma']
    T1 = ps['T1']
    Teq = ps['Teq']
    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):
    """求解ODE，ps为当前数据集的参数字典"""
    x = odeint(f, x0, t, args=(ps,))
    return x

def testmodel(x0, b, ps, ts):
    """生成单个数据集的模型值"""
    ode_result = g(ts, x0, ps)
    # 计算模型表达式
    model = (((ode_result[:,0] - ode_result[:,1] + ode_result[:,2] + b*ode_result[:,3]))**2) / (ode_result[0,0]**2)
    return model

def testmodel_dataset(params, i, ts):
    """处理单个数据集的拟合逻辑"""
    # 获取当前数据集的初始状态
    x0 = (params[f's0_{i+1}'].value, 
          params[f's1_{i+1}'].value, 
          params[f's2_{i+1}'].value, 
          params[f's3_{i+1}'].value)
    # 获取当前数据集的参数，打包成字典传给ODE
    ps_dict = {
        'amp': params[f'amp_{i+1}'].value,
        'center': params[f'center_{i+1}'].value,
        'sigma': params[f'sigma_{i+1}'].value,
        'T1': params[f'T1_{i+1}'].value,
        'Teq': params[f'Teq_{i+1}'].value
    }
    b = params[f'b_{i+1}'].value
    # 生成模型值
    return testmodel(x0, b, ps_dict, ts)

def objective(params, ts, data):
    """全局拟合的目标函数，计算所有数据集的残差"""
    ndata, _ = data.shape
    resid = np.zeros_like(data)
    for i in range(ndata):
        model = testmodel_dataset(params, i, ts)
        resid[i, :] = data[i, :] - model
    return resid.flatten()

# 构建全局参数集
fit_params = Parameters()
n_datasets = data.shape[0]
for iy in range(n_datasets):
    # 初始状态参数
    fit_params.add(f's0_{iy+1}', value=1, vary=False)  # 若需要拟合可设vary=True
    fit_params.add(f's1_{iy+1}', value=0, vary=False)
    fit_params.add(f's2_{iy+1}', value=0, vary=False)
    fit_params.add(f's3_{iy+1}', value=0, vary=False)
    # ODE相关参数
    fit_params.add(f'amp_{iy+1}', value=0.5)
    fit_params.add(f'center_{iy+1}', value=5)
    fit_params.add(f'sigma_{iy+1}', value=0.1)
    fit_params.add(f'T1_{iy+1}', value=0.3)
    fit_params.add(f'Teq_{iy+1}', value=0.7)
    # 模型系数
    fit_params.add(f'b_{iy+1}', value=-1)

# 耦合多个数据集的sigma参数（让它们共享同一个值）
for iy in range(1, n_datasets):
    fit_params[f'sigma_{iy+1}'].expr = 'sigma_1'

# 执行全局拟合
out = minimize(objective, fit_params, args=(t, data), method='leastsq')
report_fit(out.params)

关键修改点 #

1. 重构testmodel函数：不再从全局参数中读取值，改为直接接收当前数据集的初始状态x0、系数b、ODE参数字典ps和时间序列ts，避免参数名称不匹配问题。
2. 修改f函数：简化参数读取逻辑，直接从传入的参数字典中获取值，适配全局拟合的参数传递方式。
3. 优化testmodel_dataset：收集当前数据集的所有参数，打包成字典传递给ODE求解流程，确保每个数据集使用自己的参数（或耦合后的共享参数）。
4. 修正参数引用：所有参数引用都加上数据集后缀，确保与全局参数集中的名称一致。

内容的提问来源于stack exchange，提问作者Keflux