search

opticspy / lightpipes

LightPipes for Python, "Pure Python version"
https://opticspy.github.io/lightpipes/
BSD 3-Clause "New" or "Revised" License
227 stars 52 forks source link

extend with Numba? #43

Open raghavchhetri opened 3 years ago

raghavchhetri commented 3 years ago

First of all, thank you for providing this wonderful toolbox. I only recently started exploring it and I like it quite a bit!

I was wondering if there is any consideration to using Numba which uses LLVM for JIT compilation of the python code to machine code? Since Numba also supports JIT compilation of code to run on GPUs, significant speed-ups might be possible. For instance, I've noticed that the Steps function takes a really long time at the moment (Forward too, but not as much as Steps). Would love to hear the future roadmap of this project, and I look forward to being a regular user. Thanks again for providing this user-friendly toolbox in Python.

FredvanGoor commented 3 years ago

Sounds interesting, but I tried it and encountered a lot of errors. So it it is not easy. Before version 2.0 we wrote the package in C++ which should be the fastest. The pure-python versions turned out to be as fast as the C++ version thanks to numpy, so can we achieve a better performance using Numba? You, or anybody, could try to modify the code using Numba and report the results to me. Thanks for the idea anyway!

Fred.

ldoyle commented 3 years ago

Indeed I made some quick tests with Numba while porting the code from Cpp to Python. However for the few examples that I tested, it was a little complicated (e.g. you cannot allocate new arrays inside a @jit function, so you have to refactor some code) and gave no or no significant speedup. I might have been using it wrong or at the wrong places where not much is to be gained.

If you can make it work as a proof of principle with e.g. Forward, Forvard or Steps, that would definitely be cool.

On the long run, it would be interesting to see if it can be done, however I would suggest to keep numba dependency optional, since the only strict dependencies so far are numpy and scipy, both of which should be widely available on any Python platform (e.g. Raspi). pyfftw is also faster than numpy.fft, but optional since it's hardly any code change.

Regarding Forward, as mentioned in the manual, this becomes impractically slow very quickly, but Steps could be optimized more I'm pretty sure. For most methods I compared the performance of the Cpp version (up to LightPipes 1.2) and the Python/numpy version, however for Steps I might have never done it, I do think it is a little slower than the Cpp counterpart.