I'm using following code but it takes about an hour. Function chi and integrand are working fine, but chi_dop is taking too much time.
How to make it faster? Any better way to integrate other than scipy.integrate.quad? Which is faster either to integrate imaginary part or first integrating complex equation then taking imaginary part?
hbar = 1.05e-34
c = 3e8
ep = 8.854187817e-12
td = 23.93e-9
slt = 24.24e-9
dt=np.linspace(-6, 16, 12027)*1e9
dw = 200 #m/s
@njit
def chi(r, dt, dm, slt, td, tf):
A = np.array([[-1/slt, 1j*r, -1j*r], [-1j*r/2, 0, (1j*dt - 1/td)], [1j*r/2, (-1j*dt - 1/td), 0]])
B = np.array([0, -1j*r/2, 1j*r/2])
x = np.linalg.solve(A, B)
return (2*tf/c)* 2*dm**2*np.imag(x[2])/(hbar*r*ep)*1e14
def integrand(v, r, dt, dw, dm, slt, td, tf):
int = chi(r,dt-v/c*tf, dm, slt, td, tf)
return int* np.exp(-v**2 / (2 * dw**2)) / (np.sqrt(2 * np.pi) * dw)
def chi_dop(r, dt, dw, dm, slt, td,tf):
chi_int = sp.integrate.quad(integrand, -10*dw, 10*dw, args=(r, dt, dw, dm, slt, td,tf), limit=100, complex_func=True)[0]
return 1e-14*chi_int
chi=np.vectorize(chi)
chi_db = np.vectorize(chi_dop)
F1= chi_db(1/9, dt, dw, 1/18*3.3029259999999997e-29, slt, td, 2*np.pi*384.2276916107209e12) + chi_db(5/9, dt-157e6*2*np.pi, dw, 5/18*3.3029259999999997e-29, slt, td, 2*np.pi*384.2278485512209e12) + chi_db(14/9, dt-423e6*2*np.pi, dw, 7/9*3.3029259999999997e-29, slt, td, 2*np.pi*384.2281881125209e12)
F2 = chi_db(2/9, dt-6765e6*2*np.pi, dw, 1/9*3.3029259999999997e-29,slt, td, 2*np.pi*384.2344540713318e12) + chi_db(5/9, d-6834te6*2*np.pi, dw, 5/18*3.3029259999999997e-29, slt, td, 2*np.pi*384.23452629333184e12) + chi_db(5/9, dt-6992e6*2*np.pi, dw, 5/18*3.3029259999999997e-29, slt, td, 2*np.pi*384.2349498859318e12)
F3 = chi_db(20/81, dt-1443e6*2*np.pi, dw, 10/81*3.3029259999999997e-29, slt, td, 2*np.pi*384.2320640089228e12) + chi_db(70/81, dt-1510e6*2*np.pi, dw, 35/81*3.3029259999999997e-29, slt, td, 2*np.pi*384.2320933819228e12) + chi_db(2, dt-1630e6*2*np.pi, dw, 1*3.3029259999999997e-29, slt, td, 2*np.pi*384.2321567819228e12)
F4 = chi_db(2/3, dt-4436e6*2*np.pi, dw, 1/3*3.3029259999999997e-29,slt, td, 2*np.pi*384.2290576494837e12) + chi_db(70/81, dt-4478e6*2*np.pi, dw, 35/81*3.3029259999999997e-29, slt, td, 2*np.pi*384.2291210494837e12) + chi_db(56/81, dt-4545e6*2*np.pi, dw, 28/81*3.3029259999999997e-29, slt, t,d 2*np.pi*384.2292416894837e12)
I would suggest three modifications: removing vectorize, adding an @njit annotation to integrand(), and using a closed form solution for your matrix problem.
Let's start with removing these lines:
chi=np.vectorize(chi)
chi_db = np.vectorize(chi_dop)
You are not using these functions in a vectorized fashion, and even if you were, it would be better to use numba.vectorize rather than numpy.vectorize, as the latter is implemented using a Python loop.
I benchmarked running the integrand() function, with and without this change.
using np.vectorize
35.4 µs ± 1.91 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
after removing np.vectorize
6.62 µs ± 30.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
Next, let's tell Numba to compile your integrand function as well. Add the @njit annotation to integrand().
@njit
def integrand(v, r, dt, dw, dm, slt, td, tf):
int = chi(r,dt-v/c*tf, dm, slt, td, tf)
return int* np.exp(-v**2 / (2 * dw**2)) / (np.sqrt(2 * np.pi) * dw)
Benchmark:
Previous result
6.62 µs ± 30.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
After @njit
2.18 µs ± 11.6 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
Next, let's find a closed form for the solution of your matrix problem. I'm going to use SymPy, which can find an algebraic solution that we can re-use for many values. The final solution won't use SymPy.
import sympy
slt = sympy.Symbol('slt')
r = sympy.Symbol('r')
dt = sympy.Symbol('dt')
td = sympy.Symbol('td')
A = sympy.Matrix([[-1/slt, 1j*r, -1j*r], [-1j*r/2, 0, (1j*dt - 1/td)], [1j*r/2, (-1j*dt - 1/td), 0]])
B = sympy.Matrix([0, -1j*r/2, 1j*r/2])
x = A.solve(B)
print(x.tolist()[2])
Output:
[(-0.5*dt*r*td**2 + 0.5*I*r*td)/(1.0*dt**2*td**2 + 1.0*r**2*slt*td + 1.0)]
This is the algebraic value of x[2]. Directly calculating this will be faster than solving the matrix each time. To use this in Python, we must replace I, which represents an imaginary number, with 1j.
New definition of chi(), using the solution SymPy found:
@njit
def chi(r, dt, dm, slt, td, tf):
x = (-0.5*dt*r*td**2 + 0.5*1j*r*td)/(1.0*dt**2*td**2 + 1.0*r**2*slt*td + 1.0)
return (2*tf/c)* 2*dm**2*np.imag(x)/(hbar*r*ep)*1e14
Benchmark:
Matrix Solve
2.18 µs ± 11.6 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
Closed Form
671 ns ± 15.7 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
This is a 50x speedup overall.