I have a function, which I am trying to speed up using the @jit decorator from Numba module. For me it is essential to speed this up as much as possible, because my main code ca
There may be other issues, but one is that referencing an array in a module namespace seems to currently be unsupported (simple repro below). Try importing omega
as a name.
In [14]: %%file Sweep.py
...: import numpy as np
...: constant_val = 0.5
...: constant_arr = np.array([0, 1.5, 2.])
Overwriting Sweep.py
In [15]: Sweep.constant_val
Out[15]: 0.5
In [16]: Sweep.constant_arr
Out[16]: array([ 0. , 1.5, 2. ])
In [17]: @njit
...: def f(value):
...: return value + Sweep.constant_val
...:
In [18]: f(100)
Out[18]: 100.5
In [19]: @njit
...: def f(value):
...: return value + Sweep.constant_arr[0]
In [20]: f(100)
LoweringError: Failed at nopython (nopython mode backend)
'NoneType' object has no attribute 'module'
File "<ipython-input-19-0a259ade6b9e>", line 3
[1] During: lowering "$0.3 = getattr(value=$0.2, attr=constant_arr)" at <ipython-input-19-0a259ade6b9e> (3)
I can't see from your code why this isn't vectorizable. Vectorizing can speed up this kind of Python code by around 100x. Not sure of how it does relative to jit.
It looks like you could, for instance, take your dEdt out of the loop, and compute it in one step with something like :
dEdt = 0.5 * (Cmx[:, :, :, 0] * dg * (1+A*1j) * Nf[:] * Ef[:, None] * np.exp( 1j* (Sweep.omega[None, :, None, None]-Sweep.omega) *tijd)).sum(axis=2).sum(axis=1) - 0.5*(pd-g)*Ef + fbr*Ef2 + Kinj*EAinj*(1 + np.exp(1j*(u+Vmzm)) )
(Though I don't really know what the dimensionality of your Sweet.omega is).