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.
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.