Fit for a parameter when the function is obtained by numerical integration in Python

Question

I have the code below in python. What it does is to integrate numerically the function func between 2 values and save the last value in counts_list. One of the parameters of func is omega_Rabi. What I need to do is, after I obtain counts_list I would like to perform a fit of counts_list vs delta_list (knowing that counts_list was obtained from func) and get omega_Rabi. Basically I assume I know counts_list and I want to find omega_Rabi that generate it (which now I assume I don't know). How can I do this? I usually use curve_fit but the stuff I tried so far didn't work. Thank you!

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

pi = np.pi

omega = 2*pi*50000
omega_Rabi = 2*pi*170*6 

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

t_init = 0
t_fin = 0.00005
t_eval =np.arange(t_init,t_fin,t_fin/10000)

delta_list = 2*pi*np.arange(-10000,10001,4000) 
delta_list = delta_list[np.where(delta_list != 0)] 

counts_list = np.zeros(len(delta_list))

y0 = [0+0j,1+0j]
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, t_eval=t_eval, args=(omega,delta,omega_Rabi))
    y_up = abs(sol["y"][0])**2 

    counts_list[i] = y_up[-1]

Answer 1

The following function will do the trick:

def model(x, w):
    y = []
    for z in x:
        sol = integrate.solve_ivp(func, [t_init, t_fin], y0, rtol=1e-9, atol=1e-11, t_eval=t_eval, args=(omega, z, w))
        y.append(np.abs(sol["y"][0][-1])**2)
    return np.array(y)

popt, pcov = optimize.curve_fit(model, delta_list, counts_list)
# (array([6408.84901332]), array([[1.80007176e-24]]))