Hacker News new | ask | show | jobs
by Armavica 1127 days ago
diffrax is absolutely magical. I had to integrate a lot of ODEs during my PhD, so I spent quite some time choosing and tuning the scipy solvers for my problems, and I thought that I came close to the fastest I could do in Python. Recently, out of curiosity I rewrote a stiff system that I was studying to solve it with diffrax, and was astonished when I saw it being solved 150x faster.
2 comments
patrickkidger 1127 days ago
This is awesome to hear!

And if you're doing implicit methods, then you may be interested to hear that today's release of Diffrax now includes IMEX solvers! Sil3, KenCarp3, KenCarp4, KenCarp5.

link
Armavica 1127 days ago
I did see that on my GitHub feed this morning, it was the good news of the day :) Thank you for creating and maintaining this library! I also really enjoyed your blogpost "How to succeed in a machine learning PhD" which is a trove of interesting things to learn.
link
geysersam 1127 days ago
That's insane. May I ask what scipy solver you used and what you switched to in diffrax? What kind of ODE were you working with?
link
Armavica 1127 days ago
It is a system of four coupled ODEs from biology. I was using LSODA with relatively tight rtol and atol. Now in diffrax I use Tsit5, with adaptive step size control.

Something really annoying about this system is that for a non-negligible portion of the parameter space, an adaptive step solver would fail and give up at some point during the integration because the step size converges to zero. This was preventing me from doing large parameter searches. Now, because diffrax makes it easy to specify the step size controller, I first try to solve the system with a PID step size control, and if this fails, I rerun it with small fixed steps, which is slower but always goes through the end. This guarantees that I will always get a complete solution, and that it will still be fast in 95% of cases, which is really a huge improvement.

link
geysersam 1127 days ago
Okay, interesting. For such a small system I can imagine the jit compilation helps a lot.

Using Numba to jit compile the rhs might speed things up in the non-diffrax version.

Does the LSODA solver not take a minimum time step parameter? Setting that to your small fixed step might also help.

link
patrickkidger 1127 days ago
Incidentally `diffrax.PIDController` also has a `dtmin` argument that could probably be used here instead of re-running things. :)
link