Python中基于数值积分结果拟合omega_Rabi参数的方法

要通过curve_fit从counts_list和delta_list反推omega_Rabi，核心是把每个delta对应的积分计算过程包装成curve_fit可调用的拟合函数。以下是具体实现方案：

关键思路 #

curve_fit要求拟合函数的输入是自变量（此处为delta）和待拟合参数（omega_Rabi），输出是对应自变量的预测值（即模拟的counts）。因此需要编写一个包装函数，输入delta和omega_Rabi后，内部调用solve_ivp完成积分，返回最终的上态布居数（即counts值）。

完整实现代码 #

import numpy as np
from scipy.integrate import solve_ivp
from scipy.optimize import curve_fit

pi = np.pi

# 固定已知参数
omega = 2*pi*50000
t_init = 0
t_fin = 0.00005
y0 = [0+0j, 1+0j]

# 原微分方程函数
def func(t, y, omega, delta, omega_Rabi): 
    c_up, c_down = y
    dydt = [
        -1j*(omega_Rabi*np.sin(omega*t))*c_down,
        -1j*(omega_Rabi*np.sin(omega*t))*c_up - 1j*delta*c_down
    ]
    return dydt

# 拟合用包装函数：输入delta数组和omega_Rabi，输出对应的counts数组
def fit_func(delta_array, omega_Rabi):
    counts_pred = []
    for delta in delta_array:
        # 对每个delta求解微分方程（仅需最终时刻结果，无需t_eval）
        sol = solve_ivp(
            func, [t_init, t_fin], y0,
            rtol=1e-9, atol=1e-11,
            args=(omega, delta, omega_Rabi)
        )
        # 取最终时刻的上态布居数
        y_up_final = abs(sol["y"][0][-1])**2
        counts_pred.append(y_up_final)
    return np.array(counts_pred)

# ---------------------- 生成模拟实验数据（假设此部分为已知输入） ----------------------
true_omega_Rabi = 2*pi*170*6 
delta_list = 2*pi*np.arange(-10000, 10001, 4000) 
delta_list = delta_list[delta_list != 0] 

counts_list = np.zeros(len(delta_list))
for i in range(len(delta_list)):
    delta = delta_list[i]
    sol = solve_ivp(func, [t_init,t_fin], y0,rtol=1e-9, atol=1e-11, args=(omega,delta,true_omega_Rabi))
    counts_list[i] = abs(sol["y"][0][-1])**2

# ---------------------- 执行拟合 ----------------------
# 给omega_Rabi一个接近真实值的初始猜测（非线性拟合对初始值敏感）
initial_guess = [2*pi*1000]

# 调用curve_fit
popt, pcov = curve_fit(fit_func, delta_list, counts_list, p0=initial_guess)

# 输出结果
print(f"拟合得到的omega_Rabi: {popt[0]}")
print(f"真实omega_Rabi: {true_omega_Rabi}")
print(f"相对误差: {abs(popt[0]-true_omega_Rabi)/true_omega_Rabi*100:.2f}%")

注意事项 #

• 初始猜测值：非线性拟合对初始值敏感，p0需给出接近真实值的猜测（可根据实验范围估算），否则可能拟合失败或陷入局部最优。
• 积分精度：保持solve_ivp的rtol和atol足够高，避免积分误差干扰拟合结果。
• 复数处理：微分方程的解为复数，但最终取模平方得到实数布居数，需确保此步骤计算正确。
• 效率优化：若delta_list长度较大，拟合过程会多次调用solve_ivp，可去掉t_eval参数（仅需最终时刻结果，让solve_ivp自动选择时间步），提升运行速度。

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