hgrecco / numbakit-ode

Leveraging numba to speed up ODE integration
Other
68 stars 3 forks source link

Experiments with DOP853 and numba #30

Open s-m-e opened 8 months ago

s-m-e commented 8 months ago

Hi folks, in recent years you had some encounters with the developers of poliastro via the bug tracker (e.g. #11 or #13) and other places (numba's discourse e.g. here). Since poliastro has been archived, I am continuing its development in a fork. A fast implementation of DOP853 (faster than the one in scipy) somehow compiled by numba has been a long standing issue and there were attempts to port poliastro from relying on scipy over to numbakit-ode, which have stalled for a number of reasons.

I recently spent some time on isolating the relevant portions of scipy and turning them into a purely functional implementation compiled entirely by numba - with identical yet much faster results as far as my use-case is concerned. The idea is that the solver can run in parallel on both CPUs and GPUs, the latter via CUDA (thanks to numba's support for it), unaltered with an identical code base. My work might not be applicable to your package, but it might be an interesting inspiration never the less. I'd also be interested in reviews / comments on the implementation, still WIP.

The entire thing uses tuples instead of arrays (limitation imposed by CUDA compatibility), so it's over 3k lines of code ... The funny thing though is that this also makes its significantly (up to two orders of magnitude) faster single-threaded on CPUs.