speed-up trials with numba

[cosimoNigro]

As we discussed in last week's call @pawel21 will take care of trying to speed-up the SED computations using numba.
I suggested him to copy, in a script or in a notebook, one of the evaluate_sed_flux (e.g. the synchrotron one which is the simplest) and try to decorate it with numba.
After my latest modifications (PR #52) functions are @classmethods, so you do not need any other agnpy class (like the blob) to run them. They take directly as input all the physical parameters necessary to compute the SED (before some of them where provided by the blob - like spectral parameters and magnetic field).

[cosimoNigro]

@pawel21 sometimes ago I sent you an invitation to become a collaborator of the project. It is still pending, can you search in your mailbox and accept it. If you accept it you'll have more edit privileges and I can assign you issues 😄

[pawel21]

@pawel21 sometimes ago I sent you an invitation to become a collaborator of the project. It is still pending, can you search in your mailbox and accept it. If you accept it you'll have more edit privileges and I can assign you issues smile

Sorry, I missed it. Could you send me again? (because where I tried accept I got error: "This invitation has expired. ")

[pawel21]

Hi,

I modified code to speed-up fit with numba. You can find my code at branch speed_up_fit branch . I created new file synchotron_jit, synchrotron_self_compton_jit.py and kernels_jit.py, where I added numba code. Notebook: ssc fit jit . Let me know, what do you think about my solution.

[cosimoNigro]

Thanks a lot @pawel21, I checked it.
Would you mind, instead of using the fitting notebook for testing your new class, to set-up a proper comparison just invoking the evaluate_sed_flux functions?
You can take a look into experiments/profiling where I already set-up some small script to do test and profiling.
It would be interesting to check first the results from the non-decorated and decorated class using the basic functions computing the SED Synchrotron.evaluate_sed_flux. This function is calling the other functions you have decorated and all the other functions to compute the SED (e.g. Synchrotron.sed_flux(nu)) call this one anyhow.
Thanks a lot for your work!

[pawel21]

Hi, I added two scripts to profile the code with numba profile_ec_jit and profile_ssc_jit. I added numba code for EC with ring torus. Progress in time execution time is ~ 2 times faster. What do you think? So far, I created new python files with numba code external_comptin_jit, synchotron_self_compton_jit and synchotron_jit.py, but I think code should't be divide into separate file in final state. What should be next step? Add class ExternalComptonJit (and another) to external_compton.py, modified class ExternalCompton?

[cosimoNigro]

Thanks @pawel21 and sorry for having ignored you.
Last week we had the software school and I had no time at all to take a look at this.

I had a similar improvement in execution times with numba, I remember something like half of it.

Now, concerning your tests I see two problems:

1. a particular one: your fork is 277 commits behind master 😨 and I did a lot of refactoring and improvements that perhaps will change (though for sure not dramatically the results). Would you mind syncing your fork and retry? If you can't or don't want we can start with your decorators from scratch in a new branch in the master and work on it together;
2. a general one: adding numba will imply to have another dependency and the improvement does not look so dramatic either.
   We should decide which is a decent speed-up factor that will convince us that is worth adding a new dependency.

Thanks again and let me know how you want to proceed.

[jsitarek]

Hi @pawel21 @cosimoNigro
looking at the code of Paweł, I saw that the most external functions:
compton_kernel_jit and isotropic_kernel_jit do not have any of the numba decorators, however those are being calculated over an array of electron gamma values and an array of nu values, and thus would paralellize nicely (instead of paralelizing individual subfunctions). Was it done on purpose?

Another thing that I was thinking about is cache-ing the function values, there is a decorator for this as well. I'm not sure if this will help here, because we do not have any single parameter functions that is executed multiple times for the same values, but if it is easy we can try it with cos_psi_jit which has 3 angle parameters, but it is executed multiple times for different nus and gammas.

The third thing, compton_kernel_jit is using a number of small one-line functions, get_y, get_y_1, get_y_2, get_values that I do not think are used anywhere else. I guess there should be some overhead in Python from executing a function (especially if we use extra decorators) and those are executed all the time, so I would naively assume that moving those lines inside the main compton_kernel_jit function would speed some somewhat the code (and make it shorter as well.